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 <[email protected]>

Author: tzolov

mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -389,7 +389,12 @@ public Mono<McpSchema.InitializeResult> initialize() {
 		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 @@ private Mono<McpSchema.InitializeResult> doInitialize(McpClientSession mcpClient
 
 			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 @@ private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono
 
 			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)

mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java

@@ -229,27 +229,27 @@ private String generateRequestId() {
 	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));
 				}
 			}
 		});