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