Preserve init request on reinitialization

Author: nquinquenel

docs/client.md

@@ -149,7 +149,7 @@ The async client exposes the same overload and returns `Mono<InitializeResult>`.
 
 If another client method triggers lazy initialization first, the default request is sent instead. Call `initialize(request)` before `listTools()`, `callTool()`, or similar operations when custom initialize metadata is required.
 
-> **Note:** The custom request is only used for the initial handshake. If the transport session is later re-established (for example after a `McpTransportSessionNotFoundException`), the client re-initializes with the default request and the custom `_meta` is not resent.
+After a successful `initialize(InitializeRequest)`, the client remembers that request and resends it when the transport session is re-established (for example after a `McpTransportSessionNotFoundException`). The stored request is cleared when the client is closed.
 
 ## Client Transport
 

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

@@ -95,6 +95,8 @@ class LifecycleInitializer {
 
 	private final AtomicReference<DefaultInitialization> initializationRef = new AtomicReference<>();
 
+	private final AtomicReference<McpSchema.InitializeRequest> storedCustomInitializeRequest = new AtomicReference<>();
+
 	/**
 	 * The max timeout to await for the client-server connection to be initialized.
 	 */
@@ -259,7 +261,8 @@ public void handleException(Throwable t) {
 			}
 			// Providing an empty operation since we are only interested in triggering
 			// the implicit initialization step.
-			this.withInitialization("re-initializing", result -> Mono.empty()).subscribe();
+			this.withInitialization(this.storedCustomInitializeRequest.get(), "re-initializing", result -> Mono.empty())
+				.subscribe();
 		}
 	}
 
@@ -381,6 +384,10 @@ private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization init
 		}).flatMap(initializeResult -> {
 			initialization.cacheResult(initializeResult);
 			return postInitOperation.apply(initialization).thenReturn(initializeResult);
+		}).doOnNext(initializeResult -> {
+			if (customRequest != null) {
+				this.storedCustomInitializeRequest.set(customRequest);
+			}
 		}).doOnNext(initialization::complete).onErrorResume(ex -> {
 			initialization.error(ex);
 			return Mono.error(ex);
@@ -395,6 +402,7 @@ public void close() {
 		if (current != null) {
 			current.close();
 		}
+		this.storedCustomInitializeRequest.set(null);
 	}
 
 	/**
@@ -404,6 +412,7 @@ public void close() {
 	public Mono<?> closeGracefully() {
 		return Mono.defer(() -> {
 			DefaultInitialization current = this.initializationRef.getAndSet(null);
+			this.storedCustomInitializeRequest.set(null);
 			Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
 			return sessionClose;
 		});

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

@@ -477,7 +477,9 @@ public Mono<McpSchema.InitializeResult> initialize() {
 	 * Use this overload to control the full initialize request payload, including the
 	 * optional {@code _meta} field. Call this method before any other client operation to
 	 * ensure the custom request is sent; lazy initialization triggered by other methods
-	 * uses the default request built from client builder settings.
+	 * uses the default request built from client builder settings. After a successful
+	 * call, the custom request is remembered and resent on transport session recovery; it
+	 * is cleared when the client is closed.
 	 * @param initializeRequest the initialize request to send
 	 * @return the initialize result
 	 * @see #initialize()

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

@@ -5,6 +5,7 @@
 package io.modelcontextprotocol.client;
 
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -391,6 +392,114 @@ void shouldHandleTransportSessionNotFoundException() {
 		verify(mockSessionSupplier, times(2)).apply(any(ContextView.class));
 	}
 
+	@Test
+	void shouldReuseCustomInitializeRequestOnReinitialization() {
+		List<McpSchema.InitializeRequest> capturedRequests = new ArrayList<>();
+
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
+			capturedRequests.add(invocation.getArgument(1));
+			return Mono.just(MOCK_INIT_RESULT);
+		});
+
+		McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
+			.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
+			.meta(Map.of("server_id", "proxy-1"))
+			.build();
+
+		StepVerifier
+			.create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		initializer.handleException(new McpTransportSessionNotFoundException("Session not found"));
+
+		assertThat(capturedRequests).hasSize(2);
+		assertThat(capturedRequests.get(0).meta()).containsEntry("server_id", "proxy-1");
+		assertThat(capturedRequests.get(1).meta()).containsEntry("server_id", "proxy-1");
+		verify(mockSessionSupplier, times(2)).apply(any(ContextView.class));
+	}
+
+	@Test
+	void shouldUseDefaultRequestOnReinitializationWhenNoCustomRequestStored() {
+		List<McpSchema.InitializeRequest> capturedRequests = new ArrayList<>();
+
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
+			capturedRequests.add(invocation.getArgument(1));
+			return Mono.just(MOCK_INIT_RESULT);
+		});
+
+		StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		initializer.handleException(new McpTransportSessionNotFoundException("Session not found"));
+
+		assertThat(capturedRequests).hasSize(2);
+		assertThat(capturedRequests.get(1).protocolVersion()).isEqualTo("2.0.0");
+		assertThat(capturedRequests.get(1).capabilities()).isEqualTo(CLIENT_CAPABILITIES);
+		assertThat(capturedRequests.get(1).clientInfo()).isEqualTo(CLIENT_INFO);
+		assertThat(capturedRequests.get(1).meta()).isNull();
+	}
+
+	@Test
+	void shouldClearStoredCustomRequestOnClose() {
+		AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
+
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
+			capturedRequest.set(invocation.getArgument(1));
+			return Mono.just(MOCK_INIT_RESULT);
+		});
+
+		McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
+			.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
+			.meta(Map.of("server_id", "proxy-1"))
+			.build();
+
+		StepVerifier
+			.create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		initializer.close();
+
+		StepVerifier
+			.create(initializer.withInitialization("test after close", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		assertThat(capturedRequest.get().meta()).isNull();
+		assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0");
+	}
+
+	@Test
+	void shouldClearStoredCustomRequestOnCloseGracefully() {
+		AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
+
+		when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
+			capturedRequest.set(invocation.getArgument(1));
+			return Mono.just(MOCK_INIT_RESULT);
+		});
+
+		McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
+			.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
+			.meta(Map.of("server_id", "proxy-1"))
+			.build();
+
+		StepVerifier
+			.create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		StepVerifier.create(initializer.closeGracefully()).verifyComplete();
+
+		StepVerifier
+			.create(initializer.withInitialization("test after close", init -> Mono.just(init.initializeResult())))
+			.expectNext(MOCK_INIT_RESULT)
+			.verifyComplete();
+
+		assertThat(capturedRequest.get().meta()).isNull();
+	}
+
 	@Test
 	void shouldHandleOtherExceptions() {
 		// Simulate a successful initialization first