fix: share terminal initialization outcome

Author: DragonFSKY

mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java

@@ -213,7 +213,7 @@ private Mono<McpSchema.InitializeResult> await() {
 
 		private void complete(McpSchema.InitializeResult initializeResult) {
 			// inform all the subscribers waiting for the initialization
-			this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST);
+			this.initSink.tryEmitValue(initializeResult);
 		}
 
 		private void cacheResult(McpSchema.InitializeResult initializeResult) {
@@ -222,14 +222,18 @@ private void cacheResult(McpSchema.InitializeResult initializeResult) {
 		}
 
 		private void error(Throwable t) {
-			this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
+			this.initSink.tryEmitError(t);
 		}
 
 		private void close() {
 			this.mcpSession().close();
 		}
 
 		private void terminate(Throwable cause) {
+			// Initialization has a single shared outcome. Publish the terminal failure
+			// even when the session has not been installed yet so that both the owner
+			// and all concurrent joiners observe it immediately.
+			this.initSink.tryEmitError(cause);
 			McpClientSession mcpClientSession = this.mcpSession();
 			if (mcpClientSession != null) {
 				mcpClientSession.terminate(cause);
@@ -304,8 +308,20 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
 			boolean needsToInitialize = previous == null;
 			logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
 
-			Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
-					? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await();
+			Mono<McpSchema.InitializeResult> initializationJob;
+			if (needsToInitialize) {
+				// The work branch only publishes into the shared sink. Keeping it from
+				// winning directly makes the owner and all joiners consume the same
+				// first terminal signal.
+				Mono<McpSchema.InitializeResult> initializationWork = this
+					.doInitialize(newInit, this.postInitializationHook, ctx)
+					.onErrorComplete()
+					.then(Mono.never());
+				initializationJob = Mono.firstWithSignal(newInit.await(), initializationWork);
+			}
+			else {
+				initializationJob = previous.await();
+			}
 
 			return initializationJob.map(initializeResult -> this.initializationRef.get())
 				.timeout(this.initializationTimeout)
@@ -322,46 +338,45 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
 	private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization,
 			Function<Initialization, Mono<Void>> postInitOperation, ContextView ctx) {
 
-		initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
+		return Mono.defer(() -> {
+			initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
 
-		McpClientSession mcpClientSession = initialization.mcpSession();
-		Throwable terminal = this.terminalFailure.get();
-		if (terminal != null) {
-			mcpClientSession.terminate(terminal);
-			return Mono.error(terminal);
-		}
+			McpClientSession mcpClientSession = initialization.mcpSession();
+			Throwable terminal = this.terminalFailure.get();
+			if (terminal != null) {
+				mcpClientSession.terminate(terminal);
+				return Mono.error(terminal);
+			}
 
-		String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
+			String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
 
-		McpSchema.InitializeRequest initializeRequest = McpSchema.InitializeRequest
-			.builder(latestVersion, this.clientCapabilities, this.clientInfo)
-			.build();
+			McpSchema.InitializeRequest initializeRequest = McpSchema.InitializeRequest
+				.builder(latestVersion, this.clientCapabilities, this.clientInfo)
+				.build();
 
-		Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
-				initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF);
+			Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
+					initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF);
 
-		return result.flatMap(initializeResult -> {
-			logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
-					initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
-					initializeResult.instructions());
+			return result.flatMap(initializeResult -> {
+				logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
+						initializeResult.protocolVersion(), initializeResult.capabilities(),
+						initializeResult.serverInfo(), initializeResult.instructions());
 
-			if (!this.protocolVersions.contains(initializeResult.protocolVersion())) {
-				return Mono.error(McpError.builder(-32602)
-					.message("Unsupported protocol version")
-					.data("Unsupported protocol version from the server: " + initializeResult.protocolVersion())
-					.build());
-			}
+				if (!this.protocolVersions.contains(initializeResult.protocolVersion())) {
+					return Mono.error(McpError.builder(-32602)
+						.message("Unsupported protocol version")
+						.data("Unsupported protocol version from the server: " + initializeResult.protocolVersion())
+						.build());
+				}
 
-			return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
-				.contextWrite(
-						c -> c.put(McpAsyncClient.NEGOTIATED_PROTOCOL_VERSION, initializeResult.protocolVersion()))
-				.thenReturn(initializeResult);
-		}).flatMap(initializeResult -> {
-			initialization.cacheResult(initializeResult);
-			return postInitOperation.apply(initialization).thenReturn(initializeResult);
-		}).doOnNext(initialization::complete).onErrorResume(ex -> {
-			initialization.error(ex);
-			return Mono.error(ex);
+				return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
+					.contextWrite(
+							c -> c.put(McpAsyncClient.NEGOTIATED_PROTOCOL_VERSION, initializeResult.protocolVersion()))
+					.thenReturn(initializeResult);
+			}).flatMap(initializeResult -> {
+				initialization.cacheResult(initializeResult);
+				return postInitOperation.apply(initialization).thenReturn(initializeResult);
+			}).doOnNext(initialization::complete).doOnError(initialization::error);
 		});
 	}
 

mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerPostInitializationHookTests.java

@@ -6,6 +6,8 @@
 
 import java.time.Duration;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -14,11 +16,13 @@
 import io.modelcontextprotocol.spec.McpClientSession;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
+import io.modelcontextprotocol.spec.McpTransportTerminatedException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 import reactor.util.context.ContextView;
@@ -164,6 +168,38 @@ void shouldFailInitializationWhenPostInitializationHookFails() {
 		verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
 	}
 
+	@Test
+	void shouldFailInitializationWhenTransportTerminatesDuringPostInitializationHook() throws Exception {
+		var cause = new McpTransportTerminatedException("Transport terminated during post-initialization hook");
+		var hookEntered = new CountDownLatch(1);
+		var hookGate = Sinks.<Void>one();
+
+		when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> Mono.defer(() -> {
+			hookEntered.countDown();
+			return hookGate.asMono();
+		}));
+
+		var initialization = initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
+			.materialize()
+			.toFuture();
+		assertThat(hookEntered.await(1, TimeUnit.SECONDS)).isTrue();
+
+		initializer.handleException(cause);
+
+		try {
+			var signal = initialization.get(1, TimeUnit.SECONDS);
+			assertThat(signal.isOnError()).isTrue();
+			assertThat(signal.getThrowable()).hasCause(cause);
+		}
+		finally {
+			hookGate.tryEmitEmpty();
+		}
+
+		assertThat(initializer.isInitialized()).isFalse();
+		assertThat(initializer.currentInitializationResult()).isNull();
+		verify(mockClientSession).terminate(cause);
+	}
+
 	@Test
 	void shouldNotInvokePostInitializationHookWhenInitializationFails() {
 		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))

mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java

@@ -6,6 +6,8 @@
 
 import java.time.Duration;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
@@ -305,15 +307,19 @@ void shouldHandleOtherExceptions() {
 	}
 
 	@Test
-	void shouldTerminateInProgressInitializationOnTransportTermination() {
+	void shouldTerminateInProgressInitializationOnTransportTermination() throws Exception {
 		var cause = new McpTransportTerminatedException("Transport terminated");
 		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenReturn(Mono.never());
 
-		var subscription = initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
-			.subscribe();
+		var initialization = initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
+			.materialize()
+			.toFuture();
 
 		initializer.handleException(cause);
 
+		var signal = initialization.get(1, TimeUnit.SECONDS);
+		assertThat(signal.isOnError()).isTrue();
+		assertThat(signal.getThrowable()).hasCause(cause);
 		verify(mockClientSession).terminate(cause);
 		assertThat(initializer.isInitialized()).isFalse();
 		assertThat(initializer.currentInitializationResult()).isNull();
@@ -323,7 +329,6 @@ void shouldTerminateInProgressInitializationOnTransportTermination() {
 			.verify();
 
 		verify(mockSessionSupplier, times(1)).apply(any(ContextView.class));
-		subscription.dispose();
 	}
 
 	@Test
@@ -342,6 +347,49 @@ void shouldApplyTerminationThatArrivesBeforeSessionRegistration() {
 		verify(mockClientSession, never()).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any());
 	}
 
+	@Test
+	void shouldTerminateWinnerAndJoinerWhenTerminationArrivesBeforeSessionRegistration() throws Exception {
+		var cause = new McpTransportTerminatedException("Transport terminated before session registration");
+		var supplierEntered = new CountDownLatch(1);
+		var releaseSupplier = new CountDownLatch(1);
+
+		when(mockSessionSupplier.apply(any(ContextView.class))).thenAnswer(invocation -> {
+			supplierEntered.countDown();
+			if (!releaseSupplier.await(5, TimeUnit.SECONDS)) {
+				throw new IllegalStateException("Timed out waiting to release session supplier");
+			}
+			return mockClientSession;
+		});
+
+		var winner = initializer.withInitialization("winner", init -> Mono.just(init.initializeResult()))
+			.subscribeOn(Schedulers.boundedElastic())
+			.materialize()
+			.toFuture();
+		assertThat(supplierEntered.await(1, TimeUnit.SECONDS)).isTrue();
+
+		var joiner = initializer.withInitialization("joiner", init -> Mono.just(init.initializeResult()))
+			.materialize()
+			.toFuture();
+
+		initializer.handleException(cause);
+
+		try {
+			var winnerSignal = winner.get(1, TimeUnit.SECONDS);
+			var joinerSignal = joiner.get(1, TimeUnit.SECONDS);
+
+			assertThat(winnerSignal.isOnError()).isTrue();
+			assertThat(winnerSignal.getThrowable()).hasCause(cause);
+			assertThat(joinerSignal.isOnError()).isTrue();
+			assertThat(joinerSignal.getThrowable()).hasCause(cause);
+		}
+		finally {
+			releaseSupplier.countDown();
+		}
+
+		verify(mockSessionSupplier, times(1)).apply(any(ContextView.class));
+		verify(mockClientSession, never()).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any());
+	}
+
 	@Test
 	void shouldIgnoreGenericTransportExceptionDuringInitialization() {
 		var cause = new McpTransportException("Transport closed");