Skip to content

Commit 65fb680

Browse files
committed
Add initialize(InitializeRequest) overload to MCP client
Expose a caller-controlled initialize request on sync and async clients so _meta and other fields can be sent on the handshake without transport wrappers. Fixes #940
1 parent 30f1adf commit 65fb680

7 files changed

Lines changed: 284 additions & 6 deletions

File tree

docs/client.md

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,24 @@ 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+
134152
## Client Transport
135153

136154
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: 48 additions & 6 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,6 +272,22 @@ 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; must not be null
283+
* @param actionName The action to perform when the client is initialized
284+
* @param operation The operation to execute when the client is initialized
285+
* @return A Mono that completes with the result of the operation
286+
*/
287+
public <T> Mono<T> withInitialization(McpSchema.InitializeRequest initializeRequest, String actionName,
288+
Function<Initialization, Mono<T>> operation) {
289+
McpSchema.InitializeRequest sanitizedRequest = initializeRequest != null
290+
? sanitizeInitializeRequest(initializeRequest) : null;
273291
return Mono.deferContextual(ctx -> {
274292
DefaultInitialization newInit = new DefaultInitialization();
275293
DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
@@ -278,7 +296,7 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
278296
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
279297

280298
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
281-
? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await();
299+
? this.doInitialize(newInit, sanitizedRequest, this.postInitializationHook, ctx) : previous.await();
282300

283301
return initializationJob.map(initializeResult -> this.initializationRef.get())
284302
.timeout(this.initializationTimeout)
@@ -292,18 +310,42 @@ public <T> Mono<T> withInitialization(String actionName, Function<Initialization
292310
});
293311
}
294312

313+
private static McpSchema.InitializeRequest sanitizeInitializeRequest(McpSchema.InitializeRequest request) {
314+
if (request.meta() == null) {
315+
return request;
316+
}
317+
return McpSchema.InitializeRequest
318+
.builder(request.protocolVersion(), request.capabilities(), request.clientInfo())
319+
.meta(Collections.unmodifiableMap(new HashMap<>(request.meta())))
320+
.build();
321+
}
322+
323+
private McpSchema.InitializeRequest buildInitializeRequest(McpSchema.InitializeRequest customRequest) {
324+
if (customRequest != null) {
325+
return customRequest;
326+
}
327+
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
328+
return McpSchema.InitializeRequest.builder(latestVersion, this.clientCapabilities, this.clientInfo).build();
329+
}
330+
295331
private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization,
296-
Function<Initialization, Mono<Void>> postInitOperation, ContextView ctx) {
332+
McpSchema.InitializeRequest customRequest, Function<Initialization, Mono<Void>> postInitOperation,
333+
ContextView ctx) {
297334

298335
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
299336

300337
McpClientSession mcpClientSession = initialization.mcpSession();
301338

302-
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
339+
McpSchema.InitializeRequest initializeRequest = this.buildInitializeRequest(customRequest);
303340

304-
McpSchema.InitializeRequest initializeRequest = McpSchema.InitializeRequest
305-
.builder(latestVersion, this.clientCapabilities, this.clientInfo)
306-
.build();
341+
if (!this.protocolVersions.contains(initializeRequest.protocolVersion())) {
342+
McpError error = McpError.builder(-32602)
343+
.message("Unsupported protocol version")
344+
.data("Unsupported protocol version in initialize request: " + initializeRequest.protocolVersion())
345+
.build();
346+
initialization.error(error);
347+
return Mono.error(error);
348+
}
307349

308350
Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
309351
initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF);

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/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,9 @@
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;
@@ -149,6 +151,99 @@ void shouldUseLatestProtocolVersionInInitializeRequest() {
149151
.verifyComplete();
150152
}
151153

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