add McpTransportContext capability to HttpServletSseServerTransportProvider

Author: stantonk

mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpTransportContext.java

@@ -1,6 +1,7 @@
 package io.modelcontextprotocol.server;
 
 import java.util.Map;
+import java.util.StringJoiner;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -42,4 +43,13 @@ public McpTransportContext copy() {
 		return new DefaultMcpTransportContext(new ConcurrentHashMap<>(this.storage));
 	}
 
+	// TODO for debugging
+
+	@Override
+	public String toString() {
+		return new StringJoiner(", ", DefaultMcpTransportContext.class.getSimpleName() + "[", "]")
+			.add("storage=" + storage)
+			.toString();
+	}
+
 }

mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java

@@ -14,6 +14,9 @@
 
 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;
@@ -99,6 +102,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);
 
@@ -141,7 +146,7 @@ public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String m
 	@Deprecated
 	public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint,
 			String sseEndpoint) {
-		this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null);
+		this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null, null);
 	}
 
 	/**
@@ -160,11 +165,33 @@ public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String b
 	@Deprecated
 	public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint,
 			String sseEndpoint, Duration keepAliveInterval) {
+		this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, null);
+	}
+
+	/**
+	 * 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
+	 * @param contextExtractor The extractor for transport context from the request.
+	 * keep-alive functionality
+	 * @deprecated Use the builder {@link #builder()} instead for better configuration
+	 * options.
+	 */
+	@Deprecated
+	public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint,
+			String sseEndpoint, Duration keepAliveInterval,
+			McpTransportContextExtractor<HttpServletRequest> contextExtractor) {
 
 		this.objectMapper = objectMapper;
 		this.baseUrl = baseUrl;
 		this.messageEndpoint = messageEndpoint;
 		this.sseEndpoint = sseEndpoint;
+		this.contextExtractor = contextExtractor;
 
 		if (keepAliveInterval != null) {
 
@@ -336,10 +363,15 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 				body.append(line);
 			}
 
+			final McpTransportContext transportContext = 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
+			session.handle(message).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); // Block
+																														// for
+																														// Servlet
+																														// compatibility
 
 			response.setStatus(HttpServletResponse.SC_OK);
 		}
@@ -531,6 +563,8 @@ public static class Builder {
 
 		private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
 
+		private McpTransportContextExtractor<HttpServletRequest> contextExtractor = (serverRequest, context) -> context;
+
 		private Duration keepAliveInterval;
 
 		/**
@@ -580,6 +614,19 @@ public Builder sseEndpoint(String sseEndpoint) {
 			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>
@@ -606,7 +653,7 @@ public HttpServletSseServerTransportProvider build() {
 				throw new IllegalStateException("MessageEndpoint must be set");
 			}
 			return new HttpServletSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint,
-					keepAliveInterval);
+					keepAliveInterval, contextExtractor);
 		}
 
 	}

mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java

@@ -194,7 +194,9 @@ public Mono<Void> sendNotification(String method, Object params) {
 	 * @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) {
@@ -210,7 +212,7 @@ public Mono<Void> handle(McpSchema.JSONRPCMessage message) {
 			}
 			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));
@@ -223,7 +225,7 @@ else if (message instanceof McpSchema.JSONRPCNotification notification) {
 				// 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 {
@@ -236,9 +238,11 @@ else if (message instanceof McpSchema.JSONRPCNotification notification) {
 	/**
 	 * 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())) {
@@ -262,7 +266,11 @@ private Mono<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCR
 									error.message(), error.data())));
 				}
 
-				resultMono = this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, request.params()));
+				resultMono = this.exchangeSink.asMono().flatMap(exchange -> {
+					McpAsyncServerExchange newExchange = new McpAsyncServerExchange(exchange.sessionId(), this,
+							exchange.getClientCapabilities(), exchange.getClientInfo(), transportContext);
+					return handler.handle(newExchange, request.params());
+				});
 			}
 			return resultMono
 				.map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null))
@@ -276,24 +284,30 @@ private Mono<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCR
 	/**
 	 * 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());
 			if (handler == null) {
 				logger.error("No handler registered for notification method: {}", notification.method());
 				return Mono.empty();
 			}
-			return this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, notification.params()));
+			return this.exchangeSink.asMono().flatMap(exchange -> {
+				McpAsyncServerExchange newExchange = new McpAsyncServerExchange(exchange.sessionId(), this,
+						exchange.getClientCapabilities(), exchange.getClientInfo(), transportContext);
+				return handler.handle(newExchange, notification.params());
+			});
 		});
 	}