LifecycleInitializer recovers from init failures (#549)

Signed-off-by: Daniel Garnier-Moiroux <[email protected]>

Author: Kehrlann

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

@@ -289,6 +289,7 @@ public <T> Mono<T> withIntitialization(String actionName, Function<Initializatio
 			return initializationJob.map(initializeResult -> this.initializationRef.get())
 				.timeout(this.initializationTimeout)
 				.onErrorResume(ex -> {
+					this.initializationRef.compareAndSet(newInit, null);
 					return Mono.error(new RuntimeException("Client failed to initialize " + actionName, ex));
 				})
 				.flatMap(operation);

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

@@ -230,14 +230,26 @@ void shouldHandleConcurrentInitializationRequests() {
 	@Test
 	void shouldHandleInitializationFailure() {
 		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
-			.thenReturn(Mono.error(new RuntimeException("Connection failed")));
+			// fail once
+			.thenReturn(Mono.error(new RuntimeException("Connection failed")))
+			// succeeds on the second call
+			.thenReturn(Mono.just(MOCK_INIT_RESULT));
 
 		StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
 			.expectError(RuntimeException.class)
 			.verify();
 
 		assertThat(initializer.isInitialized()).isFalse();
 		assertThat(initializer.currentInitializationResult()).isNull();
+
+		// The initializer can recover from previous errors
+		StepVerifier
+			.create(initializer.withIntitialization("successful init", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		assertThat(initializer.isInitialized()).isTrue();
+		assertThat(initializer.currentInitializationResult()).isEqualTo(MOCK_INIT_RESULT);
 	}
 
 	@Test