Fix stream leak when resuming streams with LAST_EVENT_ID

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux 2026-09-08 14:33:25 +02:00 • committed by Daniel Garnier-Moiroux
commit 4186ca13d1
3 changed files with 171 additions and 53 deletions

View file

@ -324,61 +324,44 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
sessionId, asyncContext, response.getWriter());
// Check if this is a replay request
if (request.getHeader(HttpHeaders.LAST_EVENT_ID) != null) {
String lastId = request.getHeader(HttpHeaders.LAST_EVENT_ID);
try {
session.replay(lastId)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.toIterable()
.forEach(message -> {
try {
sessionTransport.sendMessage(message)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.block();
}
catch (Exception e) {
logger.error("Failed to replay message: {}", e.getMessage());
asyncContext.complete();
}
});
}
catch (Exception e) {
logger.error("Failed to replay messages: {}", e.getMessage());
asyncContext.complete();
}
// Replay the messages the client missed while its stream was broken
String lastEventId = request.getHeader(HttpHeaders.LAST_EVENT_ID);
if (lastEventId != null
&& !this.tryReplayMissedMessages(session, lastEventId, sessionTransport, transportContext)) {
// The replay failed and already closed the transport
return;
}
else {
// Establish new listening stream
McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
.listeningStream(sessionTransport);
asyncContext.addListener(new jakarta.servlet.AsyncListener() {
@Override
public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection completed for session: {}", sessionId);
listeningStream.close();
}
// Establish the listening stream. Resumed streams are registered too, so
// that the session keeps delivering messages to the reconnected client and
// the async context is completed once the client goes away.
McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
.listeningStream(sessionTransport);
@Override
public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection timed out for session: {}", sessionId);
listeningStream.close();
}
asyncContext.addListener(new jakarta.servlet.AsyncListener() {
@Override
public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection completed for session: {}", sessionId);
listeningStream.close();
}
@Override
public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection error for session: {}", sessionId);
listeningStream.close();
}
@Override
public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection timed out for session: {}", sessionId);
listeningStream.close();
}
@Override
public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
// No action needed
}
});
}
@Override
public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection error for session: {}", sessionId);
listeningStream.close();
}
@Override
public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
// No action needed
}
});
}
catch (Exception e) {
logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage());
@ -386,6 +369,34 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
}
}
/**
* Replays the messages the client missed while its SSE stream was broken.
* @param session the session the client is resuming
* @param lastEventId the ID of the last event received by the client
* @param sessionTransport the transport of the resumed SSE stream
* @param transportContext the context extracted from the request
* @return {@code true} if the replay completed, {@code false} if it failed, in which
* case the transport has been closed
*/
private boolean tryReplayMissedMessages(McpStreamableServerSession session, String lastEventId,
McpStreamableServerTransport sessionTransport, McpTransportContext transportContext) {
try {
for (McpSchema.JSONRPCMessage message : session.replay(lastEventId)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.toIterable()) {
sessionTransport.sendMessage(message)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.block();
}
return true;
}
catch (Exception e) {
logger.error("Failed to replay messages for session {}: {}", session.getId(), e.getMessage());
sessionTransport.close();
return false;
}
}
/**
* Handles POST requests for incoming JSON-RPC messages from clients.
* @param request The HTTP servlet request containing the JSON-RPC message

View file

@ -179,14 +179,20 @@ public class McpStreamableServerSession implements McpLoggableSession {
}
/**
* Create a listening stream (the generic HTTP GET request without Last-Event-ID
* header).
* Create a listening stream (the generic HTTP GET request, with or without a
* Last-Event-ID header). A session addresses a single listening stream at a time, so
* the stream being replaced, if any, is closed: no message would ever be sent to it
* again, and leaving it open would leak the underlying connection.
* @param transport The dedicated SSE transport stream
* @return a stream representation
*/
public McpStreamableServerSessionStream listeningStream(McpStreamableServerTransport transport) {
McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport);
this.listeningStreamRef.set(listeningStream);
McpLoggableSession replaced = this.listeningStreamRef.getAndSet(listeningStream);
if (replaced instanceof McpStreamableServerSessionStream replacedStream) {
logger.debug("Closing the listening stream replaced in session {}", this.id);
replacedStream.close();
}
return listeningStream;
}

View file

@ -12,6 +12,10 @@ import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Stream;
@ -24,6 +28,7 @@ import io.modelcontextprotocol.server.McpServer.AsyncSpecification;
import io.modelcontextprotocol.server.McpServer.SyncSpecification;
import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider;
import io.modelcontextprotocol.server.transport.TomcatTestUtil;
import io.modelcontextprotocol.spec.HttpHeaders;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
@ -218,4 +223,100 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg
assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_REQUEST_ENTITY_TOO_LARGE);
}
@Test
void resumedStreamReceivesServerNotifications() throws Exception {
prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
var httpClient = HttpClient.newHttpClient();
var sessionId = initializeSession(httpClient);
// Resume the stream the way a client does once its SSE connection broke. The
// resumed stream must become the session listening stream, otherwise the
// reconnected client never receives anything again.
var stream = openListeningStream(httpClient, sessionId, sessionId + "_0");
awaitStreamOpen(stream);
awaitNotification(stream.events());
}
@Test
void replacedListeningStreamIsClosed() throws Exception {
prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build();
var httpClient = HttpClient.newHttpClient();
var sessionId = initializeSession(httpClient);
var firstStream = openListeningStream(httpClient, sessionId, null);
awaitStreamOpen(firstStream);
awaitNotification(firstStream.events());
// stream keeps receiving pings, so we just ensure we've removed the notification
firstStream.events().clear();
assertThat(firstStream.events()).noneMatch(line -> line.contains("notifications/resources/list_changed"));
// Resuming installs a new listening stream. The session can no longer
// address the first one, so it must not be left open.
var secondStream = openListeningStream(httpClient, sessionId, sessionId + "_0");
assertThat(firstStream.streamFuture()).succeedsWithin(Duration.ofSeconds(5));
awaitStreamOpen(secondStream);
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null)
.block();
assertThat(secondStream.events()).anyMatch(line -> line.contains("notifications/resources/list_changed"));
assertThat(firstStream.events()).noneMatch(line -> line.contains("notifications/resources/list_changed"));
});
}
private String initializeSession(HttpClient httpClient) throws Exception {
var initialize = HttpRequest.newBuilder()
.uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT))
.header("Content-Type", "application/json")
.header("Accept", "text/event-stream, application/json")
.POST(HttpRequest.BodyPublishers.ofString("""
{"jsonrpc":"2.0","id":"init","method":"initialize","params":{
"protocolVersion":"2025-06-18","capabilities":{},
"clientInfo":{"name":"test-client","version":"1.0.0"}}}"""))
.build();
var response = httpClient.send(initialize, HttpResponse.BodyHandlers.ofString());
assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK);
return response.headers().firstValue(HttpHeaders.MCP_SESSION_ID).orElseThrow();
}
/**
* Opens an SSE listening stream with a GET request, collecting the received lines.
* @return a future completing once the server closes the stream
*/
private StreamResponse openListeningStream(HttpClient httpClient, String sessionId, String lastEventId) {
var get = HttpRequest.newBuilder()
.uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT))
.header("Accept", "text/event-stream")
.header(HttpHeaders.MCP_SESSION_ID, sessionId);
if (lastEventId != null) {
get.header(HttpHeaders.LAST_EVENT_ID, lastEventId);
}
Queue<String> events = new ConcurrentLinkedQueue<>();
var eventsReceived = new AtomicBoolean(false);
var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofLines())
.thenAccept(response -> {
eventsReceived.set(true);
response.body().forEach(events::add);
});
return new StreamResponse(clientFuture, eventsReceived, events);
}
private void awaitNotification(Queue<String> events) {
mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block();
await().atMost(Duration.ofSeconds(5)).pollDelay(Duration.ofMillis(100)).untilAsserted(() -> {
assertThat(events).anyMatch(line -> line.contains("notifications/resources/list_changed"));
});
}
private static void awaitStreamOpen(StreamResponse stream) {
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> assertThat(stream.isOpen()).isTrue());
}
record StreamResponse(CompletableFuture<Void> streamFuture, AtomicBoolean isOpen, Queue<String> events) {
}
}