From 50991f9b6cda24d9f4915ff6f5c36196e0e51802 Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Mon, 14 Sep 2026 16:00:34 +0200 Subject: [PATCH] Remove inactive sessions with a "sweeper" mechanism - Sessions are marked active when a client makes a request - A session with an open stream is also considered active - Inactive sessions are removed at a regular interval Signed-off-by: Daniel Garnier-Moiroux --- ...vletStreamableServerTransportProvider.java | 238 +++++++++++++----- .../spec/McpStreamableServerSession.java | 14 +- .../util/KeepAliveScheduler.java | 39 ++- ...HttpServletStreamableIntegrationTests.java | 229 +++++++++++++++-- 4 files changed, 431 insertions(+), 89 deletions(-) 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 661e963..90e42ed 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 @@ -9,7 +9,9 @@ import java.io.PrintWriter; import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantLock; import io.modelcontextprotocol.common.McpTransportContext; @@ -28,6 +30,8 @@ import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; import jakarta.servlet.AsyncContext; +import jakarta.servlet.AsyncEvent; +import jakarta.servlet.AsyncListener; import jakarta.servlet.ServletException; import jakarta.servlet.annotation.WebServlet; import jakarta.servlet.http.HttpServlet; @@ -35,8 +39,10 @@ import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; /** * Server-side implementation of the Model Context Protocol (MCP) streamable transport @@ -115,6 +121,13 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + /** + * IDs of the sessions which received a request since the last sweep. The set is + * swapped for an empty one on every sweep, so it only ever holds the activity of the + * current interval. + */ + private final AtomicReference> activeSessions = new AtomicReference<>(ConcurrentHashMap.newKeySet()); + private McpTransportContextExtractor contextExtractor; /** @@ -128,6 +141,12 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private KeepAliveScheduler keepAliveScheduler; + /** + * Periodic eviction of the sessions no client came back to. Activated if + * sessionSweepInterval is set. + */ + private Disposable sessionSweeper; + /** * Security validator for validating HTTP requests. */ @@ -146,11 +165,14 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet * @param httpHeaderValidator The HTTP header validator for validating HTTP requests. * @param requestMaxSize The maximum size, in bytes, of a single request body. Must be * positive. + * @param sessionSweepInterval The interval at which idle sessions are evicted. If + * null, no sweeping will be scheduled. * @throws IllegalArgumentException if any parameter is null */ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, boolean disallowDelete, McpTransportContextExtractor contextExtractor, - Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize) { + Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize, + Duration sessionSweepInterval) { Assert.notNull(jsonMapper, "JsonMapper must not be null"); Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); Assert.notNull(contextExtractor, "Context extractor must not be null"); @@ -169,11 +191,72 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet this.keepAliveScheduler = KeepAliveScheduler.builder(this::sessionsToPing) .initialDelay(keepAliveInterval) .interval(keepAliveInterval) + .onPingFailure(session -> { + // The stream the ping was written to is dead. The session survives: + // the client may reconnect to it, and the idle timeout reclaims it if + // it never does. + if (session instanceof McpStreamableServerSession streamableSession) { + streamableSession.releaseListeningStream(); + } + }) .build(); this.keepAliveScheduler.start(); } + if (sessionSweepInterval != null) { + this.sessionSweeper = Flux.interval(sessionSweepInterval, sessionSweepInterval, Schedulers.boundedElastic()) + .doOnNext(tick -> { + // Each sweep runs in its own subscription, so that a sweep failing + // does not terminate the interval and disable sweeping altogether + Mono.fromRunnable(this::sweepSessions) + .doOnError(e -> logger.error("Session sweep failed", e)) + .onErrorComplete() + .subscribe(); + }) + .onErrorComplete(error -> { + logger.error("Session sweeper error", error); + return true; + }) + .subscribe(); + } + + } + + /** + * Evicts the sessions no client is using anymore. A session is kept if it holds an + * open stream, which a client can legitimately sit on without ever writing to it, or + * if it received a request during the interval which just elapsed. Anything else is a + * session whose client went away without deleting it: the protocol lets a client + * reconnect to a session, so nothing else ever reclaims it. + */ + private void sweepSessions() { + if (this.isClosing) { + return; + } + Set active = this.activeSessions.getAndSet(ConcurrentHashMap.newKeySet()); + this.sessions.values().removeIf(session -> { + if (session.hasOpenStream() || active.contains(session.getId())) { + return false; + } + logger.debug("Evicting idle session {}", session.getId()); + try { + session.closeGracefully().block(); + } + catch (Exception e) { + logger.warn("Failed to close idle session {}: {}", session.getId(), e.getMessage()); + } + return true; + }); + } + + /** + * Records that the given session is being used, so that the next sweep does not + * mistake it for a session whose client is gone. + * @param sessionId the session the current request belongs to + */ + private void markSessionActive(String sessionId) { + this.activeSessions.get().add(sessionId); } /** @@ -264,6 +347,9 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.sessionSweeper != null) { + this.sessionSweeper.dispose(); + } }); } @@ -323,6 +409,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet response.sendError(HttpServletResponse.SC_NOT_FOUND); return; } + this.markSessionActive(sessionId); logger.debug("Handling GET request for session: {}", sessionId); @@ -354,30 +441,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet 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(); - } - - @Override - public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException { - logger.debug("SSE connection timed out 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 onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException { - // No action needed - } - }); + registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport); } catch (Exception e) { logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage()); @@ -385,34 +449,6 @@ 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 @@ -479,6 +515,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet }); McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory .startSession(initializeRequest); + this.markSessionActive(init.session().getId()); this.sessions.put(init.session().getId(), init.session()); try { @@ -530,6 +567,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet return; } + this.markSessionActive(sessionId); + if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) { session.accept(jsonrpcResponse) .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) @@ -554,6 +593,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( sessionId, asyncContext, response.getWriter()); + registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close); try { session.responseStream(jsonrpcRequest, sessionTransport) @@ -710,6 +750,70 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet super.destroy(); } + /** + * 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 static 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; + } + } + + /** + * Registers a listener releasing the SSE stream carried by the given asynchronous + * request once the container is done with it, whether the client went away, the + * request timed out or it errored. Without this, the connection is left open, holding + * on to a socket and a container thread, until the process restarts. + * @param asyncContext the asynchronous context of the SSE request + * @param sessionId the session the stream belongs to + * @param onConnectionEnd the action releasing the stream + */ + private static void registerAsyncLifecycle(AsyncContext asyncContext, String sessionId, Runnable onConnectionEnd) { + asyncContext.addListener(new AsyncListener() { + @Override + public void onComplete(AsyncEvent event) throws IOException { + logger.debug("SSE connection completed for session: {}", sessionId); + onConnectionEnd.run(); + } + + @Override + public void onTimeout(AsyncEvent event) throws IOException { + logger.debug("SSE connection timed out for session: {}", sessionId); + onConnectionEnd.run(); + } + + @Override + public void onError(AsyncEvent event) throws IOException { + logger.debug("SSE connection error for session: {}", sessionId); + onConnectionEnd.run(); + } + + @Override + public void onStartAsync(AsyncEvent event) throws IOException { + // No action needed + } + }); + } + /** * Implementation of McpStreamableServerTransport for HttpServlet SSE sessions. This * class handles the transport-level communication for a specific client session. @@ -719,7 +823,6 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet * underlying PrintWriter to prevent race conditions when multiple threads attempt to * send messages concurrently. */ - private class HttpServletStreamableMcpSessionTransport implements McpStreamableServerTransport { private final String sessionId; @@ -783,9 +886,10 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet logger.debug("Message sent to session {} with ID {}", this.sessionId, messageId); } catch (Exception e) { + // The connection is gone, the session is not: the client may come + // back for it, and the idle timeout reclaims it if it never does logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage()); - HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId); - this.asyncContext.complete(); + this.close(); } finally { lock.unlock(); @@ -829,8 +933,6 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet } this.closed = true; - - // HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId); this.asyncContext.complete(); logger.debug("Successfully completed async context for session {}", sessionId); } @@ -875,6 +977,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet private int requestMaxSize = DEFAULT_REQUEST_MAX_SIZE; + private Duration sessionSweepInterval = Duration.ofMinutes(30); + /** * Sets the JsonMapper to use for JSON serialization/deserialization of MCP * messages. @@ -974,6 +1078,22 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet return this; } + /** + * Sets the interval at which idle sessions are evicted. A session is idle once it + * holds no open stream and has received no request for a full interval, which is + * how a session whose client went away without deleting it looks. Nothing else + * reclaims those sessions, as the protocol lets a client reconnect to a session + * it has been disconnected from. + * @param sessionSweepInterval The interval between two sweeps. If null, no + * sweeping will be scheduled and sessions are kept until they are deleted or the + * server shuts down. Defaults to 30 minutes. + * @return this builder instance + */ + public Builder sessionSweepInterval(Duration sessionSweepInterval) { + this.sessionSweepInterval = sessionSweepInterval; + return this; + } + /** * Builds a new instance of {@link HttpServletStreamableServerTransportProvider} * with the configured settings. @@ -984,7 +1104,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); return new HttpServletStreamableServerTransportProvider( jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, mcpEndpoint, disallowDelete, - contextExtractor, keepAliveInterval, httpHeaderValidator, requestMaxSize); + contextExtractor, keepAliveInterval, httpHeaderValidator, requestMaxSize, sessionSweepInterval); } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java index e015b10..63d33ea 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -41,7 +41,7 @@ public class McpStreamableServerSession implements McpLoggableSession { private static final Logger logger = LoggerFactory.getLogger(McpStreamableServerSession.class); - private final ConcurrentHashMap requestIdToStream = new ConcurrentHashMap<>(); + private final ConcurrentHashMap requestIdToStream = new ConcurrentHashMap<>(); /** * Every stream with a connection currently attached, whether the listening stream or @@ -195,6 +195,16 @@ public class McpStreamableServerSession implements McpLoggableSession { return this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream; } + /** + * Whether the session currently has at least one stream with a connection attached, + * whether the listening stream or a POST response stream. A client holding such a + * connection open is still there, however long it stays silent on it. + * @return {@code true} if the session has an open stream + */ + public boolean hasOpenStream() { + return !this.openStreams.isEmpty(); + } + /** * 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 @@ -398,7 +408,7 @@ public class McpStreamableServerSession implements McpLoggableSession { */ public final class McpStreamableServerSessionStream implements McpLoggableSession { - private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); + private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); private final McpStreamableServerTransport connection; diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java b/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java index d07abcd..b625627 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java @@ -6,6 +6,7 @@ package io.modelcontextprotocol.util; import java.time.Duration; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; import java.util.function.Supplier; import org.slf4j.Logger; @@ -57,6 +58,9 @@ public class KeepAliveScheduler { /** Supplier for reactive McpSession instances */ private final Supplier> mcpSessions; + /** Invoked with the session whose keep-alive ping went unanswered */ + private final Consumer onPingFailure; + /** * Creates a KeepAliveScheduler with a custom scheduler, initial delay, interval and a * supplier for McpSession instances. @@ -64,13 +68,15 @@ public class KeepAliveScheduler { * @param initialDelay Initial delay before the first keepAlive call * @param interval Interval between subsequent keepAlive calls * @param mcpSessions Supplier for McpSession instances + * @param onPingFailure Callback invoked with the session whose ping went unanswered */ KeepAliveScheduler(Scheduler scheduler, Duration initialDelay, Duration interval, - Supplier> mcpSessions) { + Supplier> mcpSessions, Consumer onPingFailure) { this.scheduler = scheduler; this.initialDelay = initialDelay; this.interval = interval; this.mcpSessions = mcpSessions; + this.onPingFailure = onPingFailure; } /** @@ -92,8 +98,14 @@ public class KeepAliveScheduler { .doOnNext(tick -> { this.mcpSessions.get() .flatMap(session -> session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF) - .doOnError(e -> logger.warn("Failed to send keep-alive ping to session {}: {}", session, - e.getMessage())) + // A ping has to be answered before the next one is due. The + // request timeout of the session is unrelated to keeping the + // connection alive, and is measured in hours by default. + .timeout(this.interval) + .doOnError(e -> { + logger.warn("Keep-alive ping to session {} failed: {}", session, e.getMessage()); + this.onPingFailure.accept(session); + }) .onErrorComplete()) .subscribe(); }) @@ -154,6 +166,9 @@ public class KeepAliveScheduler { private Supplier> mcpSessions; + private Consumer onPingFailure = session -> { + }; + /** * Creates a new Builder instance with a supplier for McpSession instances. * @param mcpSessions The supplier for McpSession instances @@ -204,12 +219,28 @@ public class KeepAliveScheduler { return this; } + /** + * Sets the callback invoked when a session does not answer a keep-alive ping + * within the keep-alive interval. An unanswered ping means the connection the + * ping was written to is dead, which the operating system does not necessarily + * report: writing to a connection whose peer is gone keeps succeeding until it + * resets. It does not mean the session itself is over, as the client is free to + * reconnect to it. + * @param onPingFailure The callback receiving the unresponsive session + * @return This builder instance for method chaining + */ + public Builder onPingFailure(Consumer onPingFailure) { + Assert.notNull(onPingFailure, "onPingFailure must not be null"); + this.onPingFailure = onPingFailure; + return this; + } + /** * Builds and returns a new KeepAliveScheduler instance. * @return A new KeepAliveScheduler configured with the builder's settings */ public KeepAliveScheduler build() { - return new KeepAliveScheduler(scheduler, initialDelay, interval, mcpSessions); + return new KeepAliveScheduler(scheduler, initialDelay, interval, mcpSessions, onPingFailure); } } 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 d538d9a..4b32027 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -4,18 +4,20 @@ package io.modelcontextprotocol.server; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; import java.net.URI; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; 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; @@ -50,6 +52,7 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; import static org.awaitility.Awaitility.await; @@ -67,6 +70,10 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg private HttpServletStreamableServerTransportProvider mcpServerTransportProvider; + private final Duration KEEP_ALIVE_INTERVAL = Duration.ofMillis(200); + + private final Duration SESSION_SWEEP_INTERVAL = Duration.ofMillis(200); + @Override protected void awaitClientStreamEstablished() { var timeout = Duration.ofSeconds(5); @@ -112,8 +119,9 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) - .keepAliveInterval(Duration.ofMillis(200)) + .keepAliveInterval(KEEP_ALIVE_INTERVAL) .maxRequestSize(MAX_REQUEST_SIZE) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) .build(); MCP_SERVLET.setDelegate(mcpServerTransportProvider); @@ -194,7 +202,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg // A publisher with unknown content length forces chunked transfer encoding, // bypassing the Content-Length header check and exercising the body byte // count - byte[] oversizedBody = "a".repeat(MAX_REQUEST_SIZE + 1).getBytes(StandardCharsets.UTF_8); + byte[] oversizedBody = "a".repeat(MAX_REQUEST_SIZE + 1).getBytes(UTF_8); HttpRequest.BodyPublisher chunkedPublisher = new HttpRequest.BodyPublisher() { @Override public long contentLength() { @@ -240,7 +248,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg // reconnected client never receives anything again. var stream = openListeningStream(httpClient, sessionId, sessionId + "_0"); - awaitStreamOpen(stream); + awaitClientStreamEstablished(); awaitNotification(stream.events()); } @@ -252,21 +260,21 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg var sessionId = initializeSession(httpClient); var firstStream = openListeningStream(httpClient, sessionId, null); - awaitStreamOpen(firstStream); + awaitClientStreamEstablished(); 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. + // Resuming installs a new listening stream. The session can no longer address the + // first one, so it must not be left open. the first stream is only closed when + // the second stream is turned on, so we don't need to wait for a client stream to + // be established var secondStream = openListeningStream(httpClient, sessionId, sessionId + "_0"); assertThat(firstStream.streamFuture()).succeedsWithin(Duration.ofSeconds(5)); - awaitStreamOpen(secondStream); + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block(); 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")); }); @@ -287,8 +295,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg // pinged on every keep-alive interval. initializeSession(httpClient); - // The keep-alive interval is 200ms, so this spans several intervals - Thread.sleep(1_000); + Thread.sleep(KEEP_ALIVE_INTERVAL.multipliedBy(3).toMillis()); assertThat(logAppender.list).noneMatch(event -> event.getLevel() == Level.WARN); } @@ -305,13 +312,168 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg var sessionId = initializeSession(httpClient); var stream = openListeningStream(httpClient, sessionId, null); - awaitStreamOpen(stream); + awaitClientStreamEstablished(); // Sessions with a listening stream are still pinged await().atMost(Duration.ofSeconds(1)) .untilAsserted(() -> assertThat(stream.events()).anyMatch(line -> line.contains("\"method\":\"ping\""))); } + @Test + void unansweredKeepAlivePingReleasesTheStreamButKeepsTheSession() throws Exception { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + var sessionId = initializeSession(httpClient); + var stream = openListeningStream(httpClient, sessionId, null); + awaitClientStreamEstablished(); + + // Nothing answers the pings the server writes to this stream, which is how a + // connection whose client is gone looks: the writes keep succeeding until the + // peer resets. The server must not hold on to it. + assertThat(stream.streamFuture()).succeedsWithin(Duration.ofSeconds(10)); + assertThat(stream.events()).anyMatch(line -> line.contains("\"method\":\"ping\"")); + + // The session itself survives, so a client can come back to it + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + + @Test + void sessionEviction() throws Exception { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + var sessionId = initializeSession(httpClient); + for (int i = 0; i < 4; i++) { + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + Thread.sleep(SESSION_SWEEP_INTERVAL.dividedBy(2).toMillis()); + } + // The session is active, so it's not swept + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + + // Wait for two sweep intervals to ensure the session has been swept + Thread.sleep(SESSION_SWEEP_INTERVAL.multipliedBy(2).toMillis()); + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_NOT_FOUND); + } + + @Test + void sessionHoldingAnOpenStreamIsNotEvicted() throws Exception { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + // remove keepalive to ensure the SSE stream is not removed on ping failures + .keepAliveInterval(null) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + var sessionId = initializeSession(httpClient); + openListeningStream(httpClient, sessionId, null); + awaitClientStreamEstablished(); + + // A client is allowed to hold a connection open without sending anything on it, + // so this session is not idle however long it stays silent + Thread.sleep(SESSION_SWEEP_INTERVAL.multipliedBy(2).toMillis()); + + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + + @Test + void sessionWithStreamIsNotEvicted() throws Exception { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + // remove keepalive to ensure the SSE stream is not removed on ping failures + .keepAliveInterval(null) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + var sessionId = initializeSession(httpClient); + openListeningStream(httpClient, sessionId, null); + + // Wait for two sweep intervals to ensure the sweeper has run + Thread.sleep(SESSION_SWEEP_INTERVAL.multipliedBy(2).toMillis()); + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + + @Test + void sessionEvictionAfterReleasingStream() throws Exception { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + // remove keepalive to ensure the SSE stream is not removed on ping failures + .keepAliveInterval(null) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + 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()); + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + + // Close the client stream + clientStream.closeStream(); + + await().atMost(Duration.ofSeconds(5)) + .pollDelay(Duration.ZERO) + .pollInterval(SESSION_SWEEP_INTERVAL.multipliedBy(2)) + .untilAsserted(() -> { + // send a message so the server realizes the client is gone + // might take a few tries before the internal buffer fills up and the + // connection errors + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null) + .block(); + // eventually, the session should be sweepable + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_NOT_FOUND); + }); + } + + @Test + void sessionIsNotEvictedWithoutSweepInterval() throws Exception { + // Sweeping is opt-in: a transport configured without an interval keeps its + // sessions until they are deleted or the server shuts down + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + .keepAliveInterval(KEEP_ALIVE_INTERVAL) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + + var httpClient = HttpClient.newHttpClient(); + var sessionId = initializeSession(httpClient); + + // The session holds no stream and stays silent for long enough that a scheduled + // sweeper would have evicted it several times over + Thread.sleep(SESSION_SWEEP_INTERVAL.multipliedBy(4).toMillis()); + + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + + private int postNotification(HttpClient httpClient, String sessionId) throws Exception { + var notification = HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) + .header("Content-Type", "application/json") + .header("Accept", "text/event-stream, application/json") + .header(HttpHeaders.MCP_SESSION_ID, sessionId) + .POST(HttpRequest.BodyPublishers.ofString("{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}")) + .build(); + return httpClient.send(notification, HttpResponse.BodyHandlers.ofString()).statusCode(); + } + private String initializeSession(HttpClient httpClient) throws Exception { var initialize = HttpRequest.newBuilder() .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) @@ -330,7 +492,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg /** * Opens an SSE listening stream with a GET request, collecting the received lines. - * @return a future completing once the server closes the stream + * @return a handle on the stream, whose future completes once the server closes it */ private StreamResponse openListeningStream(HttpClient httpClient, String sessionId, String lastEventId) { var get = HttpRequest.newBuilder() @@ -341,13 +503,21 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg get.header(HttpHeaders.LAST_EVENT_ID, lastEventId); } Queue events = new ConcurrentLinkedQueue<>(); - var eventsReceived = new AtomicBoolean(false); - var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofLines()) + var streamRef = new AtomicReference(); + var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofInputStream()) .thenAccept(response -> { - eventsReceived.set(true); - response.body().forEach(events::add); + streamRef.set(response.body()); + try (var r = new BufferedReader(new InputStreamReader(response.body(), UTF_8))) { + String l; + while ((l = r.readLine()) != null) { + events.add(l); + } + } + catch (IOException e) { + // "closed" here is our own stop(), not a failure + } }); - return new StreamResponse(clientFuture, eventsReceived, events); + return new StreamResponse(clientFuture, events, streamRef); } private void awaitNotification(Queue events) { @@ -357,11 +527,22 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg }); } - private static void awaitStreamOpen(StreamResponse stream) { - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> assertThat(stream.isOpen()).isTrue()); - } + record StreamResponse(CompletableFuture streamFuture, Queue events, + AtomicReference streamRef) { - record StreamResponse(CompletableFuture streamFuture, AtomicBoolean isOpen, Queue events) { + void closeStream() { + // Close listening stream. We retry a few times in case the stream was not + // established on the first try + await().atMost(Duration.ofSeconds(1)).until(() -> { + var stream = streamRef.get(); + if (stream != null) { + stream.close(); + return true; + } + return false; + + }); + } } }