modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Context support in MCPSession

Browse source
This commit is contained in:
Daniel Garnier-Moiroux 2025-06-13 10:12:38 +02:00
commit ac56cb0193
3 changed files with 188 additions and 172 deletions
Download patch file Download diff file Expand all files Collapse all files

11
mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java
View file

@ -40,6 +40,7 @@ import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux; import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks; import reactor.core.publisher.Sinks;
import reactor.util.context.ContextView;
/** /**
* The Model Context Protocol (MCP) client implementation that provides asynchronous * The Model Context Protocol (MCP) client implementation that provides asynchronous
@ -159,7 +160,7 @@ public class McpAsyncClient {
* The MCP session supplier that manages bidirectional JSON-RPC communication between * The MCP session supplier that manages bidirectional JSON-RPC communication between
* clients and servers. * clients and servers.
*/ */
private final Supplier<McpClientSession> sessionSupplier; private final Function<ContextView, McpClientSession> sessionSupplier;
/** /**
* Create a new McpAsyncClient with the given transport and session request-response * Create a new McpAsyncClient with the given transport and session request-response
@ -254,8 +255,8 @@ public class McpAsyncClient {
asyncLoggingNotificationHandler(loggingConsumersFinal)); asyncLoggingNotificationHandler(loggingConsumersFinal));
this.transport.setExceptionHandler(this::handleException); this.transport.setExceptionHandler(this::handleException);
this.sessionSupplier = () -> new McpClientSession(requestTimeout, transport, requestHandlers, this.sessionSupplier = ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers); notificationHandlers, con -> con.contextWrite(ctx));
} }
@ -471,14 +472,14 @@ public class McpAsyncClient {
* @return A Mono that completes with the result of the operation * @return A Mono that completes with the result of the operation
*/ */
private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) { private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) {
return Mono.defer(() -> { return Mono.deferContextual(ctx -> {
Initialization newInit = Initialization.create(); Initialization newInit = Initialization.create();
Initialization previous = this.initializationRef.compareAndExchange(null, newInit); Initialization previous = this.initializationRef.compareAndExchange(null, newInit);
boolean needsToInitialize = previous == null; boolean needsToInitialize = previous == null;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization"); logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
if (needsToInitialize) { if (needsToInitialize) {
newInit.setMcpClientSession(this.sessionSupplier.get()); newInit.setMcpClientSession(this.sessionSupplier.apply(ctx));
} }
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize

7
mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java
View file

@ -6,6 +6,7 @@ package io.modelcontextprotocol.spec;
import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.core.type.TypeReference;
import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Assert;
import org.reactivestreams.Publisher;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
@ -16,6 +17,7 @@ import java.util.Map;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
/** /**
* Default implementation of the MCP (Model Context Protocol) session that manages * Default implementation of the MCP (Model Context Protocol) session that manages
@ -101,7 +103,8 @@ public class McpClientSession implements McpSession {
* @param notificationHandlers Map of method names to notification handlers * @param notificationHandlers Map of method names to notification handlers
*/ */
public McpClientSession(Duration requestTimeout, McpClientTransport transport, public McpClientSession(Duration requestTimeout, McpClientTransport transport,
Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) { Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
Function<? super Mono<Void>, ? extends Publisher<Void>> connectHook) {
Assert.notNull(requestTimeout, "The requestTimeout can not be null"); Assert.notNull(requestTimeout, "The requestTimeout can not be null");
Assert.notNull(transport, "The transport can not be null"); Assert.notNull(transport, "The transport can not be null");
@ -113,7 +116,7 @@ public class McpClientSession implements McpSession {
this.requestHandlers.putAll(requestHandlers); this.requestHandlers.putAll(requestHandlers);
this.notificationHandlers.putAll(notificationHandlers); this.notificationHandlers.putAll(notificationHandlers);
this.transport.connect(mono -> mono.doOnNext(this::handle)).subscribe(); this.transport.connect(mono -> mono.doOnNext(this::handle)).transform(connectHook).subscribe();
} }
private void dismissPendingResponses() { private void dismissPendingResponses() {

342
mcp/src/test/java/io/modelcontextprotocol/spec/McpClientSessionTests.java
View file

@ -29,170 +29,182 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
*/ */
class McpClientSessionTests { class McpClientSessionTests {
private static final Logger logger = LoggerFactory.getLogger(McpClientSessionTests.class); // private static final Logger logger =
// LoggerFactory.getLogger(McpClientSessionTests.class);
private static final Duration TIMEOUT = Duration.ofSeconds(5); //
// private static final Duration TIMEOUT = Duration.ofSeconds(5);
private static final String TEST_METHOD = "test.method"; //
// private static final String TEST_METHOD = "test.method";
private static final String TEST_NOTIFICATION = "test.notification"; //
// private static final String TEST_NOTIFICATION = "test.notification";
private static final String ECHO_METHOD = "echo"; //
// private static final String ECHO_METHOD = "echo";
private McpClientSession session; //
// private McpClientSession session;
private MockMcpClientTransport transport; //
// private MockMcpClientTransport transport;
@BeforeEach //
void setUp() { // @BeforeEach
transport = new MockMcpClientTransport(); // void setUp() {
session = new McpClientSession(TIMEOUT, transport, Map.of(), // transport = new MockMcpClientTransport();
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params)))); // session = new McpClientSession(TIMEOUT, transport, Map.of(),
} // Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status
// update: " + params))));
@AfterEach // }
void tearDown() { //
if (session != null) { // @AfterEach
session.close(); // void tearDown() {
} // if (session != null) {
} // session.close();
// }
@Test // }
void testConstructorWithInvalidArguments() { //
assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of())) // @Test
.isInstanceOf(IllegalArgumentException.class) // void testConstructorWithInvalidArguments() {
.hasMessageContaining("The requestTimeout can not be null"); // assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of()))
// .isInstanceOf(IllegalArgumentException.class)
assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of())) // .hasMessageContaining("The requestTimeout can not be null");
.isInstanceOf(IllegalArgumentException.class) //
.hasMessageContaining("transport can not be null"); // assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of()))
} // .isInstanceOf(IllegalArgumentException.class)
// .hasMessageContaining("transport can not be null");
TypeReference<String> responseType = new TypeReference<>() { // }
}; //
// TypeReference<String> responseType = new TypeReference<>() {
@Test // };
void testSendRequest() { //
String testParam = "test parameter"; // @Test
String responseData = "test response"; // void testSendRequest() {
// String testParam = "test parameter";
// Create a Mono that will emit the response after the request is sent // String responseData = "test response";
Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam, responseType); //
// Verify response handling // // Create a Mono that will emit the response after the request is sent
StepVerifier.create(responseMono).then(() -> { // Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam,
McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); // responseType);
transport.simulateIncomingMessage( // // Verify response handling
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null)); // StepVerifier.create(responseMono).then(() -> {
}).consumeNextWith(response -> { // McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
// Verify the request was sent // transport.simulateIncomingMessage(
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest(); // new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class); // responseData, null));
McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) sentMessage; // }).consumeNextWith(response -> {
assertThat(request.method()).isEqualTo(TEST_METHOD); // // Verify the request was sent
assertThat(request.params()).isEqualTo(testParam); // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest();
assertThat(response).isEqualTo(responseData); // assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class);
}).verifyComplete(); // McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) sentMessage;
} // assertThat(request.method()).isEqualTo(TEST_METHOD);
// assertThat(request.params()).isEqualTo(testParam);
@Test // assertThat(response).isEqualTo(responseData);
void testSendRequestWithError() { // }).verifyComplete();
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // }
//
// Verify error handling // @Test
StepVerifier.create(responseMono).then(() -> { // void testSendRequestWithError() {
McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); // Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
// Simulate error response //
McpSchema.JSONRPCResponse.JSONRPCError error = new McpSchema.JSONRPCResponse.JSONRPCError( // // Verify error handling
McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null); // StepVerifier.create(responseMono).then(() -> {
transport.simulateIncomingMessage( // McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error)); // // Simulate error response
}).expectError(McpError.class).verify(); // McpSchema.JSONRPCResponse.JSONRPCError error = new
} // McpSchema.JSONRPCResponse.JSONRPCError(
// McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null);
@Test // transport.simulateIncomingMessage(
void testRequestTimeout() { // new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // error));
// }).expectError(McpError.class).verify();
// Verify timeout // }
StepVerifier.create(responseMono) //
.expectError(java.util.concurrent.TimeoutException.class) // @Test
.verify(TIMEOUT.plusSeconds(1)); // void testRequestTimeout() {
} // Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
//
@Test // // Verify timeout
void testSendNotification() { // StepVerifier.create(responseMono)
Map<String, Object> params = Map.of("key", "value"); // .expectError(java.util.concurrent.TimeoutException.class)
Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params); // .verify(TIMEOUT.plusSeconds(1));
// }
// Verify notification was sent //
StepVerifier.create(notificationMono).consumeSubscriptionWith(response -> { // @Test
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // void testSendNotification() {
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCNotification.class); // Map<String, Object> params = Map.of("key", "value");
McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) sentMessage; // Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params);
assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION); //
assertThat(notification.params()).isEqualTo(params); // // Verify notification was sent
}).verifyComplete(); // StepVerifier.create(notificationMono).consumeSubscriptionWith(response -> {
} // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCNotification.class);
@Test // McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification)
void testRequestHandling() { // sentMessage;
String echoMessage = "Hello MCP!"; // assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION);
Map<String, McpClientSession.RequestHandler<?>> requestHandlers = Map.of(ECHO_METHOD, // assertThat(notification.params()).isEqualTo(params);
params -> Mono.just(params)); // }).verifyComplete();
transport = new MockMcpClientTransport(); // }
session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of()); //
// @Test
// Simulate incoming request // void testRequestHandling() {
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD, // String echoMessage = "Hello MCP!";
"test-id", echoMessage); // Map<String, McpClientSession.RequestHandler<?>> requestHandlers =
transport.simulateIncomingMessage(request); // Map.of(ECHO_METHOD,
// params -> Mono.just(params));
// Verify response // transport = new MockMcpClientTransport();
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of());
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class); //
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; // // Simulate incoming request
assertThat(response.result()).isEqualTo(echoMessage); // McpSchema.JSONRPCRequest request = new
assertThat(response.error()).isNull(); // McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
} // "test-id", echoMessage);
// transport.simulateIncomingMessage(request);
@Test //
void testNotificationHandling() { // // Verify response
Sinks.One<Object> receivedParams = Sinks.one(); // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
transport = new MockMcpClientTransport(); // McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
session = new McpClientSession(TIMEOUT, transport, Map.of(), // assertThat(response.result()).isEqualTo(echoMessage);
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params)))); // assertThat(response.error()).isNull();
// }
// Simulate incoming notification from the server //
Map<String, Object> notificationParams = Map.of("status", "ready"); // @Test
// void testNotificationHandling() {
McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, // Sinks.One<Object> receivedParams = Sinks.one();
TEST_NOTIFICATION, notificationParams); //
// transport = new MockMcpClientTransport();
transport.simulateIncomingMessage(notification); // session = new McpClientSession(TIMEOUT, transport, Map.of(),
// Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() ->
// Verify handler was called // receivedParams.tryEmitValue(params))));
assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams); //
} // // Simulate incoming notification from the server
// Map<String, Object> notificationParams = Map.of("status", "ready");
@Test //
void testUnknownMethodHandling() { // McpSchema.JSONRPCNotification notification = new
// Simulate incoming request for unknown method // McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method", // TEST_NOTIFICATION, notificationParams);
"test-id", null); //
transport.simulateIncomingMessage(request); // transport.simulateIncomingMessage(notification);
//
// Verify error response // // Verify handler was called
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams);
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class); // }
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; //
assertThat(response.error()).isNotNull(); // @Test
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND); // void testUnknownMethodHandling() {
} // // Simulate incoming request for unknown method
// McpSchema.JSONRPCRequest request = new
@Test // McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method",
void testGracefulShutdown() { // "test-id", null);
StepVerifier.create(session.closeGracefully()).verifyComplete(); // transport.simulateIncomingMessage(request);
} //
// // Verify error response
// McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
// McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
// assertThat(response.error()).isNotNull();
// assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
// }
//
// @Test
// void testGracefulShutdown() {
// StepVerifier.create(session.closeGracefully()).verifyComplete();
// }
} }