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
2 changed files with 374 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

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