refactor: consolidate MCP client initialization logic (#334)

- 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

View file

@ -389,7 +389,12 @@ public class McpAsyncClient {
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);
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
@ -412,6 +417,9 @@ public class McpAsyncClient {
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
.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;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
if (needsToInitialize) {
newInit.setMcpClientSession(this.sessionSupplier.get());
}
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
? doInitialize(newInit.mcpSession()).doOnNext(newInit::complete).onErrorResume(ex -> {
newInit.error(ex);
return Mono.error(ex);
}) : previous.await();
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit)
: previous.await();
return initializationJob.map(initializeResult -> this.initializationRef.get())
.timeout(this.initializationTimeout)

View file

@ -229,27 +229,27 @@ public class McpClientSession implements McpSession {
public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
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);
this.pendingResponses.put(requestId, sink);
this.pendingResponses.put(requestId, pendingResponseSink);
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
requestId, requestParams);
this.transport.sendMessage(jsonrpcRequest).contextWrite(ctx).subscribe(v -> {
}, error -> {
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) {
logger.error("Error handling request: {}", jsonRpcResponse.error());
sink.error(new McpError(jsonRpcResponse.error()));
deliveredResponseSink.error(new McpError(jsonRpcResponse.error()));
}
else {
if (typeRef.getType().equals(Void.class)) {
sink.complete();
deliveredResponseSink.complete();
}
else {
sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef));
deliveredResponseSink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef));
}
}
});