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

@@ -108,7 +108,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 +123,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.connection = this.transport.connect(mono -> mono.doOnNext((msg) -> handle(msg).subscribe())).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();
+			}
+		});
 	}
 
 	/**

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)