diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletResponseUtils.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletResponseUtils.java index 40add938c..f52737c7d 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletResponseUtils.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletResponseUtils.java @@ -4,6 +4,8 @@ package io.modelcontextprotocol.server.transport; +import java.io.IOException; + import jakarta.servlet.http.HttpServletResponse; /** @@ -15,22 +17,23 @@ final class HttpServletResponseUtils { private static final String TEXT_EVENT_STREAM = "text/event-stream"; - private static final String UTF_8 = "UTF-8"; - private HttpServletResponseUtils() { } /** - * Prepares the response for streaming Server-Sent Events: sets the content type, - * character encoding and the headers that keep intermediaries from caching, - * transforming or buffering the stream. + * Prepares the response for streaming Server-Sent Events: sets the content type and + * the headers that keep intermediaries from caching, transforming or buffering the + * stream. + *

+ * The headers are flushed immediately so the client receives something even when the + * server delays sending data. * @param response The HTTP servlet response */ - static void prepareSseResponse(HttpServletResponse response) { + static void prepareSseResponse(HttpServletResponse response) throws IOException { response.setContentType(TEXT_EVENT_STREAM); - response.setCharacterEncoding(UTF_8); response.setHeader("Cache-Control", "no-cache, no-transform"); response.setHeader("X-Accel-Buffering", "no"); + response.flushBuffer(); } } 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 index 4bf2f0f47..a80d694be 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java @@ -300,13 +300,13 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) return; } - HttpServletResponseUtils.prepareSseResponse(response); - String sessionId = UUID.randomUUID().toString(); AsyncContext asyncContext = request.startAsync(); asyncContext.setTimeout(0); + response.setCharacterEncoding(UTF_8); PrintWriter writer = response.getWriter(); + HttpServletResponseUtils.prepareSseResponse(response); // Create a new session transport HttpServletMcpSessionTransport sessionTransport = new HttpServletMcpSessionTransport(sessionId, asyncContext, 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 index 7ce830499..ffc9cbdf4 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java @@ -425,14 +425,27 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) McpTransportContext transportContext = this.contextExtractor.extract(request); + PrintWriter writer; + AsyncContext asyncContext; try { - HttpServletResponseUtils.prepareSseResponse(response); - - AsyncContext asyncContext = request.startAsync(); + response.setCharacterEncoding(UTF_8); + writer = response.getWriter(); + asyncContext = request.startAsync(); asyncContext.setTimeout(0); + } + catch (Exception e) { + // Nothing has been written yet, so an error status can still be sent + logger.error("Failed to start the GET stream for session {}: {}", sessionId, e.getMessage()); + response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); + return; + } - HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( - sessionId, asyncContext, response.getWriter()); + HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( + sessionId, asyncContext, writer); + try { + // Send headers as soon as the stream is open: even if the server doesn't send + // data, the client will receive the headers + HttpServletResponseUtils.prepareSseResponse(response); // Replay the messages the client missed while its stream was broken String lastEventId = request.getHeader(HttpHeaders.LAST_EVENT_ID); @@ -450,9 +463,14 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport); } + catch (IOException e) { + // Only flushing the headers throws IOException + logger.debug("Failed to send the GET stream headers for session {}: {}", sessionId, e.getMessage()); + sessionTransport.close(); + } catch (Exception e) { logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage()); - response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); + sessionTransport.close(); } } @@ -598,27 +616,48 @@ else if (message instanceof McpSchema.JSONRPCNotification jsonrpcNotification) { } else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { // For streaming responses, we need to return SSE - HttpServletResponseUtils.prepareSseResponse(response); - - AsyncContext asyncContext = request.startAsync(); - asyncContext.setTimeout(0); + PrintWriter writer; + AsyncContext asyncContext; + try { + response.setCharacterEncoding(UTF_8); + writer = response.getWriter(); + asyncContext = request.startAsync(); + asyncContext.setTimeout(0); + } + catch (Exception e) { + // Nothing has been written yet, so an error status can still be sent + logger.error("Failed to start the request stream for session {}: {}", sessionId, e.getMessage()); + response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); + return; + } HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( - sessionId, asyncContext, response.getWriter()); - - // The listener is given the stream rather than its transport, so that the - // end of the connection detaches the stream from the session instead of - // only dropping the socket: a stream outliving its connection keeps the - // session looking busy and spares it from the sweeper - McpStreamableServerSession.McpStreamableServerSessionStream responseStream = session - .responseStream(sessionTransport); - registerAsyncLifecycle(asyncContext, sessionId, responseStream::releaseTransport); - + sessionId, asyncContext, writer); try { + // Send headers as soon as the stream is open: even if the server + // doesn't send + // data, the client will receive the headers + HttpServletResponseUtils.prepareSseResponse(response); + + // The listener is given the stream rather than its transport, so that + // the end of the connection detaches the stream from the session + // instead of only dropping the socket: a stream outliving its + // connection keeps the session looking busy and spares it from the + // sweeper + McpStreamableServerSession.McpStreamableServerSessionStream responseStream = session + .responseStream(sessionTransport); + registerAsyncLifecycle(asyncContext, sessionId, responseStream::releaseTransport); + responseStream.handle(jsonrpcRequest) .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) .block(); } + catch (IOException e) { + // Only flushing the headers throws IOException + logger.debug("Failed to send the request stream headers for session {}: {}", sessionId, + e.getMessage()); + sessionTransport.close(); + } catch (Exception e) { logger.error("Failed to handle request stream: {}", e.getMessage()); // The transport may already have completed the async context, when diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java index ad2a215fe..135e7ced5 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -19,6 +19,7 @@ import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; @@ -39,6 +40,7 @@ import io.modelcontextprotocol.spec.HttpHeaders; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.util.KeepAliveScheduler; +import io.modelcontextprotocol.util.ToolsUtils; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; import org.apache.catalina.LifecycleException; @@ -55,6 +57,7 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; import reactor.core.scheduler.Schedulers; import reactor.test.StepVerifier; import reactor.test.scheduler.VirtualTimeScheduler; @@ -205,8 +208,9 @@ void testMissingHandlerReturnsMethodNotFoundError() { } @Test - void listeningStreamHasStreamingHeaders() throws Exception { - prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + void listeningStreamSendsHeadersBeforeFirstEvent() throws Exception { + // Without keep-alive nothing is written, so only an early flush sends the headers + withTransportProvider(null, null); var sessionId = initializeSession(httpClient); var get = HttpRequest.newBuilder() @@ -216,8 +220,10 @@ void listeningStreamHasStreamingHeaders() throws Exception { .GET() .build(); - var response = httpClient.send(get, HttpResponse.BodyHandlers.ofInputStream()); + // Completes once the headers arrive, before any event + var response = httpClient.sendAsync(get, HttpResponse.BodyHandlers.ofInputStream()).get(5, TimeUnit.SECONDS); try { + assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK); assertSseHeaders(response); } finally { @@ -226,17 +232,28 @@ void listeningStreamHasStreamingHeaders() throws Exception { } @Test - void toolCallSseResponseHasStreamingHeaders() throws Exception { - prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") - .capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) + void toolCallSseResponseSendsHeadersBeforeResult() throws Exception { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + // Without keep-alive nothing is written until the tool returns + .keepAliveInterval(null) + .sessionSweepInterval(null) + .build(); + + mcpServlet.setDelegate(mcpServerTransportProvider); + + // Returns once the test completes toolResult + Sinks.One toolResult = Sinks.one(); + prepareAsyncServerBuilder().capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) .tools(McpServerFeatures.AsyncToolSpecification.builder() - .tool(McpSchema.Tool.builder("echo", EMPTY_JSON_SCHEMA).description("returns immediately").build()) - .callHandler((exchange, - request) -> Mono.just(McpSchema.CallToolResult.builder() - .content(List.of(McpSchema.TextContent.builder("ok").build())) - .isError(false) - .build())) + .tool(McpSchema.Tool.builder("slow", ToolsUtils.EMPTY_JSON_SCHEMA) + .description("returns when released") + .build()) + .callHandler((exchange, request) -> toolResult.asMono()) .build()) + .serverInfo("test-server", "1.0.0") .build(); var sessionId = initializeSession(httpClient); @@ -247,15 +264,27 @@ void toolCallSseResponseHasStreamingHeaders() throws Exception { .header(HttpHeaders.MCP_SESSION_ID, sessionId) .POST(HttpRequest.BodyPublishers .ofString("{\"jsonrpc\":\"2.0\",\"id\":\"call-1\",\"method\":\"tools/call\"," - + "\"params\":{\"name\":\"echo\",\"arguments\":{}}}")) + + "\"params\":{\"name\":\"slow\",\"arguments\":{}}}")) .build(); - var response = httpClient.send(post, HttpResponse.BodyHandlers.ofInputStream()); try { - assertSseHeaders(response); + // Completes once the headers arrive, before any event + var response = httpClient.sendAsync(post, HttpResponse.BodyHandlers.ofInputStream()) + .get(5, TimeUnit.SECONDS); + try { + assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK); + assertSseHeaders(response); + } + finally { + response.body().close(); + } } finally { - response.body().close(); + // Release the Tomcat thread blocked in doPost until the tool returns + toolResult.tryEmitValue(McpSchema.CallToolResult.builder() + .content(List.of(McpSchema.TextContent.builder("done").build())) + .isError(false) + .build()); } } @@ -563,10 +592,6 @@ void sessionEvictionAfterReleasingStream() { var sessionId = initializeSession(httpClient); var clientStream = openListeningStream(httpClient, sessionId, null); awaitClientStreamEstablished(); - // "prep" the client so we can close the stream by sending data: the - // subscription - // in the clientStream is only present when the client has received data - awaitNotification(clientStream.events()); virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); assertThat(postNotification(httpClient, sessionId)).as("Session should be active when a stream is open") .isEqualTo(HttpServletResponse.SC_ACCEPTED); @@ -621,17 +646,22 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() { try { withTransportProvider(null, SESSION_SWEEP_INTERVAL); var toolEmissionInterval = Duration.ofMillis(10); + var toolCalled = new AtomicBoolean(); prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") .capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) .tools(McpServerFeatures.AsyncToolSpecification.builder() .tool(McpSchema.Tool.builder("hangs", EMPTY_JSON_SCHEMA).description("never returns").build()) - .callHandler((exchange, request) -> Flux - .interval(toolEmissionInterval, toolEmissionInterval, Schedulers.newSingle("test-tool")) - .flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.INFO) - .data("x".repeat(64 * 1024)) - .build())) - .then(Mono.never())) + .callHandler((exchange, request) -> { + toolCalled.set(true); + // Keep writing: a disconnect is only noticed on write + return Flux + .interval(toolEmissionInterval, toolEmissionInterval, Schedulers.newSingle("test-tool")) + .flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.INFO) + .data("x".repeat(64 * 1024)) + .build())) + .then(Mono.never()); + }) .build()) .build(); var sessionId = initializeSession(httpClient); @@ -640,18 +670,7 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() { // stream for as long as the call is in flight var responseStream = postToolCall(httpClient, sessionId, "hangs"); - await().atMost(Duration.ofSeconds(5)) - .pollDelay(Duration.ZERO) - .pollInterval(Duration.ofMillis(10)) - .untilAsserted(() -> { - // We do this in a loop for the wire-side of the connection - data has - // to be buffered then flushed - - // Advance time so the tool emits some data - virtualTime.advanceTimeBy(toolEmissionInterval); - // Verify the data is received - assertThat(responseStream.events()).anyMatch(line -> line.contains("notifications/message")); - }); + await().atMost(Duration.ofSeconds(5)).pollInterval(Duration.ofMillis(20)).untilTrue(toolCalled); virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); assertThat(postNotification(httpClient, sessionId)).as("Session with POST stream is active") .isEqualTo(HttpServletResponse.SC_ACCEPTED); @@ -781,8 +800,7 @@ record StreamResponse(CompletableFuture streamFuture, Queue events AtomicReference streamRef) { void closeStream() { - // Close listening stream. We retry a few times in case the stream was not - // established on the first try + // streamRef is set asynchronously, once the headers arrive await().pollDelay(Duration.ZERO).atMost(Duration.ofSeconds(1)).until(() -> { var stream = streamRef.get(); if (stream != null) {