refactor:The connection does not need to be in McpClientSession

Signed-off-by: JermaineHua <[email protected]>

Author: CrazyHZM

mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java

@@ -62,8 +62,6 @@ public class McpClientSession implements McpSession {
 	/** Atomic counter for generating unique request IDs */
 	private final AtomicLong requestCounter = new AtomicLong(0);
 
-	private final Disposable connection;
-
 	/**
 	 * Functional interface for handling incoming JSON-RPC requests. Implementations
 	 * should process the request parameters and return a response.
@@ -108,7 +106,7 @@ public interface NotificationHandler {
 	public McpClientSession(Duration requestTimeout, McpClientTransport transport,
 			Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) {
 
-		Assert.notNull(requestTimeout, "The requstTimeout can not be null");
+		Assert.notNull(requestTimeout, "The requestTimeout can not be null");
 		Assert.notNull(transport, "The transport can not be null");
 		Assert.notNull(requestHandlers, "The requestHandlers can not be null");
 		Assert.notNull(notificationHandlers, "The notificationHandlers can not be null");
@@ -123,33 +121,41 @@ public McpClientSession(Duration requestTimeout, McpClientTransport transport,
 		// Observation associated with the individual message - it can be used to
 		// create child Observation and emit it together with the message to the
 		// consumer
-		this.connection = this.transport.connect(mono -> mono.doOnNext(message -> {
+		this.transport.connect(mono -> mono.doOnNext((msg) -> handle(msg).subscribe()));
+	}
+
+	public Mono<Void> handle(McpSchema.JSONRPCMessage message) {
+		return Mono.defer(() -> {
 			if (message instanceof McpSchema.JSONRPCResponse response) {
 				logger.debug("Received Response: {}", response);
 				var sink = pendingResponses.remove(response.id());
 				if (sink == null) {
-					logger.warn("Unexpected response for unkown id {}", response.id());
+					logger.warn("Unexpected response for unknown id {}", response.id());
 				}
 				else {
 					sink.success(response);
 				}
+				return Mono.empty();
 			}
 			else if (message instanceof McpSchema.JSONRPCRequest request) {
 				logger.debug("Received request: {}", request);
-				handleIncomingRequest(request).flatMap(transport::sendMessage).onErrorResume(error -> {
+				return handleIncomingRequest(request).onErrorResume(error -> {
 					var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
 							new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
 									error.getMessage(), null));
-					return transport.sendMessage(errorResponse);
-				}).subscribe();
-
+					return this.transport.sendMessage(errorResponse).then(Mono.empty());
+				}).flatMap(this.transport::sendMessage);
 			}
 			else if (message instanceof McpSchema.JSONRPCNotification notification) {
 				logger.debug("Received notification: {}", notification);
-				handleIncomingNotification(notification).subscribe(null,
-						error -> logger.error("Error handling notification: {}", error.getMessage()));
+				return handleIncomingNotification(notification)
+					.doOnError(error -> logger.error("Error handling notification: {}", error.getMessage()));
 			}
-		})).subscribe();
+			else {
+				logger.warn("Received unknown message type: {}", message);
+				return Mono.empty();
+			}
+		});
 	}
 
 	/**
@@ -271,18 +277,14 @@ public Mono<Void> sendNotification(String method, Map<String, Object> params) {
 	 */
 	@Override
 	public Mono<Void> closeGracefully() {
-		return Mono.defer(() -> {
-			this.connection.dispose();
-			return transport.closeGracefully();
-		});
+		return Mono.defer(transport::closeGracefully);
 	}
 
 	/**
 	 * Closes the session immediately, potentially interrupting pending operations.
 	 */
 	@Override
 	public void close() {
-		this.connection.dispose();
 		transport.close();
 	}
 

mcp/src/test/java/io/modelcontextprotocol/MockMcpClientTransport.java

@@ -71,10 +71,11 @@ public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchem
 			return Mono.error(new IllegalStateException("Already connected"));
 		}
 		connected = true;
-		return inbound.asFlux()
+		inbound.asFlux()
 			.flatMap(message -> Mono.just(message).transform(handler))
 			.doFinally(signal -> connected = false)
-			.then();
+			.subscribe();
+		return Mono.empty();
 	}
 
 	@Override

mcp/src/test/java/io/modelcontextprotocol/spec/McpClientSessionTests.java

@@ -61,7 +61,7 @@ void tearDown() {
 	void testConstructorWithInvalidArguments() {
 		assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of()))
 			.isInstanceOf(IllegalArgumentException.class)
-			.hasMessageContaining("requstTimeout can not be null");
+			.hasMessageContaining("requestTimeout can not be null");
 
 		assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of()))
 			.isInstanceOf(IllegalArgumentException.class)