More WIP: renaming, steteless server abstractions, client context extraction

Signed-off-by: Dariusz Jędrzejczyk <[email protected]>

Author: chemicL

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

@@ -2,11 +2,13 @@
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import io.modelcontextprotocol.spec.DefaultMcpTransportContext;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpServerTransport;
 import io.modelcontextprotocol.spec.McpStreamableServerSession;
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
+import io.modelcontextprotocol.spec.McpTransportContext;
 import io.modelcontextprotocol.util.Assert;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,6 +27,7 @@
 
 import java.io.IOException;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 
 public class WebFluxStreamableServerTransportProvider implements McpStreamableServerTransportProvider {
 
@@ -48,6 +51,9 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe
 
 	private final ConcurrentHashMap<String, McpStreamableServerSession> sessions = new ConcurrentHashMap<>();
 
+	// TODO: add means to specify this
+	private Function<ServerRequest, McpTransportContext> contextExtractor = req -> new DefaultMcpTransportContext();
+
 	/**
 	 * Flag indicating if the transport is shutting down.
 	 */
@@ -183,6 +189,8 @@ private Mono<ServerResponse> handleGet(ServerRequest request) {
 			return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
 		}
 
+		McpTransportContext transportContext = this.contextExtractor.apply(request);
+
 		return Mono.defer(() -> {
 			if (!request.headers().asHttpHeaders().containsKey("mcp-session-id")) {
 				return ServerResponse.badRequest().build(); // TODO: say we need a session id
@@ -204,11 +212,11 @@ private Mono<ServerResponse> handleGet(ServerRequest request) {
 			return ServerResponse.ok().contentType(MediaType.TEXT_EVENT_STREAM)
 					.body(Flux.<ServerSentEvent<?>>create(sink -> {
 						WebFluxStreamableMcpSessionTransport sessionTransport = new WebFluxStreamableMcpSessionTransport(sink);
-						McpStreamableServerSession.McpStreamableServerSessionStream genericStream = session.newStream(sessionTransport);
-						sink.onDispose(genericStream::close);
+						McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session.listeningStream(sessionTransport);
+						sink.onDispose(listeningStream::close);
 					}), ServerSentEvent.class);
 
-		});
+		}).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext));
 	}
 
 	/**
@@ -231,6 +239,8 @@ private Mono<ServerResponse> handlePost(ServerRequest request) {
 			return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
 		}
 
+		McpTransportContext transportContext = this.contextExtractor.apply(request);
+
 		return request.bodyToMono(String.class).<ServerResponse>flatMap(body -> {
 			try {
 				McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body);
@@ -261,7 +271,7 @@ private Mono<ServerResponse> handlePost(ServerRequest request) {
 					return ServerResponse.ok().contentType(MediaType.TEXT_EVENT_STREAM)
 							.body(Flux.<ServerSentEvent<?>>create(sink -> {
 								WebFluxStreamableMcpSessionTransport st = new WebFluxStreamableMcpSessionTransport(sink);
-								Mono<Void> stream = session.handleStream(jsonrpcRequest, st);
+								Mono<Void> stream = session.responseStream(jsonrpcRequest, st);
 								Disposable streamSubscription = stream
 										.doOnError(err -> sink.error(err))
 										.contextWrite(sink.contextView())
@@ -276,7 +286,7 @@ private Mono<ServerResponse> handlePost(ServerRequest request) {
 				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 WebFluxStreamableMcpSessionTransport implements McpServerTransport {

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

@@ -16,7 +16,7 @@
 import java.util.function.BiFunction;
 
 import io.modelcontextprotocol.spec.DefaultMcpStreamableServerSessionFactory;
-import io.modelcontextprotocol.spec.McpSingleSessionServerTransportProvider;
+import io.modelcontextprotocol.spec.McpServerTransportProviderBase;
 import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -89,7 +89,7 @@ public class McpAsyncServer {
 
 	private static final Logger logger = LoggerFactory.getLogger(McpAsyncServer.class);
 
-	private final McpServerTransportProvider mcpTransportProvider;
+	private final McpServerTransportProviderBase mcpTransportProvider;
 
 	private final ObjectMapper objectMapper;
 
@@ -126,7 +126,7 @@ public class McpAsyncServer {
 	 * @param features The MCP server supported features.
 	 * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
 	 */
-	McpAsyncServer(McpSingleSessionServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
+	McpAsyncServer(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
 				   McpServerFeatures.Async features, Duration requestTimeout,
 				   McpUriTemplateManagerFactory uriTemplateManagerFactory, JsonSchemaValidator jsonSchemaValidator) {
 		this.mcpTransportProvider = mcpTransportProvider;

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

@@ -8,11 +8,13 @@
 import java.util.Collections;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+import io.modelcontextprotocol.spec.DefaultMcpTransportContext;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
 import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
 import io.modelcontextprotocol.spec.McpSession;
+import io.modelcontextprotocol.spec.McpTransportContext;
 import io.modelcontextprotocol.util.Assert;
 import reactor.core.publisher.Mono;
 
@@ -31,6 +33,8 @@ public class McpAsyncServerExchange {
 
 	private final McpSchema.Implementation clientInfo;
 
+	private final McpTransportContext transportContext;
+
 	private volatile LoggingLevel minLoggingLevel = LoggingLevel.INFO;
 
 	private static final TypeReference<McpSchema.CreateMessageResult> CREATE_MESSAGE_RESULT_TYPE_REF = new TypeReference<>() {
@@ -57,6 +61,23 @@ public McpAsyncServerExchange(McpSession session, McpSchema.ClientCapabilities c
 		this.session = session;
 		this.clientCapabilities = clientCapabilities;
 		this.clientInfo = clientInfo;
+		this.transportContext = new DefaultMcpTransportContext();
+	}
+
+	/**
+	 * Create a new asynchronous exchange with the client.
+	 * @param session The server session representing a 1-1 interaction.
+	 * @param clientCapabilities The client capabilities that define the supported
+	 * features and functionality.
+	 * @param transportContext context associated with the client as extracted from the transport
+	 * @param clientInfo The client implementation information.
+	 */
+	public McpAsyncServerExchange(McpSession session, McpSchema.ClientCapabilities clientCapabilities,
+								  McpSchema.Implementation clientInfo, McpTransportContext transportContext) {
+		this.session = session;
+		this.clientCapabilities = clientCapabilities;
+		this.clientInfo = clientInfo;
+		this.transportContext = transportContext;
 	}
 
 	/**
@@ -75,6 +96,10 @@ public McpSchema.Implementation getClientInfo() {
 		return this.clientInfo;
 	}
 
+	public McpTransportContext transportContext() {
+		return this.transportContext;
+	}
+
 	/**
 	 * Create a new message using the sampling capabilities of the client. The Model
 	 * Context Protocol (MCP) provides a standardized way for servers to request LLM