From 1ba72bd4b559195c350d8b98c596932c29e00b81 Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Tue, 8 Sep 2026 17:56:05 +0200 Subject: [PATCH 1/2] KeepAliveScheduler only pings sessions with an open stream Signed-off-by: Daniel Garnier-Moiroux --- ...vletStreamableServerTransportProvider.java | 20 +++++++- .../spec/McpStreamableServerSession.java | 11 +++++ ...HttpServletStreamableIntegrationTests.java | 47 ++++++++++++++++++- 3 files changed, 75 insertions(+), 3 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 cacb30522..661e96376 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 @@ -20,6 +20,7 @@ import io.modelcontextprotocol.spec.HttpHeaders; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpSession; import io.modelcontextprotocol.spec.McpStreamableServerSession; import io.modelcontextprotocol.spec.McpStreamableServerTransport; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; @@ -165,8 +166,7 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S if (keepAliveInterval != null) { - this.keepAliveScheduler = KeepAliveScheduler - .builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(sessions.values())) + this.keepAliveScheduler = KeepAliveScheduler.builder(this::sessionsToPing) .initialDelay(keepAliveInterval) .interval(keepAliveInterval) .build(); @@ -176,6 +176,22 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S } + /** + * Returns the sessions a keep-alive ping can be sent to, that is the sessions having + * a listening stream. A session without one, e.g. a client which only ever issues + * POST requests, has nothing to write a ping to: pinging it would fail on every + * interval without ever telling us anything about the client being alive. + * @return the sessions to ping + */ + private Flux sessionsToPing() { + if (this.isClosing) { + return Flux.empty(); + } + return Flux.fromIterable(this.sessions.values()) + .filter(McpStreamableServerSession::hasListeningStream) + .cast(McpSession.class); + } + @Override public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) { this.sessionFactory = sessionFactory; 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 7f892df09..f41fb8d2e 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -178,6 +178,17 @@ public Mono delete() { })); } + /** + * Whether the session currently has a listening stream, that is a stream the server + * can send its own requests and notifications to. Sessions have none until the client + * issues the GET request establishing one, and clients are not required to ever issue + * it. + * @return {@code true} if the session has a listening stream + */ + public boolean hasListeningStream() { + return this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream; + } + /** * 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 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 0a918b6df..d538d9a91 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -20,6 +20,9 @@ import java.util.function.Function; import java.util.stream.Stream; +import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.read.ListAppender; import io.modelcontextprotocol.AbstractMcpClientServerIntegrationTests; import io.modelcontextprotocol.client.McpClient; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; @@ -30,6 +33,7 @@ import io.modelcontextprotocol.server.transport.TomcatTestUtil; import io.modelcontextprotocol.spec.HttpHeaders; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.util.KeepAliveScheduler; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; import org.apache.catalina.LifecycleException; @@ -42,6 +46,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.provider.Arguments; +import org.slf4j.LoggerFactory; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -107,7 +112,7 @@ public void before() { mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) - .keepAliveInterval(Duration.ofSeconds(1)) + .keepAliveInterval(Duration.ofMillis(200)) .maxRequestSize(MAX_REQUEST_SIZE) .build(); MCP_SERVLET.setDelegate(mcpServerTransportProvider); @@ -267,6 +272,46 @@ void replacedListeningStreamIsClosed() throws Exception { }); } + @Test + void keepAliveSkipsSessionsWithoutListeningStream() throws Exception { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + var keepAliveLogger = (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(KeepAliveScheduler.class); + ListAppender logAppender = new ListAppender<>(); + logAppender.start(); + keepAliveLogger.addAppender(logAppender); + + try { + // A client is free to never issue the GET request establishing a listening + // stream. Such a session has nothing to write a ping to, so it must not be + // pinged on every keep-alive interval. + initializeSession(httpClient); + + // The keep-alive interval is 200ms, so this spans several intervals + Thread.sleep(1_000); + + assertThat(logAppender.list).noneMatch(event -> event.getLevel() == Level.WARN); + } + finally { + keepAliveLogger.detachAppender(logAppender); + logAppender.stop(); + } + } + + @Test + void keepAlivePingsSessionsWithListeningStream() throws Exception { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + var sessionId = initializeSession(httpClient); + var stream = openListeningStream(httpClient, sessionId, null); + awaitStreamOpen(stream); + + // Sessions with a listening stream are still pinged + await().atMost(Duration.ofSeconds(1)) + .untilAsserted(() -> assertThat(stream.events()).anyMatch(line -> line.contains("\"method\":\"ping\""))); + } + private String initializeSession(HttpClient httpClient) throws Exception { var initialize = HttpRequest.newBuilder() .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) From bf508cacda6bba7a3609a3e87394b0017fff5645 Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Wed, 9 Sep 2026 15:40:50 +0200 Subject: [PATCH 2/2] wip evict sessions --- ...vletStreamableServerTransportProvider.java | 226 +++++++++++++----- .../spec/McpStreamableServerSession.java | 101 ++++++-- .../util/KeepAliveScheduler.java | 39 ++- .../spec/McpStreamableServerSessionTests.java | 125 ++++++++++ ...HttpServletStreamableIntegrationTests.java | 107 ++++++++- 5 files changed, 511 insertions(+), 87 deletions(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java 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 661e96376..fd8a96d60 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 @@ -28,6 +28,8 @@ 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 +37,10 @@ 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 @@ -90,6 +94,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet public static final String FAILED_TO_SEND_ERROR_RESPONSE = "Failed to send error response: {}"; + private static final Duration MAX_IDLE_SWEEP_INTERVAL = Duration.ofMinutes(1); + /** * The endpoint URI where clients should send their JSON-RPC messages. Defaults to * "/mcp". @@ -128,6 +134,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private KeepAliveScheduler keepAliveScheduler; + /** + * How long a session may stay idle before being evicted. Disabled when {@code null}. + */ + private final Duration sessionIdleTimeout; + + /** + * Subscription of the periodic sweep evicting idle sessions. Only set when + * {@link #sessionIdleTimeout} is configured. + */ + private Disposable idleSessionSweep; + /** * Security validator for validating HTTP requests. */ @@ -143,6 +160,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet * @param contextExtractor The extractor for transport context from the request. * @param keepAliveInterval The interval for keep-alive pings. If null, no keep-alive * will be scheduled. + * @param sessionIdleTimeout How long a session may stay idle before being evicted. If + * null, idle sessions are never evicted. * @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. @@ -150,7 +169,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, boolean disallowDelete, McpTransportContextExtractor contextExtractor, - Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize) { + Duration keepAliveInterval, Duration sessionIdleTimeout, ServerHttpHeaderValidator httpHeaderValidator, + int requestMaxSize) { 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"); @@ -163,12 +183,35 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S this.contextExtractor = contextExtractor; this.httpHeaderValidator = httpHeaderValidator; this.requestMaxSize = requestMaxSize; + this.sessionIdleTimeout = sessionIdleTimeout; + + if (sessionIdleTimeout != null) { + Assert.isTrue(!sessionIdleTimeout.isNegative() && !sessionIdleTimeout.isZero(), + "Session idle timeout must be positive"); + // Sweeping more often than the timeout bounds how long an idle session + // outlives it, without making the sweep itself expensive + Duration sweepInterval = sessionIdleTimeout.compareTo(MAX_IDLE_SWEEP_INTERVAL) < 0 ? sessionIdleTimeout + : MAX_IDLE_SWEEP_INTERVAL; + // Evicting closes connections, which is not work for the parallel scheduler + this.idleSessionSweep = Flux.interval(sweepInterval, Schedulers.boundedElastic()) + .doOnNext(tick -> this.evictIdleSessions()) + .onErrorContinue((error, tick) -> logger.warn("Idle session sweep failed: {}", error.getMessage())) + .subscribe(); + } if (keepAliveInterval != null) { 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(); @@ -192,6 +235,40 @@ private Flux sessionsToPing() { .cast(McpSession.class); } + /** + * Removes a session and releases everything it holds, in particular the connections + * of its streams. Together with the client asking for it through a DELETE request, + * this is the only way a session leaves this transport. + * @param sessionId the session to evict + * @param reason why the session is being evicted, for logging + */ + private void evictSession(String sessionId, String reason) { + McpStreamableServerSession session = this.sessions.remove(sessionId); + if (session == null) { + return; + } + logger.info("Evicting session {}: {}", sessionId, reason); + session.close(); + } + + /** + * Evicts the sessions idle for longer than the configured session idle timeout. + * Clients are not required to send a DELETE request when they go away, so a client + * which simply vanishes, whether killed, disconnected or suspended, would otherwise + * keep its session, and the resources behind it, until the server restarts. + */ + private void evictIdleSessions() { + if (this.isClosing) { + return; + } + this.sessions.values() + .stream() + .filter(session -> session.isIdleFor(this.sessionIdleTimeout)) + .map(McpStreamableServerSession::getId) + .toList() + .forEach(sessionId -> this.evictSession(sessionId, "idle for more than " + this.sessionIdleTimeout)); + } + @Override public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) { this.sessionFactory = sessionFactory; @@ -264,6 +341,9 @@ public Mono closeGracefully() { if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.idleSessionSweep != null) { + this.idleSessionSweep.dispose(); + } }); } @@ -354,30 +434,7 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) 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 +442,6 @@ public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException { } } - /** - * 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 @@ -554,6 +583,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( sessionId, asyncContext, response.getWriter()); + registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close); try { session.responseStream(jsonrpcRequest, sessionTransport) @@ -710,6 +740,70 @@ public void destroy() { 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. @@ -783,9 +877,10 @@ public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId 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 +924,6 @@ public void close() { } this.closed = true; - - // HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId); this.asyncContext.complete(); logger.debug("Successfully completed async context for session {}", sessionId); } @@ -871,6 +964,8 @@ public static class Builder { private Duration keepAliveInterval; + private Duration sessionIdleTimeout; + private ServerHttpHeaderValidator httpHeaderValidator = ServerHttpHeaderValidator.NOOP; private int requestMaxSize = DEFAULT_REQUEST_MAX_SIZE; @@ -929,6 +1024,21 @@ public Builder contextExtractor(McpTransportContextExtractor * keep-alive will be scheduled. * @return this builder instance */ + /** + * Sets how long a session may stay idle, that is hold no connection and serve no + * request, before the server evicts it. Clients are expected to end their session + * with a DELETE request, but nothing guarantees they get the chance to, so + * without this the sessions of vanished clients are kept until the server + * restarts. + * @param sessionIdleTimeout the idle timeout. If null, idle sessions are never + * evicted. + * @return this builder instance + */ + public Builder sessionIdleTimeout(Duration sessionIdleTimeout) { + this.sessionIdleTimeout = sessionIdleTimeout; + return this; + } + public Builder keepAliveInterval(Duration keepAliveInterval) { this.keepAliveInterval = keepAliveInterval; return this; @@ -984,7 +1094,7 @@ public HttpServletStreamableServerTransportProvider build() { 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, sessionIdleTimeout, httpHeaderValidator, requestMaxSize); } } 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 f41fb8d2e..1f7a444cc 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -5,7 +5,9 @@ package io.modelcontextprotocol.spec; import java.time.Duration; +import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; @@ -43,6 +45,12 @@ public class McpStreamableServerSession implements McpLoggableSession { private final ConcurrentHashMap requestIdToStream = new ConcurrentHashMap<>(); + /** + * Every stream with a connection currently attached, whether the listening stream or + * a POST response stream, so that they can all be released when the session ends. + */ + private final Set openStreams = ConcurrentHashMap.newKeySet(); + private final String id; private final Duration requestTimeout; @@ -63,6 +71,8 @@ public class McpStreamableServerSession implements McpLoggableSession { private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO; + private volatile long lastActivityNanos = System.nanoTime(); + private final Supplier> onClose; private final JsonSchemaValidator jsonSchemaValidator; @@ -178,6 +188,30 @@ public Mono delete() { })); } + /** + * Records that the client just interacted with this session. Called by the session + * itself whenever it serves the client, so that every transport gets the accounting + * for free. + */ + private void markActive() { + this.lastActivityNanos = System.nanoTime(); + } + + /** + * Whether the session has been idle for longer than the given duration, that is it + * has no connection attached and has not served the client within it. + *

+ * A session holding an open stream is never idle, however long it has been silent: a + * client is allowed to keep a connection open without using it. Noticing that such a + * connection is in fact dead is the job of the keep-alive pings, which release the + * stream and thereby let the session go idle here. + * @param duration the idleness threshold + * @return {@code true} if the session has been idle for longer than {@code duration} + */ + public boolean isIdleFor(Duration duration) { + return this.openStreams.isEmpty() && System.nanoTime() - this.lastActivityNanos > duration.toNanos(); + } + /** * Whether the session currently has a listening stream, that is a stream the server * can send its own requests and notifications to. Sessions have none until the client @@ -189,6 +223,18 @@ public boolean hasListeningStream() { return this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream; } + /** + * Releases the connection of the listening stream, if one is attached, leaving the + * session without one until the client establishes a new stream. Used when the + * connection turns out to be dead, typically because a keep-alive ping went + * unanswered, so that the socket behind it is not held on to for nothing. + */ + public void releaseListeningStream() { + if (this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream stream) { + stream.releaseTransport(); + } + } + /** * 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 @@ -198,11 +244,12 @@ public boolean hasListeningStream() { * @return a stream representation */ public McpStreamableServerSessionStream listeningStream(McpStreamableServerTransport transport) { + markActive(); McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport); McpLoggableSession replaced = this.listeningStreamRef.getAndSet(listeningStream); if (replaced instanceof McpStreamableServerSessionStream replacedStream) { - logger.debug("Closing the listening stream replaced in session {}", this.id); - replacedStream.close(); + logger.debug("Releasing the connection of the listening stream replaced in session {}", this.id); + replacedStream.releaseTransport(); } return listeningStream; } @@ -210,6 +257,7 @@ public McpStreamableServerSessionStream listeningStream(McpStreamableServerTrans // TODO: keep track of history by keeping a map from eventId to stream and then // iterate over the events using the lastEventId public Flux replay(Object lastEventId) { + markActive(); return Flux.empty(); } @@ -220,15 +268,13 @@ public Flux replay(Object lastEventId) { * @return Mono which completes once the processing is done */ public Mono responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, McpStreamableServerTransport transport) { + markActive(); return Mono.deferContextual(ctx -> { McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); McpStreamableServerSessionStream stream = new McpStreamableServerSessionStream(transport); McpRequestHandler requestHandler = McpStreamableServerSession.this.requestHandlers .get(jsonrpcRequest.method()); - // TODO: delegate to stream, which upon successful response should close - // remove itself from the registry and also close the underlying transport - // (sink) if (requestHandler == null) { MethodNotFoundError error = getMethodNotFoundError(jsonrpcRequest.method()); return transport @@ -237,7 +283,7 @@ public Mono responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, McpStr .error(jsonrpcRequest.id(), new McpSchema.JSONRPCResponse.JSONRPCError( McpSchema.ErrorCodes.METHOD_NOT_FOUND, error.message(), error.data()))) - .then(transport.closeGracefully()); + .then(stream.closeGracefully()); } return requestHandler .handle(new McpAsyncServerExchange(this.id, stream, clientCapabilities.get(), clientInfo.get(), @@ -253,7 +299,7 @@ public Mono responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, McpStr return Mono.just(errorResponse); }) .flatMap(transport::sendMessage) - .then(transport.closeGracefully()); + .then(stream.closeGracefully()); }); } @@ -263,6 +309,7 @@ public Mono responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, McpStr * @return Mono which completes upon succesful handling */ public Mono accept(McpSchema.JSONRPCNotification notification) { + markActive(); return Mono.deferContextual(ctx -> { McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); McpNotificationHandler notificationHandler = this.notificationHandlers.get(notification.method()); @@ -284,6 +331,7 @@ public Mono accept(McpSchema.JSONRPCNotification notification) { * @return Mono which completes upon successful processing */ public Mono accept(McpSchema.JSONRPCResponse response) { + markActive(); return Mono.defer(() -> { logger.debug("Received response: {}", response); @@ -324,20 +372,18 @@ private MethodNotFoundError getMethodNotFoundError(String method) { @Override public Mono closeGracefully() { return this.onClose.get().onErrorComplete().then(Mono.defer(() -> { - McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); - return listeningStream.closeGracefully(); - // TODO: Also close all the open streams + this.listeningStreamRef.set(this.missingMcpTransportSession); + return Flux.fromIterable(List.copyOf(this.openStreams)) + .flatMap(McpStreamableServerSessionStream::closeGracefully) + .then(); })); } @Override public void close() { this.onClose.get().onErrorComplete().subscribe(); - McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); - if (listeningStream != null) { - listeningStream.close(); - } - // TODO: Also close all open streams + this.listeningStreamRef.set(this.missingMcpTransportSession); + List.copyOf(this.openStreams).forEach(McpStreamableServerSessionStream::close); } /** @@ -399,6 +445,7 @@ public final class McpStreamableServerSessionStream implements McpLoggableSessio */ public McpStreamableServerSessionStream(McpStreamableServerTransport transport) { this.transport = transport; + McpStreamableServerSession.this.openStreams.add(this); this.transportId = UUID.randomUUID().toString(); // This ID design allows for a constant-time extraction of the history by // precisely identifying the SSE stream using the first component @@ -430,7 +477,9 @@ public Mono sendRequest(String method, Object requestParams, TypeRef t // TODO: store message in history this.transport.sendMessage(jsonrpcRequest, messageId).subscribe(v -> { }, sink::error); - }).timeout(requestTimeout).doOnError(e -> { + }).timeout(requestTimeout).doFinally(signal -> { + // Also on completion and cancellation: a resolved request keeps no state, + // and a deadline imposed by the caller cancels rather than errors this.pendingResponses.remove(requestId); McpStreamableServerSession.this.requestIdToStream.remove(requestId); }).handle((jsonRpcResponse, sink) -> { @@ -459,6 +508,7 @@ public Mono sendNotification(String method, Object params) { @Override public Mono closeGracefully() { return Mono.defer(() -> { + McpStreamableServerSession.this.openStreams.remove(this); this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); this.pendingResponses.clear(); // If this was the generic stream, reset it @@ -471,6 +521,7 @@ public Mono closeGracefully() { @Override public void close() { + McpStreamableServerSession.this.openStreams.remove(this); this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); this.pendingResponses.clear(); // If this was the generic stream, reset it @@ -480,6 +531,24 @@ public void close() { this.transport.close(); } + /** + * Releases the connection carrying this stream, detaching the stream from the + * session, but keeps its pending server-initiated requests resolvable: the client + * answers those with a separate HTTP POST request, which outlives the SSE stream + * the request was sent on. + *

+ * This is the counterpart of {@link #close()} for the end of a connection rather + * than the end of the session: an SSE stream going away, whether replaced, + * disconnected or timed out, does not invalidate the requests sent on it. + */ + public void releaseTransport() { + McpStreamableServerSession.this.openStreams.remove(this); + // If this was the generic stream, reset it + McpStreamableServerSession.this.listeningStreamRef.compareAndExchange(this, + McpStreamableServerSession.this.missingMcpTransportSession); + this.transport.close(); + } + } } 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 6d53ed516..c9967045a 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 @@ 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 Disposable start() { .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 static class Builder { 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 Builder interval(Duration interval) { 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-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java new file mode 100644 index 000000000..c82bc2002 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java @@ -0,0 +1,125 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + +package io.modelcontextprotocol.spec; + +import java.time.Duration; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; + +import io.modelcontextprotocol.json.TypeRef; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link McpStreamableServerSession}. + */ +class McpStreamableServerSessionTests { + + private static final Duration TIMEOUT = Duration.ofSeconds(5); + + private McpStreamableServerSession session() { + return new McpStreamableServerSession("session-1", McpSchema.ClientCapabilities.builder().build(), + new McpSchema.Implementation("test-client", "1.0.0"), TIMEOUT, Map.of(), Map.of()); + } + + @Test + void replacedListeningStreamHasItsConnectionReleased() { + var session = session(); + var firstTransport = new RecordingTransport(); + + session.listeningStream(firstTransport); + assertThat(firstTransport.closed).isFalse(); + + session.listeningStream(new RecordingTransport()); + assertThat(firstTransport.closed).isTrue(); + } + + @Test + void replacingListeningStreamKeepsItsPendingRequestsResolvable() { + var session = session(); + var firstTransport = new RecordingTransport(); + session.listeningStream(firstTransport); + + // Server-initiated requests (sampling, elicitation, roots/list) are sent on the + // listening SSE stream, but the client answers them with a separate HTTP POST + // which outlives that stream. + var pending = session.sendRequest("sampling/createMessage", null, new TypeRef() { + }).toFuture(); + assertThat(firstTransport.sent).hasSize(1); + var requestId = ((McpSchema.JSONRPCRequest) firstTransport.sent.peek()).id(); + + // The client reconnects with a Last-Event-ID header, replacing the listening + // stream. The request sent on the replaced stream must stay pending. + session.listeningStream(new RecordingTransport()); + assertThat(pending).isNotDone(); + + session.accept(McpSchema.JSONRPCResponse.result(requestId, "response-value")).block(TIMEOUT); + + assertThat(pending).succeedsWithin(TIMEOUT).isEqualTo("response-value"); + } + + @Test + void sessionHoldingAnOpenStreamIsNeverIdle() { + var session = session(); + assertThat(session.isIdleFor(Duration.ZERO)).isTrue(); + + var listeningStream = session.listeningStream(new RecordingTransport()); + assertThat(session.isIdleFor(Duration.ZERO)).isFalse(); + + // Releasing the connection is what lets the session go idle, and eventually be + // evicted, once its client turns out to be gone + listeningStream.releaseTransport(); + assertThat(session.isIdleFor(Duration.ZERO)).isTrue(); + } + + @Test + void closingTheSessionReleasesTheConnectionOfItsStreams() { + var session = session(); + var transport = new RecordingTransport(); + session.listeningStream(transport); + + session.close(); + + assertThat(transport.closed).isTrue(); + } + + static class RecordingTransport implements McpStreamableServerTransport { + + final Queue sent = new ConcurrentLinkedQueue<>(); + + volatile boolean closed; + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId) { + return Mono.fromRunnable(() -> this.sent.add(message)); + } + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message) { + return sendMessage(message, null); + } + + @Override + public Mono closeGracefully() { + return Mono.fromRunnable(() -> this.closed = true); + } + + @Override + public void close() { + this.closed = true; + } + + @SuppressWarnings("unchecked") + @Override + public T unmarshalFrom(Object data, TypeRef typeRef) { + return (T) data; + } + + } + +} 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 d538d9a91..5af4a8e9d 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -230,6 +230,7 @@ public void cancel() { @Test void resumedStreamReceivesServerNotifications() throws Exception { + withoutKeepAlive(); prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); var httpClient = HttpClient.newHttpClient(); @@ -240,35 +241,33 @@ void resumedStreamReceivesServerNotifications() throws Exception { // reconnected client never receives anything again. var stream = openListeningStream(httpClient, sessionId, sessionId + "_0"); - awaitStreamOpen(stream); awaitNotification(stream.events()); } @Test void replacedListeningStreamIsClosed() throws Exception { + withoutKeepAlive(); 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); + + // Its body is over, so whatever the first stream received is now final and can + // be compared against once the session starts using the second stream + var receivedByFirstStream = firstStream.events().size(); 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")); + assertThat(firstStream.events()).hasSize(receivedByFirstStream); }); } @@ -312,6 +311,95 @@ void keepAlivePingsSessionsWithListeningStream() throws Exception { .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); + awaitStreamOpen(stream); + + // 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 idleSessionIsEvicted() throws Exception { + useProvider(null, Duration.ofMillis(500)); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + var sessionId = initializeSession(httpClient); + + // Clients are not required to send a DELETE request, so a client which simply + // vanishes must not keep its session, and the resources behind it, forever. + // Probing only once: any request would count as activity and reset the clock. + Thread.sleep(2_000); + + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_NOT_FOUND); + } + + @Test + void sessionHoldingAnOpenStreamIsNotEvicted() throws Exception { + useProvider(null, Duration.ofMillis(500)); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + var httpClient = HttpClient.newHttpClient(); + + var sessionId = initializeSession(httpClient); + var stream = openListeningStream(httpClient, sessionId, null); + awaitNotification(stream.events()); + + // 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(2_000); + + 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(); + } + + /** + * Swaps in a transport provider without keep-alive, for the tests driving raw SSE + * streams which, unlike a real client, never answer the server's pings. + */ + private void withoutKeepAlive() { + useProvider(null, null); + } + + /** + * Swaps in a transport provider with the given lifecycle settings, replacing the one + * the fixture set up. + * @param keepAliveInterval the keep-alive interval, or null to disable keep-alive + * @param sessionIdleTimeout the session idle timeout, or null to disable eviction + */ + private void useProvider(Duration keepAliveInterval, Duration sessionIdleTimeout) { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + .keepAliveInterval(keepAliveInterval) + .sessionIdleTimeout(sessionIdleTimeout) + .maxRequestSize(MAX_REQUEST_SIZE) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + } + private String initializeSession(HttpClient httpClient) throws Exception { var initialize = HttpRequest.newBuilder() .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) @@ -351,8 +439,9 @@ private StreamResponse openListeningStream(HttpClient httpClient, String session } private void awaitNotification(Queue events) { - mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block(); await().atMost(Duration.ofSeconds(5)).pollDelay(Duration.ofMillis(100)).untilAsserted(() -> { + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null) + .block(); assertThat(events).anyMatch(line -> line.contains("notifications/resources/list_changed")); }); }