From 3d063e3e4b3b21aaca207bf090fb07777a3d89e3 Mon Sep 17 00:00:00 2001 From: Nicolas QUINQUENEL Date: Mon, 6 Jul 2026 13:19:01 +0200 Subject: [PATCH] Add initialize(InitializeRequest) overload to MCP client --- docs/client.md | 20 ++ .../client/LifecycleInitializer.java | 85 ++++++- .../client/McpAsyncClient.java | 19 ++ .../client/McpSyncClient.java | 17 ++ .../client/LifecycleInitializerTests.java | 217 ++++++++++++++++++ .../client/McpAsyncClientTest.java | 48 ++++ .../client/McpSyncClientTest.java | 43 ++++ 7 files changed, 438 insertions(+), 11 deletions(-) diff --git a/docs/client.md b/docs/client.md index 199a9d34e..be1587925 100644 --- a/docs/client.md +++ b/docs/client.md @@ -131,6 +131,26 @@ The client provides both synchronous and asynchronous APIs for flexibility in di .subscribe(); ``` +### Custom Initialize Request + +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: + +```java +InitializeRequest request = InitializeRequest + .builder(ProtocolVersions.MCP_2025_11_25, client.getClientCapabilities(), client.getClientInfo()) + .meta(Map.of("server_id", "proxy-1", "invocation_id", "abc-123")) + .build(); + +// Call before any other client operation so the custom request is sent. +client.initialize(request); +``` + +The async client exposes the same overload and returns `Mono`. + +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. + +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 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. diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java index f62cd7c71..4d004a6bf 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java @@ -7,7 +7,9 @@ import java.time.Duration; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -93,6 +95,8 @@ class LifecycleInitializer { private final AtomicReference initializationRef = new AtomicReference<>(); + private final AtomicReference storedCustomInitializeRequest = new AtomicReference<>(); + /** * The max timeout to await for the client-server connection to be initialized. */ @@ -257,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(); } } @@ -270,6 +275,26 @@ public void handleException(Throwable t) { * @return A Mono that completes with the result of the operation */ public Mono withInitialization(String actionName, Function> operation) { + return this.withInitialization(null, actionName, operation); + } + + /** + * Utility method to ensure the initialization is established before executing an + * operation, using a caller-provided initialize request. + * @param The type of the result Mono + * @param initializeRequest The initialize request to send; may be null to use the + * default. Only applied when this call triggers the initialization, otherwise + * ignored. + * @param actionName The action to perform when the client is initialized + * @param operation The operation to execute when the client is initialized + * @return A Mono that completes with the result of the operation + */ + public Mono withInitialization(McpSchema.InitializeRequest initializeRequest, String actionName, + Function> operation) { + // Snapshot eagerly so mutating the source _meta map before subscription cannot + // change what is sent. + McpSchema.InitializeRequest sanitizedRequest = initializeRequest != null + ? sanitizeInitializeRequest(initializeRequest) : null; return Mono.deferContextual(ctx -> { DefaultInitialization newInit = new DefaultInitialization(); DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit); @@ -277,8 +302,16 @@ public Mono withInitialization(String actionName, Function initializationJob = needsToInitialize - ? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await(); + Mono initializationJob; + if (needsToInitialize) { + initializationJob = this.doInitialize(newInit, sanitizedRequest, this.postInitializationHook, ctx); + } + else { + if (sanitizedRequest != null) { + logger.debug("Custom initialize request ignored; client is already initialized"); + } + initializationJob = previous.await(); + } return initializationJob.map(initializeResult -> this.initializationRef.get()) .timeout(this.initializationTimeout) @@ -292,19 +325,43 @@ public Mono withInitialization(String actionName, Function(request.meta()))) + .build(); + } + + private McpSchema.InitializeRequest buildInitializeRequest(McpSchema.InitializeRequest customRequest) { + if (customRequest != null) { + return customRequest; + } + String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1); + return McpSchema.InitializeRequest.builder(latestVersion, this.clientCapabilities, this.clientInfo).build(); + } + private Mono doInitialize(DefaultInitialization initialization, - Function> postInitOperation, ContextView ctx) { + McpSchema.InitializeRequest customRequest, Function> postInitOperation, + ContextView ctx) { + + McpSchema.InitializeRequest initializeRequest = this.buildInitializeRequest(customRequest); + + if (!this.protocolVersions.contains(initializeRequest.protocolVersion())) { + McpError error = McpError.builder(-32602) + .message("Unsupported protocol version") + .data("Unsupported protocol version in initialize request: " + initializeRequest.protocolVersion()) + .build(); + initialization.error(error); + return Mono.error(error); + } initialization.setMcpClientSession(this.sessionSupplier.apply(ctx)); McpClientSession mcpClientSession = initialization.mcpSession(); - String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1); - - McpSchema.InitializeRequest initializeRequest = McpSchema.InitializeRequest - .builder(latestVersion, this.clientCapabilities, this.clientInfo) - .build(); - Mono result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE, initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF); @@ -327,6 +384,10 @@ private Mono 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); @@ -341,6 +402,7 @@ public void close() { if (current != null) { current.close(); } + this.storedCustomInitializeRequest.set(null); } /** @@ -350,9 +412,10 @@ 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; }); } -} \ No newline at end of file +} diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java index 945221bd0..63891bc08 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -471,6 +471,25 @@ public Mono initialize() { return this.initializer.withInitialization("by explicit API call", init -> Mono.just(init.initializeResult())); } + /** + * Initializes the client using a caller-provided initialize request. + *

+ * 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. 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() + */ + public Mono initialize(McpSchema.InitializeRequest initializeRequest) { + Assert.notNull(initializeRequest, "InitializeRequest must not be null"); + return this.initializer.withInitialization(initializeRequest, "by explicit API call with custom request", + init -> Mono.just(init.initializeResult())); + } + // -------------------------- // Basic Utilities // -------------------------- diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java index 7e08f83a0..fb980c4b4 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java @@ -190,6 +190,23 @@ public McpSchema.InitializeResult initialize() { return withProvidedContext(this.delegate.initialize()).block(); } + /** + * Initializes the client using a caller-provided initialize request. + *

+ * 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. 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() + */ + public McpSchema.InitializeResult initialize(McpSchema.InitializeRequest initializeRequest) { + return withProvidedContext(this.delegate.initialize(initializeRequest)).block(); + } + /** * Send a roots/list_changed notification. */ diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java index be7b7dc53..bac06738f 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java @@ -5,13 +5,17 @@ package io.modelcontextprotocol.client; import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import io.modelcontextprotocol.client.LifecycleInitializer.Initialization; import io.modelcontextprotocol.spec.McpClientSession; +import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; import org.junit.jupiter.api.BeforeEach; @@ -149,6 +153,111 @@ void shouldUseLatestProtocolVersionInInitializeRequest() { .verifyComplete(); } + @Test + void shouldUseCustomInitializeRequest() { + AtomicReference capturedRequest = new AtomicReference<>(); + + when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> { + capturedRequest.set((McpSchema.InitializeRequest) 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()))) + .assertNext(result -> { + assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0"); + assertThat(capturedRequest.get().capabilities()).isEqualTo(CLIENT_CAPABILITIES); + assertThat(capturedRequest.get().clientInfo()).isEqualTo(CLIENT_INFO); + assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1"); + }) + .verifyComplete(); + } + + @Test + void shouldDefensivelyCopyMetaFromCustomInitializeRequest() { + Map meta = new HashMap<>(); + meta.put("traceId", "abc-123"); + meta.put("client", Map.of("name", "test-client")); + + McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest + .builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO) + .meta(meta) + .build(); + + AtomicReference capturedRequest = new AtomicReference<>(); + + when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any())).thenAnswer(invocation -> { + capturedRequest.set((McpSchema.InitializeRequest) invocation.getArgument(1)); + return Mono.just(MOCK_INIT_RESULT); + }); + + Mono initialization = initializer.withInitialization(customRequest, "test", + init -> Mono.just(init.initializeResult())); + meta.put("traceId", "changed"); + + StepVerifier.create(initialization).expectNext(MOCK_INIT_RESULT).verifyComplete(); + + assertThat(capturedRequest.get().meta()).containsEntry("traceId", "abc-123") + .containsEntry("client", Map.of("name", "test-client")); + assertThatThrownBy(() -> capturedRequest.get().meta().put("new", "value")) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void shouldFailForUnsupportedProtocolVersionInCustomRequest() { + McpSchema.InitializeRequest customRequest = McpSchema.InitializeRequest + .builder("999.0.0", CLIENT_CAPABILITIES, CLIENT_INFO) + .build(); + + StepVerifier + .create(initializer.withInitialization(customRequest, "test", init -> Mono.just(init.initializeResult()))) + .verifyErrorSatisfies(error -> assertThat(error).isInstanceOf(RuntimeException.class) + .cause() + .isInstanceOf(McpError.class) + .hasMessageContaining("Unsupported protocol version")); + + verify(mockSessionSupplier, never()).apply(any(ContextView.class)); + verify(mockClientSession, never()).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()); + verify(mockClientSession, never()).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()); + } + + @Test + void shouldReuseExistingInitializationWhenCustomRequestProvided() { + AtomicReference 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 firstRequest = McpSchema.InitializeRequest + .builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO) + .meta(Map.of("server_id", "proxy-1")) + .build(); + McpSchema.InitializeRequest secondRequest = McpSchema.InitializeRequest + .builder("2.0.0", CLIENT_CAPABILITIES, CLIENT_INFO) + .meta(Map.of("server_id", "proxy-2")) + .build(); + + StepVerifier.create(initializer.withInitialization(firstRequest, "test1", init -> Mono.just("result1"))) + .expectNext("result1") + .verifyComplete(); + + StepVerifier.create(initializer.withInitialization(secondRequest, "test2", init -> Mono.just("result2"))) + .expectNext("result2") + .verifyComplete(); + + verify(mockSessionSupplier, times(1)).apply(any(ContextView.class)); + verify(mockClientSession, times(1)).sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()); + // Only the first request is sent; the second custom request is ignored. + assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1"); + } + @Test void shouldFailForUnsupportedProtocolVersion() { McpSchema.InitializeResult unsupportedResult = McpSchema.InitializeResult.builder("999.0.0", // Unsupported @@ -283,6 +392,114 @@ void shouldHandleTransportSessionNotFoundException() { verify(mockSessionSupplier, times(2)).apply(any(ContextView.class)); } + @Test + void shouldReuseCustomInitializeRequestOnReinitialization() { + List 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 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 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 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 diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java index dea7d42e9..37eace028 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTest.java @@ -4,14 +4,20 @@ package io.modelcontextprotocol.client; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import io.modelcontextprotocol.MockMcpClientTransport; import io.modelcontextprotocol.json.schema.JsonSchemaValidator; import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatCode; @@ -24,6 +30,48 @@ */ class McpAsyncClientTest { + @Nested + class Initialize { + + @Test + void customInitializeRequestIsSentOnWire() { + AtomicReference capturedRequest = new AtomicReference<>(); + McpSchema.InitializeResult initializeResult = McpSchema.InitializeResult + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ServerCapabilities.builder().build(), + McpSchema.Implementation.builder("test-server", "1.0.0").build()) + .build(); + MockMcpClientTransport transport = new MockMcpClientTransport((mockTransport, message) -> { + if (message instanceof McpSchema.JSONRPCRequest request + && McpSchema.METHOD_INITIALIZE.equals(request.method())) { + capturedRequest.set((McpSchema.InitializeRequest) request.params()); + mockTransport + .simulateIncomingMessage(McpSchema.JSONRPCResponse.result(request.id(), initializeResult)); + } + }); + + Map meta = new HashMap<>(); + meta.put("server_id", "proxy-1"); + + McpSchema.InitializeRequest request = McpSchema.InitializeRequest + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().build(), + McpSchema.Implementation.builder("test-client", "1.0.0").build()) + .meta(meta) + .build(); + + McpAsyncClient client = McpClient.async(transport) + .jsonSchemaValidator(mock(JsonSchemaValidator.class)) + .build(); + + Mono initialization = client.initialize(request); + meta.put("server_id", "changed"); + + StepVerifier.create(initialization).expectNext(initializeResult).verifyComplete(); + + assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1"); + } + + } + @Nested class ClientBuilder { diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpSyncClientTest.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpSyncClientTest.java index 9790dea6a..0fb557b53 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/McpSyncClientTest.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/McpSyncClientTest.java @@ -4,11 +4,16 @@ package io.modelcontextprotocol.client; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import io.modelcontextprotocol.MockMcpClientTransport; import io.modelcontextprotocol.json.schema.JsonSchemaValidator; import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -21,6 +26,44 @@ */ class McpSyncClientTest { + @Nested + class Initialize { + + @Test + void customInitializeRequestIsSentOnWire() { + AtomicReference capturedRequest = new AtomicReference<>(); + McpSchema.InitializeResult initializeResult = McpSchema.InitializeResult + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ServerCapabilities.builder().build(), + McpSchema.Implementation.builder("test-server", "1.0.0").build()) + .build(); + MockMcpClientTransport transport = new MockMcpClientTransport((mockTransport, message) -> { + if (message instanceof McpSchema.JSONRPCRequest request + && McpSchema.METHOD_INITIALIZE.equals(request.method())) { + capturedRequest.set((McpSchema.InitializeRequest) request.params()); + mockTransport + .simulateIncomingMessage(McpSchema.JSONRPCResponse.result(request.id(), initializeResult)); + } + }); + + Map meta = new HashMap<>(); + meta.put("server_id", "proxy-1"); + + McpSchema.InitializeRequest request = McpSchema.InitializeRequest + .builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().build(), + McpSchema.Implementation.builder("test-client", "1.0.0").build()) + .meta(meta) + .build(); + + McpSyncClient client = McpClient.sync(transport) + .jsonSchemaValidator(mock(JsonSchemaValidator.class)) + .build(); + + assertThat(client.initialize(request)).isEqualTo(initializeResult); + assertThat(capturedRequest.get().meta()).containsEntry("server_id", "proxy-1"); + } + + } + @Nested class ClientCapabilities {