From b4fef5241fa81dda87138c98566014a090250f06 Mon Sep 17 00:00:00 2001 From: Kevin Stanton Date: Sun, 10 Aug 2025 00:18:05 -0500 Subject: [PATCH] feat: Add transport context extraction support to all MCP servers (#477) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add McpTransportContextExtractor to WebFlux/WebMVC SSE and Streamable transport providers - Enable extraction of HTTP transport metadata (headers, etc.) for use during request processing - Pass transport context through reactive context chain using McpTransportContext.KEY - Add contextExtractor() builder methods for configuring custom extractors - Update HttpServlet transport providers with same context extraction capability - Modify McpServerSession to properly propagate transport context to handlers - Add test coverage with TEST_CONTEXT_EXTRACTOR in integration tests This allows MCP feature implementations to access HTTP transport level metadata that was present at request time, enabling use cases like authentication, request tracing, and custom header processing. Signed-off-by: Christian Tzolov Co-authored-by: Christian Tzolov Co-authored-by: Dariusz Jędrzejczyk --- .../WebFluxSseServerTransportProvider.java | 57 ++++++++++++++++- ...FluxStreamableServerTransportProvider.java | 13 +++- .../WebFluxSseIntegrationTests.java | 8 +++ .../WebFluxStreamableIntegrationTests.java | 8 +++ .../WebMvcSseServerTransportProvider.java | 60 ++++++++++++++++- .../server/WebMvcSseIntegrationTests.java | 7 ++ .../WebMvcSseSyncServerTransportTests.java | 5 +- .../WebMvcStreamableIntegrationTests.java | 7 ++ ...stractMcpClientServerIntegrationTests.java | 63 +++++++++++++++++- ...HttpServletSseServerTransportProvider.java | 56 +++++++++++++++- .../spec/McpServerSession.java | 36 ++++++++--- ...stractMcpClientServerIntegrationTests.java | 64 ++++++++++++++++++- .../HttpServletSseIntegrationTests.java | 7 ++ ...HttpServletStreamableIntegrationTests.java | 7 ++ 14 files changed, 374 insertions(+), 24 deletions(-) diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java index aaf7bab..ead7380 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java @@ -11,6 +11,10 @@ import java.util.concurrent.ConcurrentHashMap; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; + +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerSession; @@ -115,6 +119,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv */ private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** * Flag indicating if the transport is shutting down. */ @@ -194,15 +200,38 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv @Deprecated public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Constructs a new WebFlux SSE server transport provider instance. + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of MCP messages. Must not be null. + * @param baseUrl webflux message base path + * @param messageEndpoint The endpoint URI where clients should send their JSON-RPC + * messages. This endpoint will be communicated to clients during SSE connection + * setup. Must not be null. + * @param sseEndpoint The SSE endpoint path. Must not be null. + * @param keepAliveInterval The interval for sending keep-alive pings to clients. + * @param contextExtractor The context extractor to use for extracting MCP transport + * context from HTTP requests. Must not be null. + * @throws IllegalArgumentException if either parameter is null + */ + private WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.notNull(baseUrl, "Message base path must not be null"); Assert.notNull(messageEndpoint, "Message endpoint must not be null"); Assert.notNull(sseEndpoint, "SSE endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; this.routerFunction = RouterFunctions.route() .GET(this.sseEndpoint, this::handleSseConnection) .POST(this.messageEndpoint, this::handleMessage) @@ -315,6 +344,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down"); } + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + return ServerResponse.ok() .contentType(MediaType.TEXT_EVENT_STREAM) .body(Flux.>create(sink -> { @@ -336,7 +367,7 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv logger.debug("Session {} cancelled", sessionId); sessions.remove(sessionId); }); - }), ServerSentEvent.class); + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), ServerSentEvent.class); } /** @@ -370,6 +401,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv .bodyValue(new McpError("Session not found: " + request.queryParam("sessionId").get())); } + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + return request.bodyToMono(String.class).flatMap(body -> { try { McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); @@ -386,7 +419,7 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv logger.error("Failed to deserialize message: {}", e.getMessage()); return ServerResponse.badRequest().bodyValue(new McpError("Invalid message format")); } - }); + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)); } private class WebFluxMcpSessionTransport implements McpServerTransport { @@ -458,6 +491,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv private Duration keepAliveInterval; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + /** * Sets the ObjectMapper to use for JSON serialization/deserialization of MCP * messages. @@ -519,6 +554,22 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv return this; } + /** + * Sets the context extractor that allows providing the MCP feature + * implementations to inspect HTTP transport level metadata that was present at + * HTTP request processing time. This allows to extract custom headers and other + * useful data for use during execution later on in the process. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Builds a new instance of {@link WebFluxSseServerTransportProvider} with the * configured settings. @@ -530,7 +581,7 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv Assert.notNull(messageEndpoint, "Message endpoint must be set"); return new WebFluxSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java index f3f6c2c..963a502 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java @@ -191,7 +191,9 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe String lastId = request.headers().asHttpHeaders().getFirst(HttpHeaders.LAST_EVENT_ID); return ServerResponse.ok() .contentType(MediaType.TEXT_EVENT_STREAM) - .body(session.replay(lastId), ServerSentEvent.class); + .body(session.replay(lastId) + .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), + ServerSentEvent.class); } return ServerResponse.ok() @@ -202,7 +204,9 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session .listeningStream(sessionTransport); sink.onDispose(listeningStream::close); - }), ServerSentEvent.class); + // TODO Clarify why the outer context is not present in the + // Flux.create sink? + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), ServerSentEvent.class); }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)); } @@ -282,7 +286,10 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe return true; }).contextWrite(sink.contextView()).subscribe(); sink.onCancel(streamSubscription); - }), ServerSentEvent.class); + // TODO Clarify why the outer context is not present in the + // Flux.create sink? + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), + ServerSentEvent.class); } else { return ServerResponse.badRequest().bodyValue(new McpError("Unknown message type")); diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java index 6140fe4..c8dc6e9 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java @@ -13,6 +13,7 @@ import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.web.reactive.function.client.WebClient; import org.springframework.web.reactive.function.server.RouterFunctions; +import org.springframework.web.reactive.function.server.ServerRequest; import com.fasterxml.jackson.databind.ObjectMapper; @@ -22,6 +23,7 @@ import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport; import io.modelcontextprotocol.server.McpServer; import io.modelcontextprotocol.server.McpServer.AsyncSpecification; import io.modelcontextprotocol.server.McpServer.SingleSessionSyncSpecification; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.server.TestUtil; import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider; import reactor.netty.DisposableServer; @@ -40,6 +42,11 @@ class WebFluxSseIntegrationTests extends AbstractMcpClientServerIntegrationTests private WebFluxSseServerTransportProvider mcpServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Override protected void prepareClients(int port, String mcpEndpoint) { @@ -75,6 +82,7 @@ class WebFluxSseIntegrationTests extends AbstractMcpClientServerIntegrationTests .objectMapper(new ObjectMapper()) .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) .sseEndpoint(CUSTOM_SSE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .build(); HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransportProvider.getRouterFunction()); diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java index 9eba0e5..a7aac0f 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java @@ -13,6 +13,7 @@ import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.web.reactive.function.client.WebClient; import org.springframework.web.reactive.function.server.RouterFunctions; +import org.springframework.web.reactive.function.server.ServerRequest; import com.fasterxml.jackson.databind.ObjectMapper; @@ -22,6 +23,7 @@ import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport import io.modelcontextprotocol.server.McpServer; import io.modelcontextprotocol.server.McpServer.AsyncSpecification; import io.modelcontextprotocol.server.McpServer.SyncSpecification; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.server.TestUtil; import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider; import reactor.netty.DisposableServer; @@ -38,6 +40,11 @@ class WebFluxStreamableIntegrationTests extends AbstractMcpClientServerIntegrati private WebFluxStreamableServerTransportProvider mcpStreamableServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Override protected void prepareClients(int port, String mcpEndpoint) { @@ -71,6 +78,7 @@ class WebFluxStreamableIntegrationTests extends AbstractMcpClientServerIntegrati this.mcpStreamableServerTransportProvider = WebFluxStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .build(); HttpHandler httpHandler = RouterFunctions diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java index ff452ca..6e92cf1 100644 --- a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java @@ -13,6 +13,10 @@ import java.util.concurrent.locks.ReentrantLock; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; + +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerTransport; @@ -106,6 +110,8 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi */ private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** * Flag indicating if the transport is shutting down. */ @@ -177,7 +183,7 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi * messages via HTTP POST. This endpoint will be communicated to clients through the * SSE connection's initial endpoint event. * @param sseEndpoint The endpoint URI where clients establish their SSE connections. - * * @param keepAliveInterval The interval for sending keep-alive messages to + * @param keepAliveInterval The interval for sending keep-alive messages to clients. * @throws IllegalArgumentException if any parameter is null * @deprecated Use the builder {@link #builder()} instead for better configuration * options. @@ -185,15 +191,39 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi @Deprecated public WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Constructs a new WebMvcSseServerTransportProvider instance. + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of messages. + * @param baseUrl The base URL for the message endpoint, used to construct the full + * endpoint URL for clients. + * @param messageEndpoint The endpoint URI where clients should send their JSON-RPC + * messages via HTTP POST. This endpoint will be communicated to clients through the + * SSE connection's initial endpoint event. + * @param sseEndpoint The endpoint URI where clients establish their SSE connections. + * @param keepAliveInterval The interval for sending keep-alive messages to clients. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @throws IllegalArgumentException if any parameter is null + */ + private WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.notNull(baseUrl, "Message base URL must not be null"); Assert.notNull(messageEndpoint, "Message endpoint must not be null"); Assert.notNull(sseEndpoint, "SSE endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; this.routerFunction = RouterFunctions.route() .GET(this.sseEndpoint, this::handleSseConnection) .POST(this.messageEndpoint, this::handleMessage) @@ -367,11 +397,17 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi } try { + final McpTransportContext transportContext = this.contextExtractor.extract(request, + new DefaultMcpTransportContext()); + String body = request.body(String.class); McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); // Process the message through the session's handle method - session.handle(message).block(); // Block for WebMVC compatibility + session.handle(message).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); // Block + // for + // WebMVC + // compatibility return ServerResponse.ok().build(); } @@ -517,6 +553,8 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi private Duration keepAliveInterval; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + /** * Sets the JSON object mapper to use for message serialization/deserialization. * @param objectMapper The object mapper to use @@ -576,6 +614,22 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi return this; } + /** + * Sets the context extractor that allows providing the MCP feature + * implementations to inspect HTTP transport level metadata that was present at + * HTTP request processing time. This allows to extract custom headers and other + * useful data for use during execution later on in the process. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Builds a new instance of WebMvcSseServerTransportProvider with the configured * settings. @@ -587,7 +641,7 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi throw new IllegalStateException("MessageEndpoint must be set"); } return new WebMvcSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java index 5d04835..8cb2973 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java @@ -17,6 +17,7 @@ import org.springframework.context.annotation.Configuration; import org.springframework.web.reactive.function.client.WebClient; import org.springframework.web.servlet.config.annotation.EnableWebMvc; import org.springframework.web.servlet.function.RouterFunction; +import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.function.ServerResponse; import com.fasterxml.jackson.databind.ObjectMapper; @@ -39,6 +40,11 @@ class WebMvcSseIntegrationTests extends AbstractMcpClientServerIntegrationTests private WebMvcSseServerTransportProvider mcpServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Override protected void prepareClients(int port, String mcpEndpoint) { @@ -60,6 +66,7 @@ class WebMvcSseIntegrationTests extends AbstractMcpClientServerIntegrationTests return WebMvcSseServerTransportProvider.builder() .objectMapper(new ObjectMapper()) .messageEndpoint(MESSAGE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .build(); } diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java index 7e49ddf..101a067 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java @@ -36,7 +36,10 @@ class WebMvcSseSyncServerTransportTests extends AbstractMcpSyncServerTests { @Bean public WebMvcSseServerTransportProvider webMvcSseServerTransportProvider() { - return new WebMvcSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT); + return WebMvcSseServerTransportProvider.builder() + .objectMapper(new ObjectMapper()) + .messageEndpoint(MESSAGE_ENDPOINT) + .build(); } @Bean diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java index 16012e7..2f4c651 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java @@ -15,6 +15,7 @@ import org.junit.jupiter.api.Timeout; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.web.reactive.function.client.WebClient; +import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.config.annotation.EnableWebMvc; import org.springframework.web.servlet.function.RouterFunction; import org.springframework.web.servlet.function.ServerResponse; @@ -39,6 +40,11 @@ class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegratio private WebMvcStreamableServerTransportProvider mcpServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Configuration @EnableWebMvc static class TestConfig { @@ -47,6 +53,7 @@ class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegratio public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportProvider() { return WebMvcStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) .build(); } diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java index 8e041d9..5246c1e 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java @@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; @@ -36,6 +37,7 @@ import io.modelcontextprotocol.server.McpServer; import io.modelcontextprotocol.server.McpServerFeatures; import io.modelcontextprotocol.server.McpSyncServer; import io.modelcontextprotocol.server.McpSyncServerExchange; +import io.modelcontextprotocol.server.McpTransportContext; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolResult; @@ -55,6 +57,7 @@ import io.modelcontextprotocol.spec.McpSchema.Role; import io.modelcontextprotocol.spec.McpSchema.Root; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import io.modelcontextprotocol.spec.McpSchema.Tool; +import io.modelcontextprotocol.util.Utils; import net.javacrumbs.jsonunit.core.Option; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -750,6 +753,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { var clientBuilder = clientBuilders.get(clientType); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) @@ -763,7 +767,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { .GET() .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); - assertThat(responseBody).isNotBlank(); + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); } catch (Exception e) { e.printStackTrace(); @@ -786,6 +790,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); assertThat(response).isNotNull().isEqualTo(callResponse); } @@ -829,6 +834,62 @@ public abstract class AbstractMcpClientServerIntegrationTests { mcpServer.close(); } + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testToolCallSuccessWithTranportContextExtraction(String clientType) { + + var clientBuilder = clientBuilders.get(clientType); + + var transportContextIsNull = new AtomicBoolean(false); + var transportContextIsEmpty = new AtomicBoolean(false); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); + + var expectedCallResponse = new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=value")), null); + McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() + .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) + .callHandler((exchange, request) -> { + + McpTransportContext transportContext = exchange.transportContext(); + transportContextIsNull.set(transportContext == null); + transportContextIsEmpty.set(transportContext.equals(McpTransportContext.EMPTY)); + String ctxValue = (String) transportContext.get("important"); + + try { + String responseBody = "TOOL RESPONSE"; + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); + } + catch (Exception e) { + e.printStackTrace(); + } + + return new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=" + ctxValue)), null); + }) + .build(); + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + + assertThat(transportContextIsNull.get()).isFalse(); + assertThat(transportContextIsEmpty.get()).isFalse(); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); + assertThat(response).isNotNull().isEqualTo(expectedCallResponse); + } + + mcpServer.close(); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient", "webflux" }) void testToolListChangeHandlingSuccess(String clientType) { diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java index ceeea31..582120e 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java @@ -16,6 +16,9 @@ import java.util.concurrent.atomic.AtomicBoolean; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerSession; @@ -102,6 +105,8 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement /** Map of active client sessions, keyed by session ID */ private final Map sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** Flag indicating if the transport is in the process of shutting down */ private final AtomicBoolean isClosing = new AtomicBoolean(false); @@ -144,7 +149,7 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement @Deprecated public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint) { - this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null); + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null, (serverRequest, context) -> context); } /** @@ -163,11 +168,38 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement @Deprecated public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Creates a new HttpServletSseServerTransportProvider instance with a custom SSE + * endpoint. + * @param objectMapper The JSON object mapper to use for message + * serialization/deserialization + * @param baseUrl The base URL for the server transport + * @param messageEndpoint The endpoint path where clients will send their messages + * @param sseEndpoint The endpoint path where clients will establish SSE connections + * @param keepAliveInterval The interval for keep-alive pings, or null to disable + * keep-alive functionality + * @param contextExtractor The extractor for transport context from the request. + * @deprecated Use the builder {@link #builder()} instead for better configuration + * options. + */ + private HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { + + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.notNull(messageEndpoint, "messageEndpoint must not be null"); + Assert.notNull(sseEndpoint, "sseEndpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; if (keepAliveInterval != null) { @@ -339,10 +371,13 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement body.append(line); } + final McpTransportContext transportContext = this.contextExtractor.extract(request, + new DefaultMcpTransportContext()); McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body.toString()); // Process the message through the session's handle method - session.handle(message).block(); // Block for Servlet compatibility + // Block for Servlet compatibility + session.handle(message).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); response.setStatus(HttpServletResponse.SC_OK); } @@ -534,6 +569,8 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement private String sseEndpoint = DEFAULT_SSE_ENDPOINT; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + private Duration keepAliveInterval; /** @@ -583,6 +620,19 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement return this; } + /** + * Sets the context extractor for extracting transport context from the request. + * @param contextExtractor The context extractor to use. Must not be null. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public HttpServletSseServerTransportProvider.Builder contextExtractor( + McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Sets the interval for keep-alive pings. *

@@ -609,7 +659,7 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement throw new IllegalStateException("MessageEndpoint must be set"); } return new HttpServletSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java index 62985dc..e562ca0 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java @@ -198,7 +198,9 @@ public class McpServerSession implements McpLoggableSession { * @return a Mono that completes when the message is processed */ public Mono handle(McpSchema.JSONRPCMessage message) { - return Mono.defer(() -> { + return Mono.deferContextual(ctx -> { + McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); + // TODO handle errors for communication to without initialization happening // first if (message instanceof McpSchema.JSONRPCResponse response) { @@ -214,7 +216,7 @@ public class McpServerSession implements McpLoggableSession { } else if (message instanceof McpSchema.JSONRPCRequest request) { logger.debug("Received request: {}", request); - return handleIncomingRequest(request).onErrorResume(error -> { + return handleIncomingRequest(request, transportContext).onErrorResume(error -> { var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, error.getMessage(), null)); @@ -227,7 +229,7 @@ public class McpServerSession implements McpLoggableSession { // happening first logger.debug("Received notification: {}", notification); // TODO: in case of error, should the POST request be signalled? - return handleIncomingNotification(notification) + return handleIncomingNotification(notification, transportContext) .doOnError(error -> logger.error("Error handling notification: {}", error.getMessage())); } else { @@ -240,9 +242,11 @@ public class McpServerSession implements McpLoggableSession { /** * Handles an incoming JSON-RPC request by routing it to the appropriate handler. * @param request The incoming JSON-RPC request + * @param transportContext * @return A Mono containing the JSON-RPC response */ - private Mono handleIncomingRequest(McpSchema.JSONRPCRequest request) { + private Mono handleIncomingRequest(McpSchema.JSONRPCRequest request, + McpTransportContext transportContext) { return Mono.defer(() -> { Mono resultMono; if (McpSchema.METHOD_INITIALIZE.equals(request.method())) { @@ -266,7 +270,8 @@ public class McpServerSession implements McpLoggableSession { error.message(), error.data()))); } - resultMono = this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, request.params())); + resultMono = this.exchangeSink.asMono() + .flatMap(exchange -> handler.handle(copyExchange(exchange, transportContext), request.params())); } return resultMono .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null)) @@ -280,16 +285,18 @@ public class McpServerSession implements McpLoggableSession { /** * Handles an incoming JSON-RPC notification by routing it to the appropriate handler. * @param notification The incoming JSON-RPC notification + * @param transportContext * @return A Mono that completes when the notification is processed */ - private Mono handleIncomingNotification(McpSchema.JSONRPCNotification notification) { + private Mono handleIncomingNotification(McpSchema.JSONRPCNotification notification, + McpTransportContext transportContext) { return Mono.defer(() -> { if (McpSchema.METHOD_NOTIFICATION_INITIALIZED.equals(notification.method())) { this.state.lazySet(STATE_INITIALIZED); // FIXME: The session ID passed here is not the same as the one in the // legacy SSE transport. exchangeSink.tryEmitValue(new McpAsyncServerExchange(this.id, this, clientCapabilities.get(), - clientInfo.get(), McpTransportContext.EMPTY)); + clientInfo.get(), transportContext)); } var handler = notificationHandlers.get(notification.method()); @@ -297,10 +304,23 @@ public class McpServerSession implements McpLoggableSession { logger.warn("No handler registered for notification method: {}", notification); return Mono.empty(); } - return this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, notification.params())); + return this.exchangeSink.asMono() + .flatMap(exchange -> handler.handle(copyExchange(exchange, transportContext), notification.params())); }); } + /** + * This legacy implementation assumes an exchange is established upon the + * initialization phase see: exchangeSink.tryEmitValue(...), which creates a cached + * immutable exchange. Here, we create a new exchange and copy over everything from + * that cached exchange, and use it for a single HTTP request, with the transport + * context passed in. + */ + private McpAsyncServerExchange copyExchange(McpAsyncServerExchange exchange, McpTransportContext transportContext) { + return new McpAsyncServerExchange(exchange.sessionId(), this, exchange.getClientCapabilities(), + exchange.getClientInfo(), transportContext); + } + record MethodNotFoundError(String method, String message, Object data) { } diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java index e2adb34..acaf0c8 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java @@ -23,6 +23,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import java.util.function.Function; @@ -51,6 +52,7 @@ import io.modelcontextprotocol.spec.McpSchema.Role; import io.modelcontextprotocol.spec.McpSchema.Root; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import io.modelcontextprotocol.spec.McpSchema.Tool; +import io.modelcontextprotocol.util.Utils; import net.javacrumbs.jsonunit.core.Option; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -746,6 +748,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { var clientBuilder = clientBuilders.get(clientType); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) @@ -759,7 +762,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { .GET() .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); - assertThat(responseBody).isNotBlank(); + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); } catch (Exception e) { e.printStackTrace(); @@ -782,6 +785,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); assertThat(response).isNotNull().isEqualTo(callResponse); } @@ -825,6 +829,62 @@ public abstract class AbstractMcpClientServerIntegrationTests { mcpServer.close(); } + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient" }) + void testToolCallSuccessWithTranportContextExtraction(String clientType) { + + var clientBuilder = clientBuilders.get(clientType); + + var transportContextIsNull = new AtomicBoolean(false); + var transportContextIsEmpty = new AtomicBoolean(false); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); + + var expectedCallResponse = new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=value")), null); + McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() + .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) + .callHandler((exchange, request) -> { + + McpTransportContext transportContext = exchange.transportContext(); + transportContextIsNull.set(transportContext == null); + transportContextIsEmpty.set(transportContext.equals(McpTransportContext.EMPTY)); + String ctxValue = (String) transportContext.get("important"); + + try { + String responseBody = "TOOL RESPONSE"; + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); + } + catch (Exception e) { + e.printStackTrace(); + } + + return new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=" + ctxValue)), null); + }) + .build(); + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + + assertThat(transportContextIsNull.get()).isFalse(); + assertThat(transportContextIsEmpty.get()).isFalse(); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); + assertThat(response).isNotNull().isEqualTo(expectedCallResponse); + } + + mcpServer.close(); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient" }) void testToolListChangeHandlingSuccess(String clientType) { @@ -946,7 +1006,7 @@ public abstract class AbstractMcpClientServerIntegrationTests { List receivedNotifications = new CopyOnWriteArrayList<>(); var clientBuilder = clientBuilders.get(clientType); - ; + // Create server with a tool that sends logging notifications McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder() diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java index 823c28d..0f2991a 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java @@ -23,6 +23,7 @@ import io.modelcontextprotocol.server.McpServer.AsyncSpecification; import io.modelcontextprotocol.server.McpServer.SyncSpecification; import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider; import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import jakarta.servlet.http.HttpServletRequest; @Timeout(15) class HttpServletSseIntegrationTests extends AbstractMcpClientServerIntegrationTests { @@ -42,6 +43,7 @@ class HttpServletSseIntegrationTests extends AbstractMcpClientServerIntegrationT // Create and configure the transport provider mcpServerTransportProvider = HttpServletSseServerTransportProvider.builder() .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) .sseEndpoint(CUSTOM_SSE_ENDPOINT) .build(); @@ -92,4 +94,9 @@ class HttpServletSseIntegrationTests extends AbstractMcpClientServerIntegrationT protected void prepareClients(int port, String mcpEndpoint) { } + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + } diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java index 8a8675d..2e9b4cb 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -23,6 +23,7 @@ import io.modelcontextprotocol.server.McpServer.AsyncSpecification; import io.modelcontextprotocol.server.McpServer.SyncSpecification; import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider; import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import jakarta.servlet.http.HttpServletRequest; @Timeout(15) class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests { @@ -40,6 +41,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg // Create and configure the transport provider mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) .keepAliveInterval(Duration.ofSeconds(1)) .build(); @@ -90,4 +92,9 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg protected void prepareClients(int port, String mcpEndpoint) { } + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + }