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.Mono;
import reactor.core.publisher.Sinks;
import reactor.util.context.ContextView;
/**
* 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
* 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
@ -254,8 +255,8 @@ public class McpAsyncClient {
asyncLoggingNotificationHandler(loggingConsumersFinal));
this.transport.setExceptionHandler(this::handleException);
this.sessionSupplier = () -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers);
this.sessionSupplier = ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers, con -> con.contextWrite(ctx));
}
@ -471,14 +472,14 @@ public class McpAsyncClient {
* @return A Mono that completes with the result of the 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 previous = this.initializationRef.compareAndExchange(null, newInit);
boolean needsToInitialize = previous == null;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
if (needsToInitialize) {
newInit.setMcpClientSession(this.sessionSupplier.get());
newInit.setMcpClientSession(this.sessionSupplier.apply(ctx));
}
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 io.modelcontextprotocol.util.Assert;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
@ -16,6 +17,7 @@ import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
/**
* 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
*/
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(transport, "The transport can not be null");
@ -113,7 +116,7 @@ public class McpClientSession implements McpSession {
this.requestHandlers.putAll(requestHandlers);
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() {

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