Adding StreamableHttpServerTransportProvider class and unit tests

Author: ZachGerman

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

@@ -214,7 +214,6 @@ private Mono<McpSchema.InitializeResult> asyncInitializeRequestHandler(
 						"Client requested unsupported protocol version: {}, so the server will suggest the {} version instead",
 						initializeRequest.protocolVersion(), serverProtocolVersion);
 			}
-			System.out.println("---------------Server sending initalize response-----------");
 			return Mono.just(new McpSchema.InitializeResult(serverProtocolVersion, this.serverCapabilities,
 					this.serverInfo, this.instructions));
 		});

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

@@ -325,7 +325,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 			throws ServletException, IOException {
 
 		String requestURI = request.getRequestURI();
-		logger.debug("POST request received for URI: '{}' with headers: {}", requestURI, extractHeaders(request));
+		logger.info("POST request received for URI: '{}' with headers: {}", requestURI, extractHeaders(request));
 
 		if (!requestURI.endsWith(mcpEndpoint)) {
 			logger.debug("URI does not match MCP endpoint: '{}'", mcpEndpoint);
@@ -373,6 +373,7 @@ public void onDataAvailable() throws IOException {
 			public void onAllDataRead() throws IOException {
 				try {
 					// Parse the JSON-RPC message
+					logger.debug("Parsing JSON-RPC message: {}", body.toString());
 					JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body.toString());
 
 					// Check if this is an initialize request
@@ -398,30 +399,59 @@ public void onAllDataRead() throws IOException {
 						return;
 					}
 
+					logger.debug("Creating session for sessionId: {}, isInitialize: {}", sessionId,
+							isInitializeRequest);
 					StreamableHttpSession session = getOrCreateSession(sessionId, isInitializeRequest);
 					if (session == null) {
+						logger.error("Failed to create session for sessionId: {}", sessionId);
 						handleSessionNotFound(sessionId, request, response);
 						asyncContext.complete();
 						return;
 					}
 
+					// Detect response type early
+					StreamableHttpSession.ResponseType responseType = session.detectResponseType(message);
+					logger.debug("Detected response type: {} for message: {}", responseType, message);
+
+					// Set session ID header
+					response.setHeader(SESSION_ID_HEADER, sessionId);
+
 					// Make variables effectively final for lambda
 					final String finalSessionId = sessionId;
 					final StreamableHttpSession finalSession = session;
 					final boolean finalIsInitializeRequest = isInitializeRequest;
 
-					// Handle message and determine response type
-					session.handleMessage(message, request, response, asyncContext).flatMap(responseType -> {
-						response.setHeader(SESSION_ID_HEADER, finalSessionId);
-						if (responseType == StreamableHttpSession.ResponseType.STREAM) {
-							return setupSseResponse(finalSession, finalSessionId, response, request.getRequestId(),
-									asyncContext, finalIsInitializeRequest);
-						}
-						return setupImmediateResponse(finalSession, finalSessionId, response, asyncContext);
-					}).subscribe(null, error -> {
-						logger.error("Error handling message: {}", error.getMessage());
-						asyncContext.complete();
-					});
+					logger.debug("About to handle message with response type: {}", responseType);
+					if (responseType == StreamableHttpSession.ResponseType.STREAM) {
+						logger.debug("Handling STREAM response type");
+						// Set up SSE response first, then handle message
+						sendStreamResponse(finalSession, response, request.getRequestId(), asyncContext)
+							.doOnSuccess(v -> logger.debug("Stream response setup completed"))
+							.doOnError(e -> logger.error("Error in stream response setup: {}", e.getMessage(), e))
+							.then(session.handleMessage(message, request, response, asyncContext))
+							.doOnSuccess(v -> logger.debug("Message handling completed successfully"))
+							.doOnError(e -> logger.error("Error in message handling: {}", e.getMessage(), e))
+							.contextWrite(Context.of(MCP_SESSION_ID, sessionId))
+							.subscribe(null, error -> {
+								logger.error("Error in STREAM handling chain: {}", error.getMessage(), error);
+								asyncContext.complete();
+							});
+					}
+					else {
+						logger.debug("Handling IMMEDIATE response type");
+						// Handle message then set up immediate response
+						session.handleMessage(message, request, response, asyncContext)
+							.doOnSuccess(v -> logger.debug("Message handling completed successfully"))
+							.doOnError(e -> logger.error("Error in message handling: {}", e.getMessage(), e))
+							.then(sendImmediateResponse(finalSession, response, asyncContext))
+							.doOnSuccess(v -> logger.debug("Immediate response setup completed"))
+							.doOnError(e -> logger.error("Error in immediate response setup: {}", e.getMessage(), e))
+							.contextWrite(Context.of(MCP_SESSION_ID, sessionId))
+							.subscribe(null, error -> {
+								logger.error("Error in IMMEDIATE handling chain: {}", error.getMessage(), error);
+								asyncContext.complete();
+							});
+					}
 				}
 				catch (Exception e) {
 					logger.error("Error processing message: {}", e.getMessage());
@@ -515,10 +545,9 @@ private StreamableHttpSession getOrCreateSession(String sessionId, boolean creat
 	/**
 	 * Sets up immediate response for session-based handling.
 	 */
-	private Mono<Void> setupImmediateResponse(StreamableHttpSession session, String sessionId,
-			HttpServletResponse response, AsyncContext asyncContext) {
+	private Mono<Void> sendImmediateResponse(StreamableHttpSession session, HttpServletResponse response,
+			AsyncContext asyncContext) {
 		return Mono.fromRunnable(() -> {
-			response.setHeader(SESSION_ID_HEADER, sessionId);
 			try {
 				response.getWriter().flush();
 			}
@@ -529,25 +558,11 @@ private Mono<Void> setupImmediateResponse(StreamableHttpSession session, String
 	}
 
 	/**
-	 * Sets up SSE response
-	 */
-	private Mono<Void> setupSseResponse(StreamableHttpSession session, String sessionId, HttpServletResponse response,
-			String requestId, AsyncContext asyncContext, boolean isInitializeRequest) {
-		StreamableHttpSseStream sseStream = session.getOrCreateSseStream(requestId);
-		return setupSseResponseForStream(sseStream, response, asyncContext, isInitializeRequest).doFirst(() -> {
-			response.setHeader(SESSION_ID_HEADER, sessionId);
-			if (isInitializeRequest) {
-				response.setHeader(SESSION_ID_HEADER, sessionId);
-			}
-		});
-	}
-
-	/**
-	 * Sets up SSE response for a specific stream.
+	 * Sends SSE response for the given requestId
 	 */
-	private Mono<Void> setupSseResponseForStream(StreamableHttpSseStream sseStream, HttpServletResponse response,
-			AsyncContext asyncContext, boolean isInitializeRequest) {
-		return Mono.fromRunnable(() -> {
+	private Mono<Void> sendStreamResponse(StreamableHttpSession session, HttpServletResponse response, String requestId,
+			AsyncContext asyncContext) {
+		return Mono.create(sink -> {
 			try {
 				// Set up SSE connection
 				response.setContentType(TEXT_EVENT_STREAM);
@@ -558,6 +573,7 @@ private Mono<Void> setupSseResponseForStream(StreamableHttpSseStream sseStream,
 				PrintWriter writer = response.getWriter();
 
 				// Subscribe to the SSE stream and write events to the response
+				StreamableHttpSseStream sseStream = session.getOrCreateSseStream(requestId);
 				sseStream.getEventFlux().doOnNext(event -> {
 					try {
 						if (event.id() != null) {
@@ -588,10 +604,14 @@ private Mono<Void> setupSseResponseForStream(StreamableHttpSseStream sseStream,
 					logger.error("Error in SSE stream: {}", e.getMessage());
 					asyncContext.complete();
 				}).subscribe();
+
+				// Signal that SSE subscription is ready
+				sink.success();
 			}
 			catch (IOException e) {
-				logger.error("Failed to setup SSE response: {}", e.getMessage());
+				logger.error("Failed to send SSE response: {}", e.getMessage());
 				asyncContext.complete();
+				sink.error(e);
 			}
 		});
 	}
@@ -765,38 +785,51 @@ public StreamableHttpSseStream getOrCreateSseStream(String streamName) {
 		/**
 		 * Handles a message using the appropriate transport based on response type.
 		 */
-		public Mono<ResponseType> handleMessage(McpSchema.JSONRPCMessage message, HttpServletRequest request,
+		public Mono<Void> handleMessage(McpSchema.JSONRPCMessage message, HttpServletRequest request,
 				HttpServletResponse response, AsyncContext asyncContext) {
 
-			// Immediate transport only needs 1 session per StreamableHttpSession
-			if (immediateSession == null) {
-				// Create session with both transports on first message
-				immediateTransport = new HttpServerTransport(objectMapper, response, asyncContext);
-				immediateSession = sessionFactory.create(immediateTransport);
-			}
-
-			// Streaming transport requires an inner session for each individual request
-			McpServerTransport streamTransport = getOrCreateSseStream(request.getRequestId()).getTransport();
-			streamSession = sessionFactory.create(streamTransport);
-			streamSessions.add(streamSession);
-
 			ResponseType responseType = detectResponseType(message);
+			logger.debug("Handling message with response type: {}, message: {}", responseType, message);
 
 			if (responseType == ResponseType.STREAM) {
-				// Handle the message with the session uses streamed responses
-				return streamSession.handle(message)
-					.then(Mono.just(responseType))
-					.onErrorReturn(ResponseType.IMMEDIATE);
+				// Streaming transport requires an inner session for each individual
+				// request
+				McpServerTransport streamTransport = getOrCreateSseStream(request.getRequestId()).getTransport();
+				streamSession = sessionFactory.create(streamTransport);
+				// Force the session to skip initialization check by setting the flag via
+				// reflection
+				try {
+					java.lang.reflect.Field skipField = streamSession.getClass()
+						.getDeclaredField("skipInitializationCheck");
+					skipField.setAccessible(true);
+					skipField.set(streamSession, true);
+				}
+				catch (Exception e) {
+					logger.warn("Could not set skipInitializationCheck flag: {}", e.getMessage());
+				}
+				// Copy client info from immediate session if available
+				if (immediateSession != null) {
+					streamSession.init(immediateSession.getClientCapabilities(), immediateSession.getClientInfo());
+				}
+				streamSessions.add(streamSession);
+				return streamSession.handle(message);
+			}
+			else {
+				// Immediate transport only needs 1 session per StreamableHttpSession
+				if (immediateSession == null) {
+					// Create session with immediate transport
+					immediateTransport = new HttpServerTransport(objectMapper, response, asyncContext);
+					immediateSession = sessionFactory.create(immediateTransport);
+				}
+				// Handle the message with the session that uses immediate responses
+				return immediateSession.handle(message);
 			}
-
-			// Handle the message with the session that uses immediate responses
-			return immediateSession.handle(message).then(Mono.just(responseType)).onErrorReturn(ResponseType.IMMEDIATE);
 		}
 
 		/**
 		 * Detects response type based on message characteristics.
 		 */
-		private ResponseType detectResponseType(McpSchema.JSONRPCMessage message) {
+		public ResponseType detectResponseType(McpSchema.JSONRPCMessage message) {
 			if (message instanceof McpSchema.JSONRPCRequest request) {
 				// Initialize requests should return immediate JSON response
 				if (McpSchema.METHOD_INITIALIZE.equals(request.method())) {
@@ -889,6 +922,7 @@ public void sendEvent(String eventType, String data) {
 			String eventId = String.valueOf(++eventCounter);
 			SseEvent event = new SseEvent(eventId, eventType, data);
 			eventHistory.put(eventId, event);
+			logger.debug("Sending SSE event {}: {}", eventId, data);
 			eventSink.tryEmitNext(event);
 		}
 
@@ -963,6 +997,7 @@ public StreamableHttpServerTransport(StreamableHttpSseStream sseStream) {
 		public Mono<Void> sendMessage(JSONRPCMessage message) {
 			try {
 				String jsonText = sseStream.objectMapper.writeValueAsString(message);
+				logger.debug("StreamableHttpServerTransport sending message: {}", jsonText);
 				sseStream.sendEvent(MESSAGE_EVENT_TYPE, jsonText);
 
 				// Complete stream after sending response to avoid hanging

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

@@ -40,7 +40,7 @@ public final class McpSchema {
 	private McpSchema() {
 	}
 
-	public static final String LATEST_PROTOCOL_VERSION = "2025-06-18";
+	public static final String LATEST_PROTOCOL_VERSION = "2024-11-05";
 
 	public static final String JSONRPC_VERSION = "2.0";
 

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

@@ -56,6 +56,8 @@ public class McpServerSession implements McpSession {
 
 	private final AtomicInteger state = new AtomicInteger(STATE_UNINITIALIZED);
 
+	private final boolean skipInitializationCheck;
+
 	/**
 	 * Creates a new server session with the given parameters and the transport to use.
 	 * @param id session id
@@ -72,13 +74,22 @@ public class McpServerSession implements McpSession {
 	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
 			InitRequestHandler initHandler, InitNotificationHandler initNotificationHandler,
 			Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) {
+		this(id, requestTimeout, transport, initHandler, initNotificationHandler, requestHandlers, notificationHandlers,
+				false);
+	}
+
+	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
+			InitRequestHandler initHandler, InitNotificationHandler initNotificationHandler,
+			Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
+			boolean skipInitializationCheck) {
 		this.id = id;
 		this.requestTimeout = requestTimeout;
 		this.transport = transport;
 		this.initRequestHandler = initHandler;
 		this.initNotificationHandler = initNotificationHandler;
 		this.requestHandlers = requestHandlers;
 		this.notificationHandlers = notificationHandlers;
+		this.skipInitializationCheck = skipInitializationCheck;
 	}
 
 	/**
@@ -104,6 +115,14 @@ public void init(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Impl
 		this.clientInfo.lazySet(clientInfo);
 	}
 
+	public McpSchema.ClientCapabilities getClientCapabilities() {
+		return this.clientCapabilities.get();
+	}
+
+	public McpSchema.Implementation getClientInfo() {
+		return this.clientInfo.get();
+	}
+
 	private String generateRequestId() {
 		return this.id + "-" + this.requestCounter.getAndIncrement();
 	}
@@ -222,7 +241,16 @@ private Mono<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCR
 									error.message(), error.data())));
 				}
 
-				resultMono = this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, request.params()));
+				if (skipInitializationCheck) {
+					// Create a dummy exchange for sessions that skip initialization
+					McpAsyncServerExchange dummyExchange = new McpAsyncServerExchange(this, clientCapabilities.get(),
+							clientInfo.get());
+					resultMono = handler.handle(dummyExchange, request.params());
+				}
+				else {
+					resultMono = this.exchangeSink.asMono()
+						.flatMap(exchange -> handler.handle(exchange, request.params()));
+				}
 			}
 			return resultMono
 				.map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null))