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

refactor: consolidate MCP client initialization logic (#334)

Browse source 
- Move session setup, completion, and error callbacks into doInitialize method
- Rename variables in McpClientSession for better clarity (sink -> pendingResponseSink/deliveredResponseSink)

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov 2025-06-21 11:54:28 +02:00
commit be719f53d8
2 changed files with 18 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

20
mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java
View file

@ -389,7 +389,12 @@ public class McpAsyncClient {
return withSession("by explicit API call", init -> Mono.just(init.get())); return withSession("by explicit API call", init -> Mono.just(init.get()));
} }
private Mono<McpSchema.InitializeResult> doInitialize(McpClientSession mcpClientSession) { private Mono<McpSchema.InitializeResult> doInitialize(Initialization initialization) {
initialization.setMcpClientSession(this.sessionSupplier.get());
McpClientSession mcpClientSession = initialization.mcpSession();
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1); String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
@ -412,6 +417,9 @@ public class McpAsyncClient {
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null) return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
.thenReturn(initializeResult); .thenReturn(initializeResult);
}).doOnNext(initialization::complete).onErrorResume(ex -> {
initialization.error(ex);
return Mono.error(ex);
}); });
} }
@ -479,15 +487,9 @@ public class McpAsyncClient {
boolean needsToInitialize = previous == null; boolean needsToInitialize = previous == null;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization"); logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
if (needsToInitialize) {
newInit.setMcpClientSession(this.sessionSupplier.get());
}
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit)
? doInitialize(newInit.mcpSession()).doOnNext(newInit::complete).onErrorResume(ex -> { : previous.await();
newInit.error(ex);
return Mono.error(ex);
}) : previous.await();
return initializationJob.map(initializeResult -> this.initializationRef.get()) return initializationJob.map(initializeResult -> this.initializationRef.get())
.timeout(this.initializationTimeout) .timeout(this.initializationTimeout)

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

@ -229,27 +229,27 @@ public class McpClientSession implements McpSession {
public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) { public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
String requestId = this.generateRequestId(); String requestId = this.generateRequestId();
return Mono.deferContextual(ctx -> Mono.<McpSchema.JSONRPCResponse>create(sink -> { return Mono.deferContextual(ctx -> Mono.<McpSchema.JSONRPCResponse>create(pendingResponseSink -> {
logger.debug("Sending message for method {}", method); logger.debug("Sending message for method {}", method);
this.pendingResponses.put(requestId, sink); this.pendingResponses.put(requestId, pendingResponseSink);
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
requestId, requestParams); requestId, requestParams);
this.transport.sendMessage(jsonrpcRequest).contextWrite(ctx).subscribe(v -> { this.transport.sendMessage(jsonrpcRequest).contextWrite(ctx).subscribe(v -> {
}, error -> { }, error -> {
this.pendingResponses.remove(requestId); this.pendingResponses.remove(requestId);
sink.error(error); pendingResponseSink.error(error);
}); });
})).timeout(this.requestTimeout).handle((jsonRpcResponse, sink) -> { })).timeout(this.requestTimeout).handle((jsonRpcResponse, deliveredResponseSink) -> {
if (jsonRpcResponse.error() != null) { if (jsonRpcResponse.error() != null) {
logger.error("Error handling request: {}", jsonRpcResponse.error()); logger.error("Error handling request: {}", jsonRpcResponse.error());
sink.error(new McpError(jsonRpcResponse.error())); deliveredResponseSink.error(new McpError(jsonRpcResponse.error()));
} }
else { else {
if (typeRef.getType().equals(Void.class)) { if (typeRef.getType().equals(Void.class)) {
sink.complete(); deliveredResponseSink.complete();
} }
else { else {
sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef)); deliveredResponseSink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef));
} }
} }
}); });