Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

package io.modelcontextprotocol.server.transport;

import java.io.IOException;

import jakarta.servlet.http.HttpServletResponse;

/**
Expand All @@ -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.
* <p>
* 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();
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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();
}
}

Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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()
Expand All @@ -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 {
Expand All @@ -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<McpSchema.CallToolResult> 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);

Expand All @@ -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());
}
}

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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.<McpSchema.CallToolResult>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.<McpSchema.CallToolResult>never());
})
.build())
.build();
var sessionId = initializeSession(httpClient);
Expand All @@ -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);
Expand Down Expand Up @@ -781,8 +800,7 @@ record StreamResponse(CompletableFuture<Void> streamFuture, Queue<String> events
AtomicReference<InputStream> 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) {
Expand Down
Loading