chore: Pass TransportContext sessionFactory

He-Pin · He-Pin · commit aee3a2b11711 · 2025-11-18T20:17:40.000+08:00
Signed-off-by: He-Pin &lt;hepin1989@gmail.com&gt;
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java
@@ -210,6 +210,7 @@ public Mono<Void> notifyClients(String method, Object params) {
 	@Override
 	protected void doGet(HttpServletRequest request, HttpServletResponse response)
 			throws ServletException, IOException {
+		final McpTransportContext transportContext = this.contextExtractor.extract(request);
 
 		String requestURI = request.getRequestURI();
 		if (!requestURI.endsWith(sseEndpoint)) {
@@ -239,7 +240,7 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 				writer);
 
 		// Create a new session using the session factory
-		McpServerSession session = sessionFactory.create(sessionTransport);
+		McpServerSession session = sessionFactory.create(transportContext, sessionTransport);
 		this.sessions.put(sessionId, session);
 
 		// Send initial endpoint event
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java
@@ -408,7 +408,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 						new TypeRef<McpSchema.InitializeRequest>() {
 						});
 				McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory
-					.startSession(initializeRequest);
+					.startSession(transportContext, initializeRequest);
 				this.sessions.put(init.session().getId(), init.session());
 
 				try {
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java
@@ -91,7 +91,7 @@ public List<String> protocolVersions() {
 	public void setSessionFactory(McpServerSession.Factory sessionFactory) {
 		// Create a single session for the stdio connection
 		var transport = new StdioMcpSessionTransport();
-		this.session = sessionFactory.create(transport);
+		this.session = sessionFactory.create(null, transport);
 		transport.initProcessing();
 	}
 
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java
@@ -4,6 +4,7 @@
 
 package io.modelcontextprotocol.spec;
 
+import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.server.McpNotificationHandler;
 import io.modelcontextprotocol.server.McpRequestHandler;
 
@@ -45,11 +46,33 @@ public DefaultMcpStreamableServerSessionFactory(Duration requestTimeout,
 
 	@Override
 	public McpStreamableServerSession.McpStreamableServerSessionInit startSession(
-			McpSchema.InitializeRequest initializeRequest) {
+			final McpSchema.InitializeRequest initializeRequest) {
+		final String sessionId = generateSessionId(null, initializeRequest);
 		return new McpStreamableServerSession.McpStreamableServerSessionInit(
-				new McpStreamableServerSession(UUID.randomUUID().toString(), initializeRequest.capabilities(),
+				new McpStreamableServerSession(sessionId, initializeRequest.capabilities(),
 						initializeRequest.clientInfo(), requestTimeout, requestHandlers, notificationHandlers),
 				this.initRequestHandler.handle(initializeRequest));
 	}
 
+	@Override
+	public McpStreamableServerSession.McpStreamableServerSessionInit startSession(
+			final McpTransportContext mcpTransportContext, final McpSchema.InitializeRequest initializeRequest) {
+		final String sessionId = generateSessionId(mcpTransportContext, initializeRequest);
+		return new McpStreamableServerSession.McpStreamableServerSessionInit(
+				new McpStreamableServerSession(sessionId, initializeRequest.capabilities(),
+						initializeRequest.clientInfo(), requestTimeout, requestHandlers, notificationHandlers),
+				this.initRequestHandler.handle(initializeRequest));
+	}
+
+	/**
+	 * An extensibility point to generate session IDs differently.
+	 * @param mcpTransportContext transport context
+	 * @param initializeRequest initialization request
+	 * @return generated session ID
+	 */
+	protected String generateSessionId(McpTransportContext mcpTransportContext,
+			McpSchema.InitializeRequest initializeRequest) {
+		return UUID.randomUUID().toString();
+	}
+
 }
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java
@@ -439,6 +439,16 @@ public interface Factory {
 		 */
 		McpServerSession create(McpServerTransport sessionTransport);
 
+		/**
+		 * Creates a new 1:1 representation of the client-server interaction.
+		 * @param mcpTransportContext the transport context associated with the client.
+		 * @param sessionTransport the transport to use for communication with the client.
+		 * @return a new server session.
+		 */
+		default McpServerSession create(McpTransportContext mcpTransportContext, McpServerTransport sessionTransport) {
+			return create(sessionTransport);
+		}
+
 	}
 
 }
diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java
@@ -298,8 +298,20 @@ public interface Factory {
 		 * @param initializeRequest the initialization request from the client
 		 * @return a composite allowing the session to start
 		 */
+		@Deprecated
 		McpStreamableServerSessionInit startSession(McpSchema.InitializeRequest initializeRequest);
 
+		/**
+		 * Given an initialize request, create a composite for the session initialization
+		 * @param mcpTransportContext the transport context for the initialization request
+		 * @param initializeRequest the initialization request from the client
+		 * @return a composite allowing the session to start
+		 */
+		default McpStreamableServerSessionInit startSession(McpTransportContext mcpTransportContext,
+				McpSchema.InitializeRequest initializeRequest) {
+			return startSession(initializeRequest);
+		}
+
 	}
 
 	/**
diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java b/mcp-core/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java
@@ -30,7 +30,7 @@ public MockMcpServerTransport getTransport() {
 	@Override
 	public void setSessionFactory(Factory sessionFactory) {
 
-		session = sessionFactory.create(transport);
+		session = sessionFactory.create(null, transport);
 	}
 
 	@Override
diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/StdioServerTransportProviderTests.java
@@ -14,6 +14,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
+import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpServerSession;
@@ -67,7 +68,8 @@ void setUp() {
 		sessionFactory = mock(McpServerSession.Factory.class);
 
 		// Configure mock behavior
-		when(sessionFactory.create(any(McpServerTransport.class))).thenReturn(mockSession);
+		when(sessionFactory.create(any(McpTransportContext.class), any(McpServerTransport.class)))
+			.thenReturn(mockSession);
 		when(mockSession.closeGracefully()).thenReturn(Mono.empty());
 		when(mockSession.sendNotification(any(), any())).thenReturn(Mono.empty());
 
diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java
@@ -278,7 +278,7 @@ private Mono<ServerResponse> handleSseConnection(ServerRequest request) {
 			.body(Flux.<ServerSentEvent<?>>create(sink -> {
 				WebFluxMcpSessionTransport sessionTransport = new WebFluxMcpSessionTransport(sink);
 
-				McpServerSession session = sessionFactory.create(sessionTransport);
+				McpServerSession session = sessionFactory.create(transportContext, sessionTransport);
 				String sessionId = session.getId();
 
 				logger.debug("Created new SSE connection for session: {}", sessionId);
diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java
@@ -239,7 +239,7 @@ private Mono<ServerResponse> handlePost(ServerRequest request) {
 					McpSchema.InitializeRequest initializeRequest = jsonMapper.convertValue(jsonrpcRequest.params(),
 							typeReference);
 					McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory
-						.startSession(initializeRequest);
+						.startSession(transportContext, initializeRequest);
 					sessions.put(init.session().getId(), init.session());
 					return init.initResult().map(initializeResult -> {
 						McpSchema.JSONRPCResponse jsonrpcResponse = new McpSchema.JSONRPCResponse(
diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java
@@ -255,6 +255,7 @@ private ServerResponse handleSseConnection(ServerRequest request) {
 			return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
 		}
 
+		McpTransportContext mcpTransportContext = this.contextExtractor.extract(request);
 		String sessionId = UUID.randomUUID().toString();
 		logger.debug("Creating new SSE connection for session: {}", sessionId);
 
@@ -271,7 +272,7 @@ private ServerResponse handleSseConnection(ServerRequest request) {
 				});
 
 				WebMvcMcpSessionTransport sessionTransport = new WebMvcMcpSessionTransport(sessionId, sseBuilder);
-				McpServerSession session = sessionFactory.create(sessionTransport);
+				McpServerSession session = sessionFactory.create(mcpTransportContext, sessionTransport);
 				this.sessions.put(sessionId, session);
 
 				try {
diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java
@@ -335,7 +335,7 @@ private ServerResponse handlePost(ServerRequest request) {
 						new TypeRef<McpSchema.InitializeRequest>() {
 						});
 				McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory
-					.startSession(initializeRequest);
+					.startSession(transportContext, initializeRequest);
 				this.sessions.put(init.session().getId(), init.session());
 
 				try {

Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ public List<String> protocolVersions() {`
`91`	`91`	`public void setSessionFactory(McpServerSession.Factory sessionFactory) {`
`92`	`92`	`// Create a single session for the stdio connection`
`93`	`93`	`var transport = new StdioMcpSessionTransport();`
`94`		`- this.session = sessionFactory.create(transport);`
	`94`	`+ this.session = sessionFactory.create(null, transport);`
`95`	`95`	`transport.initProcessing();`
`96`	`96`	`}`
`97`	`97`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ public MockMcpServerTransport getTransport() {`
`30`	`30`	`@Override`
`31`	`31`	`public void setSessionFactory(Factory sessionFactory) {`
`32`	`32`
`33`		`- session = sessionFactory.create(transport);`
	`33`	`+ session = sessionFactory.create(null, transport);`
`34`	`34`	`}`
`35`	`35`
`36`	`36`	`@Override`