fix: close terminal initialization races

Author: DragonFSKY

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

@@ -306,6 +306,14 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
 			DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
 
 			boolean needsToInitialize = previous == null;
+			DefaultInitialization activeInitialization = needsToInitialize ? newInit : previous;
+			// Complete the handoff with handleException(): if termination won before
+			// registration, this check publishes it; if registration won first, the
+			// exception handler observes the active initialization.
+			Throwable terminalAfterRegistration = this.terminalFailure.get();
+			if (terminalAfterRegistration != null) {
+				activeInitialization.terminate(terminalAfterRegistration);
+			}
 			logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
 
 			Mono<McpSchema.InitializeResult> initializationJob;
@@ -317,10 +325,10 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
 					.doInitialize(newInit, this.postInitializationHook, ctx)
 					.onErrorComplete()
 					.then(Mono.never());
-				initializationJob = Mono.firstWithSignal(newInit.await(), initializationWork);
+				initializationJob = Mono.firstWithSignal(activeInitialization.await(), initializationWork);
 			}
 			else {
-				initializationJob = previous.await();
+				initializationJob = activeInitialization.await();
 			}
 
 			return initializationJob.map(initializeResult -> this.initializationRef.get())
@@ -376,8 +384,8 @@ private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization init
 			}).flatMap(initializeResult -> {
 				initialization.cacheResult(initializeResult);
 				return postInitOperation.apply(initialization).thenReturn(initializeResult);
-			}).doOnNext(initialization::complete).doOnError(initialization::error);
-		});
+			});
+		}).doOnNext(initialization::complete).doOnError(initialization::error);
 	}
 
 	/**

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

@@ -390,6 +390,43 @@ void shouldTerminateWinnerAndJoinerWhenTerminationArrivesBeforeSessionRegistrati
 		verify(mockClientSession, never()).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any());
 	}
 
+	@Test
+	void shouldShareSynchronousSessionSupplierFailureBetweenWinnerAndJoiner() throws Exception {
+		var cause = new IllegalStateException("Session supplier failed");
+		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");
+			}
+			throw cause;
+		});
+
+		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();
+
+		releaseSupplier.countDown();
+
+		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);
+		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");