Polish modelcontextprotocol#33

Signed-off-by: Dariusz Jędrzejczyk <[email protected]>

Author: chemicL

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

@@ -122,42 +122,38 @@ 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 -> handle(message).subscribe()))
-			.subscribe();
+		this.connection = this.transport.connect(mono -> mono.doOnNext(this::handle)).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 unknown id {}", response.id());
-				}
-				else {
-					sink.success(response);
-				}
-				return Mono.empty();
-			}
-			else if (message instanceof McpSchema.JSONRPCRequest request) {
-				logger.debug("Received request: {}", request);
-				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 this.transport.sendMessage(errorResponse).then(Mono.empty());
-				}).flatMap(this.transport::sendMessage);
-			}
-			else if (message instanceof McpSchema.JSONRPCNotification notification) {
-				logger.debug("Received notification: {}", notification);
-				return handleIncomingNotification(notification)
-					.doOnError(error -> logger.error("Error handling notification: {}", error.getMessage()));
+	private void handle(McpSchema.JSONRPCMessage message) {
+		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 unknown id {}", response.id());
 			}
 			else {
-				logger.warn("Received unknown message type: {}", message);
-				return Mono.empty();
+				sink.success(response);
 			}
-		});
+		}
+		else if (message instanceof McpSchema.JSONRPCRequest request) {
+			logger.debug("Received request: {}", request);
+			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 this.transport.sendMessage(errorResponse).then(Mono.empty());
+			}).flatMap(this.transport::sendMessage).subscribe();
+		}
+		else if (message instanceof McpSchema.JSONRPCNotification notification) {
+			logger.debug("Received notification: {}", notification);
+			handleIncomingNotification(notification)
+				.doOnError(error -> logger.error("Error handling notification: {}", error.getMessage()))
+				.subscribe();
+		}
+		else {
+			logger.warn("Received unknown message type: {}", message);
+		}
 	}
 
 	/**