modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: Add transport context extraction support to all MCP servers (#477)

Browse source 
- 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 <christian.tzolov@broadcom.com>

Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com>
Co-authored-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Kevin Stanton 2025-08-10 00:18:05 -05:00 • committed by Christian Tzolov
commit b4fef5241f
14 changed files with 374 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

57
mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java
View file

@ -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<String, McpServerSession> sessions = new ConcurrentHashMap<>();
private McpTransportContextExtractor<ServerRequest> 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<ServerRequest> 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.<ServerSentEvent<?>>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<ServerRequest> 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<ServerRequest> 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);
}
}

13
mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java
View file

@ -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"));

8
mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java
View file

@ -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<ServerRequest> 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());

8
mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java
View file

@ -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<ServerRequest> 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

60
mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java
View file

@ -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<String, McpServerSession> sessions = new ConcurrentHashMap<>();
private McpTransportContextExtractor<ServerRequest> 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<ServerRequest> 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<ServerRequest> 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<ServerRequest> 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);
}
}

7
mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java
View file

@ -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<ServerRequest> 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();
}

5
mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java
View file

@ -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

7
mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java
View file

@ -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<ServerRequest> 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();
}

63
mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java
View file

@ -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) {

56
mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java
View file

@ -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<String, McpServerSession> sessions = new ConcurrentHashMap<>();
private McpTransportContextExtractor<HttpServletRequest> 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<HttpServletRequest> 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<HttpServletRequest> 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<HttpServletRequest> contextExtractor) {
Assert.notNull(contextExtractor, "Context extractor must not be null");
this.contextExtractor = contextExtractor;
return this;
}
/**
* Sets the interval for keep-alive pings.
* <p>
@ -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);
}
}

36
mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java
View file

@ -198,7 +198,9 @@ public class McpServerSession implements McpLoggableSession {
* @return a Mono that completes when the message is processed
*/
public Mono<Void> 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<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCRequest request) {
private Mono<McpSchema.JSONRPCResponse> 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<Void> handleIncomingNotification(McpSchema.JSONRPCNotification notification) {
private Mono<Void> 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) {
}

64
mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java
View file

@ -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<McpSchema.LoggingMessageNotification> 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()

7
mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java
View file

@ -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<HttpServletRequest> TEST_CONTEXT_EXTRACTOR = (r, tc) -> {
tc.put("important", "value");
return tc;
};
}

7
mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java
View file

@ -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<HttpServletRequest> TEST_CONTEXT_EXTRACTOR = (r, tc) -> {
tc.put("important", "value");
return tc;
};
}