diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java b/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java index 8630c2b..d6804c8 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -389,7 +389,12 @@ public class McpAsyncClient { return withSession("by explicit API call", init -> Mono.just(init.get())); } - private Mono doInitialize(McpClientSession mcpClientSession) { + private Mono 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 initializationJob = needsToInitialize - ? doInitialize(newInit.mcpSession()).doOnNext(newInit::complete).onErrorResume(ex -> { - newInit.error(ex); - return Mono.error(ex); - }) : previous.await(); + Mono initializationJob = needsToInitialize ? doInitialize(newInit) + : previous.await(); return initializationJob.map(initializeResult -> this.initializationRef.get()) .timeout(this.initializationTimeout) diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java index c839924..fa0853d 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java @@ -229,27 +229,27 @@ public class McpClientSession implements McpSession { public Mono sendRequest(String method, Object requestParams, TypeReference typeRef) { String requestId = this.generateRequestId(); - return Mono.deferContextual(ctx -> Mono.create(sink -> { + return Mono.deferContextual(ctx -> Mono.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)); } } });