Skip to content

Commit a48673e

Browse files
committed
Streamable HTTP: flush SSE headers early
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
1 parent 88ebcb7 commit a48673e

2 files changed

Lines changed: 132 additions & 60 deletions

File tree

‎mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java‎

Lines changed: 71 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -425,14 +425,32 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
425425

426426
McpTransportContext transportContext = this.contextExtractor.extract(request);
427427

428-
try {
429-
HttpServletResponseUtils.prepareSseResponse(response);
428+
// Before getWriter(): the character encoding is ignored once the writer is taken
429+
HttpServletResponseUtils.prepareSseResponse(response);
430430

431-
AsyncContext asyncContext = request.startAsync();
431+
PrintWriter writer;
432+
AsyncContext asyncContext;
433+
try {
434+
writer = response.getWriter();
435+
asyncContext = request.startAsync();
432436
asyncContext.setTimeout(0);
437+
}
438+
catch (Exception e) {
439+
// A deployment error, such as a servlet or filter not supporting async.
440+
// Nothing has been sent and async has not started, so an error status can
441+
// still be returned
442+
logger.error("Failed to start the GET stream for session {}: {}", sessionId, e.getMessage());
443+
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
444+
return;
445+
}
433446

434-
HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
435-
sessionId, asyncContext, response.getWriter());
447+
// From here on the headers may already be sent: a failure can only end the stream
448+
HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
449+
sessionId, asyncContext, writer);
450+
try {
451+
// Send the status line and headers now: the listening stream may stay idle
452+
// for a long time, and clients wait for the headers before reading events
453+
response.flushBuffer();
436454

437455
// Replay the messages the client missed while its stream was broken
438456
String lastEventId = request.getHeader(HttpHeaders.LAST_EVENT_ID);
@@ -450,9 +468,14 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
450468

451469
registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport);
452470
}
471+
catch (IOException e) {
472+
// Only flushing the headers throws it: the client is gone
473+
logger.debug("Failed to send the GET stream headers for session {}: {}", sessionId, e.getMessage());
474+
sessionTransport.close();
475+
}
453476
catch (Exception e) {
454477
logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage());
455-
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
478+
sessionTransport.close();
456479
}
457480
}
458481

@@ -597,28 +620,57 @@ else if (message instanceof McpSchema.JSONRPCNotification jsonrpcNotification) {
597620
response.setStatus(HttpServletResponse.SC_ACCEPTED);
598621
}
599622
else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
600-
// For streaming responses, we need to return SSE
623+
// For streaming responses, we need to return SSE. Before getWriter(): the
624+
// character encoding is ignored once the writer is taken
601625
HttpServletResponseUtils.prepareSseResponse(response);
602626

603-
AsyncContext asyncContext = request.startAsync();
604-
asyncContext.setTimeout(0);
627+
PrintWriter writer;
628+
AsyncContext asyncContext;
629+
try {
630+
writer = response.getWriter();
631+
asyncContext = request.startAsync();
632+
asyncContext.setTimeout(0);
633+
}
634+
catch (Exception e) {
635+
// A deployment error, such as a servlet or filter not supporting
636+
// async. Nothing has been sent and async has not started, so an error
637+
// status can still be returned
638+
logger.error("Failed to start the request stream for session {}: {}", sessionId, e.getMessage());
639+
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
640+
return;
641+
}
605642

643+
// From here on the headers may already be sent: a failure can only end
644+
// the stream
606645
HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
607-
sessionId, asyncContext, response.getWriter());
608-
609-
// The listener is given the stream rather than its transport, so that the
610-
// end of the connection detaches the stream from the session instead of
611-
// only dropping the socket: a stream outliving its connection keeps the
612-
// session looking busy and spares it from the sweeper
613-
McpStreamableServerSession.McpStreamableServerSessionStream responseStream = session
614-
.responseStream(sessionTransport);
615-
registerAsyncLifecycle(asyncContext, sessionId, responseStream::releaseTransport);
616-
646+
sessionId, asyncContext, writer);
617647
try {
648+
// Send the status line and headers now: the request may take a long
649+
// time to answer, and clients wait for the headers before reading
650+
// events. Flush before the stream is registered with the session, so
651+
// that no other thread can close it while the headers are written
652+
response.flushBuffer();
653+
654+
// The listener is given the stream rather than its transport, so that
655+
// the end of the connection detaches the stream from the session
656+
// instead of only dropping the socket: a stream outliving its
657+
// connection keeps the session looking busy and spares it from the
658+
// sweeper
659+
McpStreamableServerSession.McpStreamableServerSessionStream responseStream = session
660+
.responseStream(sessionTransport);
661+
registerAsyncLifecycle(asyncContext, sessionId, responseStream::releaseTransport);
662+
618663
responseStream.handle(jsonrpcRequest)
619664
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
620665
.block();
621666
}
667+
catch (IOException e) {
668+
// Only flushing the headers throws it: the client is gone, and the
669+
// request is not handled
670+
logger.debug("Failed to send the request stream headers for session {}: {}", sessionId,
671+
e.getMessage());
672+
sessionTransport.close();
673+
}
622674
catch (Exception e) {
623675
logger.error("Failed to handle request stream: {}", e.getMessage());
624676
// The transport may already have completed the async context, when

‎mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java‎

Lines changed: 61 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.Queue;
2020
import java.util.concurrent.CompletableFuture;
2121
import java.util.concurrent.ConcurrentLinkedQueue;
22+
import java.util.concurrent.TimeUnit;
2223
import java.util.concurrent.atomic.AtomicBoolean;
2324
import java.util.concurrent.atomic.AtomicReference;
2425
import java.util.function.Function;
@@ -39,6 +40,7 @@
3940
import io.modelcontextprotocol.spec.HttpHeaders;
4041
import io.modelcontextprotocol.spec.McpSchema;
4142
import io.modelcontextprotocol.util.KeepAliveScheduler;
43+
import io.modelcontextprotocol.util.ToolsUtils;
4244
import jakarta.servlet.http.HttpServletRequest;
4345
import jakarta.servlet.http.HttpServletResponse;
4446
import org.apache.catalina.LifecycleException;
@@ -55,6 +57,7 @@
5557
import org.slf4j.LoggerFactory;
5658
import reactor.core.publisher.Flux;
5759
import reactor.core.publisher.Mono;
60+
import reactor.core.publisher.Sinks;
5861
import reactor.core.scheduler.Schedulers;
5962
import reactor.test.StepVerifier;
6063
import reactor.test.scheduler.VirtualTimeScheduler;
@@ -205,8 +208,9 @@ void testMissingHandlerReturnsMethodNotFoundError() {
205208
}
206209

207210
@Test
208-
void listeningStreamHasStreamingHeaders() throws Exception {
209-
prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
211+
void listeningStreamSendsHeadersBeforeFirstEvent() throws Exception {
212+
// No keep-alive, so the server never writes to the stream
213+
withTransportProvider(null, null);
210214
var sessionId = initializeSession(httpClient);
211215

212216
var get = HttpRequest.newBuilder()
@@ -216,8 +220,10 @@ void listeningStreamHasStreamingHeaders() throws Exception {
216220
.GET()
217221
.build();
218222

219-
var response = httpClient.send(get, HttpResponse.BodyHandlers.ofInputStream());
223+
// The response completes as soon as the headers are received
224+
var response = httpClient.sendAsync(get, HttpResponse.BodyHandlers.ofInputStream()).get(5, TimeUnit.SECONDS);
220225
try {
226+
assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK);
221227
assertSseHeaders(response);
222228
}
223229
finally {
@@ -226,17 +232,29 @@ void listeningStreamHasStreamingHeaders() throws Exception {
226232
}
227233

228234
@Test
229-
void toolCallSseResponseHasStreamingHeaders() throws Exception {
230-
prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0")
231-
.capabilities(McpSchema.ServerCapabilities.builder().tools(true).build())
235+
void toolCallSseResponseSendsHeadersBeforeResult() throws Exception {
236+
mcpServerTransportProvider.closeGracefully().block();
237+
mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder()
238+
.contextExtractor(TEST_CONTEXT_EXTRACTOR)
239+
.mcpEndpoint(MESSAGE_ENDPOINT)
240+
// No keep-alive, so the server writes nothing until the tool returns
241+
.keepAliveInterval(null)
242+
// no sweeping: the session must survive every keep-alive round
243+
.sessionSweepInterval(null)
244+
.build();
245+
246+
MCP_SERVLET.setDelegate(mcpServerTransportProvider);
247+
248+
// Custom server with a tool that hangs
249+
Sinks.One<McpSchema.CallToolResult> toolResult = Sinks.one();
250+
prepareAsyncServerBuilder().capabilities(McpSchema.ServerCapabilities.builder().tools(true).build())
232251
.tools(McpServerFeatures.AsyncToolSpecification.builder()
233-
.tool(McpSchema.Tool.builder("echo", EMPTY_JSON_SCHEMA).description("returns immediately").build())
234-
.callHandler((exchange,
235-
request) -> Mono.just(McpSchema.CallToolResult.builder()
236-
.content(List.of(McpSchema.TextContent.builder("ok").build()))
237-
.isError(false)
238-
.build()))
252+
.tool(McpSchema.Tool.builder("slow", ToolsUtils.EMPTY_JSON_SCHEMA)
253+
.description("returns when released")
254+
.build())
255+
.callHandler((exchange, request) -> toolResult.asMono())
239256
.build())
257+
.serverInfo("test-server", "1.0.0")
240258
.build();
241259
var sessionId = initializeSession(httpClient);
242260

@@ -247,15 +265,27 @@ void toolCallSseResponseHasStreamingHeaders() throws Exception {
247265
.header(HttpHeaders.MCP_SESSION_ID, sessionId)
248266
.POST(HttpRequest.BodyPublishers
249267
.ofString("{\"jsonrpc\":\"2.0\",\"id\":\"call-1\",\"method\":\"tools/call\","
250-
+ "\"params\":{\"name\":\"echo\",\"arguments\":{}}}"))
268+
+ "\"params\":{\"name\":\"slow\",\"arguments\":{}}}"))
251269
.build();
252270

253-
var response = httpClient.send(post, HttpResponse.BodyHandlers.ofInputStream());
254271
try {
255-
assertSseHeaders(response);
272+
// The response completes as soon as the headers are received
273+
var response = httpClient.sendAsync(post, HttpResponse.BodyHandlers.ofInputStream())
274+
.get(5, TimeUnit.SECONDS);
275+
try {
276+
assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK);
277+
assertSseHeaders(response);
278+
}
279+
finally {
280+
response.body().close();
281+
}
256282
}
257283
finally {
258-
response.body().close();
284+
// Release the Tomcat thread blocked in doPost until the tool returns
285+
toolResult.tryEmitValue(McpSchema.CallToolResult.builder()
286+
.content(List.of(McpSchema.TextContent.builder("done").build()))
287+
.isError(false)
288+
.build());
259289
}
260290
}
261291

@@ -563,10 +593,6 @@ void sessionEvictionAfterReleasingStream() {
563593
var sessionId = initializeSession(httpClient);
564594
var clientStream = openListeningStream(httpClient, sessionId, null);
565595
awaitClientStreamEstablished();
566-
// "prep" the client so we can close the stream by sending data: the
567-
// subscription
568-
// in the clientStream is only present when the client has received data
569-
awaitNotification(clientStream.events());
570596
virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2));
571597
assertThat(postNotification(httpClient, sessionId)).as("Session should be active when a stream is open")
572598
.isEqualTo(HttpServletResponse.SC_ACCEPTED);
@@ -621,17 +647,22 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() {
621647
try {
622648
withTransportProvider(null, SESSION_SWEEP_INTERVAL);
623649
var toolEmissionInterval = Duration.ofMillis(10);
650+
var toolCalled = new AtomicBoolean();
624651
prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0")
625652
.capabilities(McpSchema.ServerCapabilities.builder().tools(true).build())
626653
.tools(McpServerFeatures.AsyncToolSpecification.builder()
627654
.tool(McpSchema.Tool.builder("hangs", EMPTY_JSON_SCHEMA).description("never returns").build())
628-
.callHandler((exchange, request) -> Flux
629-
.interval(toolEmissionInterval, toolEmissionInterval, Schedulers.newSingle("test-tool"))
630-
.flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder()
631-
.level(McpSchema.LoggingLevel.INFO)
632-
.data("x".repeat(64 * 1024))
633-
.build()))
634-
.then(Mono.<McpSchema.CallToolResult>never()))
655+
.callHandler((exchange, request) -> {
656+
toolCalled.set(true);
657+
// Write until the server notices the client is gone
658+
return Flux
659+
.interval(toolEmissionInterval, toolEmissionInterval, Schedulers.newSingle("test-tool"))
660+
.flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder()
661+
.level(McpSchema.LoggingLevel.INFO)
662+
.data("x".repeat(64 * 1024))
663+
.build()))
664+
.then(Mono.<McpSchema.CallToolResult>never());
665+
})
635666
.build())
636667
.build();
637668
var sessionId = initializeSession(httpClient);
@@ -640,18 +671,7 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() {
640671
// stream for as long as the call is in flight
641672
var responseStream = postToolCall(httpClient, sessionId, "hangs");
642673

643-
await().atMost(Duration.ofSeconds(5))
644-
.pollDelay(Duration.ZERO)
645-
.pollInterval(Duration.ofMillis(10))
646-
.untilAsserted(() -> {
647-
// We do this in a loop for the wire-side of the connection - data has
648-
// to be buffered then flushed
649-
650-
// Advance time so the tool emits some data
651-
virtualTime.advanceTimeBy(toolEmissionInterval);
652-
// Verify the data is received
653-
assertThat(responseStream.events()).anyMatch(line -> line.contains("notifications/message"));
654-
});
674+
await().atMost(Duration.ofSeconds(5)).pollInterval(Duration.ofMillis(20)).untilTrue(toolCalled);
655675
virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2));
656676
assertThat(postNotification(httpClient, sessionId)).as("Session with POST stream is active")
657677
.isEqualTo(HttpServletResponse.SC_ACCEPTED);
@@ -781,8 +801,8 @@ record StreamResponse(CompletableFuture<Void> streamFuture, Queue<String> events
781801
AtomicReference<InputStream> streamRef) {
782802

783803
void closeStream() {
784-
// Close listening stream. We retry a few times in case the stream was not
785-
// established on the first try
804+
// Close the stream. The response headers arrive asynchronously, so wait
805+
// until the client has received them
786806
await().pollDelay(Duration.ZERO).atMost(Duration.ofSeconds(1)).until(() -> {
787807
var stream = streamRef.get();
788808
if (stream != null) {

0 commit comments

Comments
 (0)