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

Optimize client nested streams in McpClientSession (#33)

Browse source 
Signed-off-by: JermaineHua <crazyhzm@apache.org>
This commit is contained in:
JermaineHua 2025-04-16 23:10:59 +08:00 • committed by Dariusz Jędrzejczyk
commit 04046ca05b
No known key found for this signature in database
GPG key ID: FD786A4E2335C448
1 changed files with 20 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -122,7 +122,12 @@ public class McpClientSession implements McpSession {
// Observation associated with the individual message - it can be used to // Observation associated with the individual message - it can be used to
// create child Observation and emit it together with the message to the // create child Observation and emit it together with the message to the
// consumer // consumer
this.connection = this.transport.connect(mono -> mono.doOnNext(message -> { this.connection = this.transport.connect(mono -> mono.doOnNext(message -> handle(message).subscribe()))
.subscribe();
}
public Mono<Void> handle(McpSchema.JSONRPCMessage message) {
return Mono.defer(() -> {
if (message instanceof McpSchema.JSONRPCResponse response) { if (message instanceof McpSchema.JSONRPCResponse response) {
logger.debug("Received Response: {}", response); logger.debug("Received Response: {}", response);
var sink = pendingResponses.remove(response.id()); var sink = pendingResponses.remove(response.id());
@ -132,23 +137,27 @@ public class McpClientSession implements McpSession {
else { else {
sink.success(response); sink.success(response);
} }
return Mono.empty();
} }
else if (message instanceof McpSchema.JSONRPCRequest request) { else if (message instanceof McpSchema.JSONRPCRequest request) {
logger.debug("Received request: {}", request); logger.debug("Received request: {}", request);
handleIncomingRequest(request).subscribe(response -> transport.sendMessage(response).subscribe(), return handleIncomingRequest(request).onErrorResume(error -> {
error -> { var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
null, new McpSchema.JSONRPCResponse.JSONRPCError( error.getMessage(), null));
McpSchema.ErrorCodes.INTERNAL_ERROR, error.getMessage(), null)); return this.transport.sendMessage(errorResponse).then(Mono.empty());
transport.sendMessage(errorResponse).subscribe(); }).flatMap(this.transport::sendMessage);
});
} }
else if (message instanceof McpSchema.JSONRPCNotification notification) { else if (message instanceof McpSchema.JSONRPCNotification notification) {
logger.debug("Received notification: {}", notification); logger.debug("Received notification: {}", notification);
handleIncomingNotification(notification).subscribe(null, return handleIncomingNotification(notification)
error -> logger.error("Error handling notification: {}", error.getMessage())); .doOnError(error -> logger.error("Error handling notification: {}", error.getMessage()));
} }
})).subscribe(); else {
logger.warn("Received unknown message type: {}", message);
return Mono.empty();
}
});
} }
/** /**