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
// 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(message -> handle(message).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());
@ -132,23 +137,27 @@ public class McpClientSession implements McpSession {
else {
sink.success(response);
}
return Mono.empty();
}
else if (message instanceof McpSchema.JSONRPCRequest request) {
logger.debug("Received request: {}", request);
handleIncomingRequest(request).subscribe(response -> transport.sendMessage(response).subscribe(),
error -> {
var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
null, new McpSchema.JSONRPCResponse.JSONRPCError(
McpSchema.ErrorCodes.INTERNAL_ERROR, error.getMessage(), null));
transport.sendMessage(errorResponse).subscribe();
});
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);
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();
}
});
}
/**