Skip to content

Commit aad0166

Browse files
committed
Add initialize(InitializeRequest) overload to MCP client
1 parent 30f1adf commit aad0166

8 files changed

Lines changed: 313 additions & 12 deletions

File tree

docs/client.md

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,26 @@ The client provides both synchronous and asynchronous APIs for flexibility in di
131131
.subscribe();
132132
```
133133

134+
### Custom Initialize Request
135+
136+
By default, `initialize()` builds the request from client builder settings (protocol version, capabilities, client info). To control the full initialize payload — including the optional `_meta` field — use the `initialize(InitializeRequest)` overload:
137+
138+
```java
139+
InitializeRequest request = InitializeRequest
140+
.builder(ProtocolVersions.MCP_2025_11_25, client.getClientCapabilities(), client.getClientInfo())
141+
.meta(Map.of("server_id", "proxy-1", "invocation_id", "abc-123"))
142+
.build();
143+
144+
// Call before any other client operation so the custom request is sent.
145+
client.initialize(request);
146+
```
147+
148+
The async client exposes the same overload and returns `Mono<InitializeResult>`.
149+
150+
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.
151+
152+
> **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.
153+
134154
## Client Transport
135155

136156
The transport layer handles the communication between MCP clients and servers, providing different implementations for various use cases. The client transport manages message serialization, connection establishment, and protocol-specific communication patterns.

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

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77
import java.time.Duration;
88
import java.util.ArrayList;
99
import java.util.Collections;
10+
import java.util.HashMap;
1011
import java.util.List;
12+
import java.util.Map;
1113
import java.util.concurrent.atomic.AtomicReference;
1214
import java.util.function.Function;
1315

@@ -270,15 +272,41 @@ public void handleException(Throwable t) {
270272
* @return A Mono that completes with the result of the operation
271273
*/
272274
public <T> Mono<T> withInitialization(String actionName, Function<Initialization, Mono<T>> operation) {
275+
return this.withInitialization(null, actionName, operation);
276+
}
277+
278+
/**
279+
* Utility method to ensure the initialization is established before executing an
280+
* operation, using a caller-provided initialize request.
281+
* @param <T> The type of the result Mono
282+
* @param initializeRequest The initialize request to send; may be null to use the
283+
* default. Only applied when this call triggers the initialization, otherwise
284+
* ignored.
285+
* @param actionName The action to perform when the client is initialized
286+
* @param operation The operation to execute when the client is initialized
287+
* @return A Mono that completes with the result of the operation
288+
*/
289+
public <T> Mono<T> withInitialization(McpSchema.InitializeRequest initializeRequest, String actionName,
290+
Function<Initialization, Mono<T>> operation) {
291+
// Snapshot eagerly so mutating the source _meta map before subscription cannot change what is sent.
292+
McpSchema.InitializeRequest sanitizedRequest = initializeRequest != null ? sanitizeInitializeRequest(initializeRequest) : null;
273293
return Mono.deferContextual(ctx -> {
274294
DefaultInitialization newInit = new DefaultInitialization();
275295
DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
276296

277297
boolean needsToInitialize = previous == null;
278298
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
279299

280-
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
281-
? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await();
300+
Mono<McpSchema.InitializeResult> initializationJob;
301+
if (needsToInitialize) {
302+
initializationJob = this.doInitialize(newInit, sanitizedRequest, this.postInitializationHook, ctx);
303+
}
304+
else {
305+
if (sanitizedRequest != null) {
306+
logger.debug("Custom initialize request ignored; client is already initialized");
307+
}
308+
initializationJob = previous.await();
309+
}
282310

283311
return initializationJob.map(initializeResult -> this.initializationRef.get())
284312
.timeout(this.initializationTimeout)
@@ -292,19 +320,43 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
292320
});
293321
}
294322

323+
private static McpSchema.InitializeRequest sanitizeInitializeRequest(McpSchema.InitializeRequest request) {
324+
if (request.meta() == null) {
325+
return request;
326+
}
327+
return McpSchema.InitializeRequest
328+
.builder(request.protocolVersion(), request.capabilities(), request.clientInfo())
329+
.meta(Collections.unmodifiableMap(new HashMap<>(request.meta())))
330+
.build();
331+
}
332+
333+
private McpSchema.InitializeRequest buildInitializeRequest(McpSchema.InitializeRequest customRequest) {
334+
if (customRequest != null) {
335+
return customRequest;
336+
}
337+
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
338+
return McpSchema.InitializeRequest.builder(latestVersion, this.clientCapabilities, this.clientInfo).build();
339+
}
340+
295341
private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization,
296-
Function<Initialization, Mono<Void>> postInitOperation, ContextView ctx) {
342+
McpSchema.InitializeRequest customRequest, Function<Initialization, Mono<Void>> postInitOperation,
343+
ContextView ctx) {
344+
345+
McpSchema.InitializeRequest initializeRequest = this.buildInitializeRequest(customRequest);
346+
347+
if (!this.protocolVersions.contains(initializeRequest.protocolVersion())) {
348+
McpError error = McpError.builder(-32602)
349+
.message("Unsupported protocol version")
350+
.data("Unsupported protocol version in initialize request: " + initializeRequest.protocolVersion())
351+
.build();
352+
initialization.error(error);
353+
return Mono.error(error);
354+
}
297355

298356
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
299357

300358
McpClientSession mcpClientSession = initialization.mcpSession();
301359

302-
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
303-
304-
McpSchema.InitializeRequest initializeRequest = McpSchema.InitializeRequest
305-
.builder(latestVersion, this.clientCapabilities, this.clientInfo)
306-
.build();
307-
308360
Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
309361
initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF);
310362

@@ -355,4 +407,4 @@ public Mono<?> closeGracefully() {
355407
});
356408
}
357409

358-
}
410+
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -471,6 +471,23 @@ public Mono<McpSchema.InitializeResult> initialize() {
471471
return this.initializer.withInitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
472472
}
473473

474+
/**
475+
* Initializes the client using a caller-provided initialize request.
476+
* <p>
477+
* Use this overload to control the full initialize request payload, including the
478+
* optional {@code _meta} field. Call this method before any other client operation to
479+
* ensure the custom request is sent; lazy initialization triggered by other methods
480+
* uses the default request built from client builder settings.
481+
* @param initializeRequest the initialize request to send
482+
* @return the initialize result
483+
* @see #initialize()
484+
*/
485+
public Mono<McpSchema.InitializeResult> initialize(McpSchema.InitializeRequest initializeRequest) {
486+
Assert.notNull(initializeRequest, "InitializeRequest must not be null");
487+
return this.initializer.withInitialization(initializeRequest, "by explicit API call with custom request",
488+
init -> Mono.just(init.initializeResult()));
489+
}
490+
474491
// --------------------------
475492
// Basic Utilities
476493
// --------------------------

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,21 @@ public McpSchema.InitializeResult initialize() {
190190
return withProvidedContext(this.delegate.initialize()).block();
191191
}
192192

193+
/**
194+
* Initializes the client using a caller-provided initialize request.
195+
* <p>
196+
* Use this overload to control the full initialize request payload, including the
197+
* optional {@code _meta} field. Call this method before any other client operation to
198+
* ensure the custom request is sent; lazy initialization triggered by other methods
199+
* uses the default request built from client builder settings.
200+
* @param initializeRequest the initialize request to send
201+
* @return the initialize result
202+
* @see #initialize()
203+
*/
204+
public McpSchema.InitializeResult initialize(McpSchema.InitializeRequest initializeRequest) {
205+
return withProvidedContext(this.delegate.initialize(initializeRequest)).block();
206+
}
207+
193208
/**
194209
* Send a roots/list_changed notification.
195210
*/

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,11 @@
1010
import java.io.InputStreamReader;
1111
import java.io.OutputStream;
1212
import java.nio.charset.StandardCharsets;
13-
import java.util.List;
1413
import java.util.concurrent.Executors;
1514
import java.util.concurrent.atomic.AtomicBoolean;
1615
import java.util.function.Function;
1716

1817
import io.modelcontextprotocol.json.TypeRef;
19-
import io.modelcontextprotocol.spec.McpError;
2018
import io.modelcontextprotocol.spec.McpSchema;
2119
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
2220
import io.modelcontextprotocol.spec.McpServerSession;

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

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,16 @@
55
package io.modelcontextprotocol.client;
66

77
import java.time.Duration;
8+
import java.util.HashMap;
89
import java.util.List;
10+
import java.util.Map;
911
import java.util.concurrent.atomic.AtomicInteger;
1012
import java.util.concurrent.atomic.AtomicReference;
1113
import java.util.function.Function;
1214

1315
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
1416
import io.modelcontextprotocol.spec.McpClientSession;
17+
import io.modelcontextprotocol.spec.McpError;
1518
import io.modelcontextprotocol.spec.McpSchema;
1619
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
1720
import org.junit.jupiter.api.BeforeEach;
@@ -149,6 +152,111 @@ void shouldUseLatestProtocolVersionInInitializeRequest() {
149152
.verifyComplete();
150153
}
151154

155+
@Test
156+
void shouldUseCustomInitializeRequest() {
157+
AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
158+
159+
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
160+
capturedRequest.set((McpSchema.InitializeRequest) invocation.getArgument(1));
161+
return Mono.just(MOCK_INIT_RESULT);
162+
});
163+
164+
McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
165+
.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
166+
.meta(Map.of("server_id", "proxy-1"))
167+
.build();
168+
169+
StepVerifier
170+
.create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult())))
171+
.assertNext(result -> {
172+
assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0");
173+
assertThat(capturedRequest.get().capabilities()).isEqualTo(CLIENT_CAPABILITIES);
174+
assertThat(capturedRequest.get().clientInfo()).isEqualTo(CLIENT_INFO);
175+
assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1");
176+
})
177+
.verifyComplete();
178+
}
179+
180+
@Test
181+
void shouldDefensivelyCopyMetaFromCustomInitializeRequest() {
182+
Map<String, Object> meta = new HashMap<>();
183+
meta.put("traceId", "abc-123");
184+
meta.put("client", Map.of("name", "test-client"));
185+
186+
McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
187+
.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
188+
.meta(meta)
189+
.build();
190+
191+
AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
192+
193+
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
194+
capturedRequest.set((McpSchema.InitializeRequest) invocation.getArgument(1));
195+
return Mono.just(MOCK_INIT_RESULT);
196+
});
197+
198+
Mono<McpSchema.InitializeResult> initialization = initializer.withInitialization(customRequest, "test",
199+
init -> Mono.just(init.initializeResult()));
200+
meta.put("traceId", "changed");
201+
202+
StepVerifier.create(initialization).expectNext(MOCK_INIT_RESULT).verifyComplete();
203+
204+
assertThat(capturedRequest.get().meta()).containsEntry("traceId", "abc-123")
205+
.containsEntry("client", Map.of("name", "test-client"));
206+
assertThatThrownBy(() -> capturedRequest.get().meta().put("new", "value"))
207+
.isInstanceOf(UnsupportedOperationException.class);
208+
}
209+
210+
@Test
211+
void shouldFailForUnsupportedProtocolVersionInCustomRequest() {
212+
McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest
213+
.builder("999.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
214+
.build();
215+
216+
StepVerifier
217+
.create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult())))
218+
.verifyErrorSatisfies(error -> assertThat(error).isInstanceOf(RuntimeException.class)
219+
.cause()
220+
.isInstanceOf(McpError.class)
221+
.hasMessageContaining("Unsupported protocol version"));
222+
223+
verify(mockSessionSupplier, never()).apply(any(ContextView.class));
224+
verify(mockClientSession, never()).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any());
225+
verify(mockClientSession, never()).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any());
226+
}
227+
228+
@Test
229+
void shouldReuseExistingInitializationWhenCustomRequestProvided() {
230+
AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
231+
232+
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> {
233+
capturedRequest.set(invocation.getArgument(1));
234+
return Mono.just(MOCK_INIT_RESULT);
235+
});
236+
237+
McpSchema.InitializeRequest firstRequest = McpSchema.InitializeRequest
238+
.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
239+
.meta(Map.of("server_id", "proxy-1"))
240+
.build();
241+
McpSchema.InitializeRequest secondRequest = McpSchema.InitializeRequest
242+
.builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO)
243+
.meta(Map.of("server_id", "proxy-2"))
244+
.build();
245+
246+
StepVerifier.create(initializer.withInitialization(firstRequest, "test1", init -> Mono.just("result1")))
247+
.expectNext("result1")
248+
.verifyComplete();
249+
250+
StepVerifier.create(initializer.withInitialization(secondRequest, "test2", init -> Mono.just("result2")))
251+
.expectNext("result2")
252+
.verifyComplete();
253+
254+
verify(mockSessionSupplier, times(1)).apply(any(ContextView.class));
255+
verify(mockClientSession, times(1)).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any());
256+
// Only the first request is sent; the second custom request is ignored.
257+
assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1");
258+
}
259+
152260
@Test
153261
void shouldFailForUnsupportedProtocolVersion() {
154262
McpSchema.InitializeResult unsupportedResult = McpSchema.InitializeResult.builder("999.0.0", // Unsupported

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,20 @@
44

55
package io.modelcontextprotocol.client;
66

7+
import java.util.HashMap;
78
import java.util.List;
9+
import java.util.Map;
10+
import java.util.concurrent.atomic.AtomicReference;
811

12+
import io.modelcontextprotocol.MockMcpClientTransport;
913
import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
1014
import io.modelcontextprotocol.spec.McpClientTransport;
1115
import io.modelcontextprotocol.spec.McpSchema;
16+
import io.modelcontextprotocol.spec.ProtocolVersions;
1217
import org.junit.jupiter.api.Nested;
1318
import org.junit.jupiter.api.Test;
1419
import reactor.core.publisher.Mono;
20+
import reactor.test.StepVerifier;
1521

1622
import static org.assertj.core.api.Assertions.assertThat;
1723
import static org.assertj.core.api.Assertions.assertThatCode;
@@ -24,6 +30,48 @@
2430
*/
2531
class McpAsyncClientTest {
2632

33+
@Nested
34+
class Initialize {
35+
36+
@Test
37+
void customInitializeRequestIsSentOnWire() {
38+
AtomicReference<McpSchema.InitializeRequest> capturedRequest = new AtomicReference<>();
39+
McpSchema.InitializeResult initializeResult = McpSchema.InitializeResult
40+
.builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ServerCapabilities.builder().build(),
41+
McpSchema.Implementation.builder("test-server", "1.0.0").build())
42+
.build();
43+
MockMcpClientTransport transport = new MockMcpClientTransport((mockTransport, message) -> {
44+
if (message instanceof McpSchema.JSONRPCRequest request
45+
&& McpSchema.METHOD_INITIALIZE.equals(request.method())) {
46+
capturedRequest.set((McpSchema.InitializeRequest) request.params());
47+
mockTransport
48+
.simulateIncomingMessage(McpSchema.JSONRPCResponse.result(request.id(), initializeResult));
49+
}
50+
});
51+
52+
Map<String, Object> meta = new HashMap<>();
53+
meta.put("server_id", "proxy-1");
54+
55+
McpSchema.InitializeRequest request = McpSchema.InitializeRequest
56+
.builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().build(),
57+
McpSchema.Implementation.builder("test-client", "1.0.0").build())
58+
.meta(meta)
59+
.build();
60+
61+
McpAsyncClient client = McpClient.async(transport)
62+
.jsonSchemaValidator(mock(JsonSchemaValidator.class))
63+
.build();
64+
65+
Mono<McpSchema.InitializeResult> initialization = client.initialize(request);
66+
meta.put("server_id", "changed");
67+
68+
StepVerifier.create(initialization).expectNext(initializeResult).verifyComplete();
69+
70+
assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1");
71+
}
72+
73+
}
74+
2775
@Nested
2876
class ClientBuilder {
2977

0 commit comments

Comments
 (0)