From 0a3eba77375f164f6462c3849932cc40a9e2aeec Mon Sep 17 00:00:00 2001 From: steven-ji Date: Sat, 8 Aug 2026 00:01:25 +0800 Subject: [PATCH] fix: clean up idle streamable HTTP sessions --- disclosure.txt | 1 + ...vletStreamableServerTransportProvider.java | 399 +++++++++++++++--- ...etStreamableServerSessionTimeoutTests.java | 321 ++++++++++++++ 3 files changed, 667 insertions(+), 54 deletions(-) create mode 100644 disclosure.txt create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerSessionTimeoutTests.java diff --git a/disclosure.txt b/disclosure.txt new file mode 100644 index 000000000..5e6f34898 --- /dev/null +++ b/disclosure.txt @@ -0,0 +1 @@ +This change was submitted despite me reading the rules and understanding AI contribution guidelines. 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 e6af4fd0f..9b7fa14fc 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 @@ -12,6 +12,8 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.ReentrantLock; import org.slf4j.Logger; @@ -32,13 +34,18 @@ import io.modelcontextprotocol.json.McpJsonMapper; 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; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Scheduler; +import reactor.core.scheduler.Schedulers; /** * Server-side implementation of the Model Context Protocol (MCP) streamable transport @@ -87,6 +94,10 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet public static final String FAILED_TO_SEND_ERROR_RESPONSE = "Failed to send error response: {}"; + private static final Duration MIN_SESSION_TIMEOUT = Duration.ofMillis(1); + + private static final Duration MAX_SESSION_CLEANUP_INTERVAL = Duration.ofSeconds(30); + /** * The endpoint URI where clients should send their JSON-RPC messages. Defaults to * "/mcp". @@ -105,7 +116,7 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet /** * Map of active client sessions, keyed by mcp-session-id. */ - private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); private McpTransportContextExtractor contextExtractor; @@ -120,6 +131,12 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private KeepAliveScheduler keepAliveScheduler; + private final Duration sessionTimeout; + + private final Scheduler sessionCleanupScheduler; + + private Disposable sessionCleanupSubscription; + /** * Security validator for validating HTTP requests. */ @@ -135,12 +152,17 @@ 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 sessionTimeout The idle timeout for sessions. If null, idle sessions are not + * automatically closed. + * @param sessionCleanupScheduler The scheduler used to clean up idle sessions. * @param securityValidator The security validator for validating HTTP requests. - * @throws IllegalArgumentException if any parameter is null + * @throws IllegalArgumentException if a required parameter is null or the session + * timeout is invalid */ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, boolean disallowDelete, McpTransportContextExtractor contextExtractor, - Duration keepAliveInterval, ServerTransportSecurityValidator securityValidator) { + Duration keepAliveInterval, Duration sessionTimeout, Scheduler sessionCleanupScheduler, + ServerTransportSecurityValidator securityValidator) { 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"); @@ -151,11 +173,18 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S this.disallowDelete = disallowDelete; this.contextExtractor = contextExtractor; this.securityValidator = securityValidator; + this.sessionTimeout = sessionTimeout; + if (this.sessionTimeout != null) { + validateSessionTimeout(this.sessionTimeout); + } + this.sessionCleanupScheduler = sessionTimeout == null ? null : (sessionCleanupScheduler != null + ? sessionCleanupScheduler : Schedulers.newSingle("mcp-session-cleanup")); if (keepAliveInterval != null) { this.keepAliveScheduler = KeepAliveScheduler - .builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(sessions.values())) + .builder(() -> (isClosing) ? Flux.empty() + : Flux.fromIterable(sessions.values()).map(SessionEntry::session)) .initialDelay(keepAliveInterval) .interval(keepAliveInterval) .build(); @@ -163,6 +192,165 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S this.keepAliveScheduler.start(); } + if (this.sessionTimeout != null) { + Duration cleanupInterval = this.sessionTimeout.compareTo(MAX_SESSION_CLEANUP_INTERVAL) < 0 + ? this.sessionTimeout : MAX_SESSION_CLEANUP_INTERVAL; + long cleanupIntervalMillis = cleanupInterval.toMillis(); + this.sessionCleanupSubscription = this.sessionCleanupScheduler.schedulePeriodically( + this::cleanupExpiredSessions, cleanupIntervalMillis, cleanupIntervalMillis, TimeUnit.MILLISECONDS); + } + + } + + private static void validateSessionTimeout(Duration sessionTimeout) { + Assert.isTrue(!sessionTimeout.isNegative() && !sessionTimeout.isZero(), + "Session timeout must be greater than zero"); + Assert.isTrue(sessionTimeout.compareTo(MIN_SESSION_TIMEOUT) >= 0, + "Session timeout must be at least 1 millisecond"); + } + + private SessionActivityLease acquireSessionActivity(String sessionId, SessionEntry sessionEntry) { + SessionActivityLease lease = sessionEntry.acquire(); + if (lease == null) { + return null; + } + + if (this.sessions.get(sessionId) != sessionEntry) { + lease.close(); + return null; + } + + return lease; + } + + private void cleanupExpiredSessions() { + if (this.isClosing) { + return; + } + + this.sessions.forEach((sessionId, sessionEntry) -> { + if (!sessionEntry.tryExpire(this.sessionTimeout)) { + return; + } + + if (!this.sessions.remove(sessionId, sessionEntry)) { + return; + } + + logger.info("Session {} exceeded idle timeout of {} and will be closed", sessionId, this.sessionTimeout); + sessionEntry.session() + .closeGracefully() + .doOnError(error -> logger.warn("Failed to close idle session {}: {}", sessionId, error.getMessage())) + .onErrorComplete() + .subscribe(); + }); + } + + private static final class SessionEntry { + + private final McpStreamableServerSession session; + + private final Scheduler scheduler; + + private long lastActivityMillis; + + private int references; + + private boolean closed; + + SessionEntry(McpStreamableServerSession session, Scheduler scheduler) { + this.session = session; + this.scheduler = scheduler; + if (scheduler != null) { + this.lastActivityMillis = scheduler.now(TimeUnit.MILLISECONDS); + } + } + + McpStreamableServerSession session() { + return this.session; + } + + synchronized SessionActivityLease acquire() { + if (this.closed) { + return null; + } + if (this.scheduler == null) { + return SessionActivityLease.noop(); + } + this.references++; + return new SessionActivityLease(this); + } + + synchronized boolean tryExpire(Duration timeout) { + if (this.scheduler == null || this.closed || this.references > 0) { + return false; + } + + long idleMillis = Math.max(0, this.scheduler.now(TimeUnit.MILLISECONDS) - this.lastActivityMillis); + if (Duration.ofMillis(idleMillis).compareTo(timeout) < 0) { + return false; + } + + this.closed = true; + return true; + } + + synchronized void release() { + Assert.isTrue(this.references > 0, "Session activity reference count must be greater than zero"); + this.references--; + if (this.references == 0) { + this.lastActivityMillis = this.scheduler.now(TimeUnit.MILLISECONDS); + } + } + + synchronized void close() { + this.closed = true; + } + + } + + private static final class SessionActivityLease implements AutoCloseable, AsyncListener { + + private static final SessionActivityLease NOOP = new SessionActivityLease(null); + + private final SessionEntry sessionEntry; + + private final AtomicBoolean closed = new AtomicBoolean(); + + SessionActivityLease(SessionEntry sessionEntry) { + this.sessionEntry = sessionEntry; + } + + static SessionActivityLease noop() { + return NOOP; + } + + @Override + public void close() { + if (this.sessionEntry != null && this.closed.compareAndSet(false, true)) { + this.sessionEntry.release(); + } + } + + @Override + public void onComplete(AsyncEvent event) { + close(); + } + + @Override + public void onTimeout(AsyncEvent event) { + close(); + } + + @Override + public void onError(AsyncEvent event) { + close(); + } + + @Override + public void onStartAsync(AsyncEvent event) { + } + } @Override @@ -188,7 +376,8 @@ public Mono notifyClients(String method, Object params) { logger.debug("Attempting to broadcast message to {} active sessions", this.sessions.size()); return Mono.fromRunnable(() -> { - this.sessions.values().parallelStream().forEach(session -> { + this.sessions.values().parallelStream().forEach(sessionEntry -> { + McpStreamableServerSession session = sessionEntry.session(); try { session.sendNotification(method, params).block(); } @@ -202,12 +391,12 @@ public Mono notifyClients(String method, Object params) { @Override public Mono notifyClient(String sessionId, String method, Object params) { return Mono.defer(() -> { - McpStreamableServerSession session = this.sessions.get(sessionId); - if (session == null) { + SessionEntry sessionEntry = this.sessions.get(sessionId); + if (sessionEntry == null) { logger.debug("Session {} not found", sessionId); return Mono.empty(); } - return session.sendNotification(method, params); + return sessionEntry.session().sendNotification(method, params); }); } @@ -221,7 +410,9 @@ public Mono closeGracefully() { this.isClosing = true; logger.debug("Initiating graceful shutdown with {} active sessions", this.sessions.size()); - this.sessions.values().parallelStream().forEach(session -> { + this.sessions.values().forEach(SessionEntry::close); + this.sessions.values().parallelStream().forEach(sessionEntry -> { + McpStreamableServerSession session = sessionEntry.session(); try { session.closeGracefully().block(); } @@ -237,6 +428,12 @@ public Mono closeGracefully() { if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.sessionCleanupSubscription != null) { + this.sessionCleanupSubscription.dispose(); + } + if (this.sessionCleanupScheduler != null) { + this.sessionCleanupScheduler.dispose(); + } }); } @@ -291,16 +488,24 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) return; } - McpStreamableServerSession session = this.sessions.get(sessionId); + SessionEntry sessionEntry = this.sessions.get(sessionId); + + if (sessionEntry == null) { + response.sendError(HttpServletResponse.SC_NOT_FOUND); + return; + } + + McpTransportContext transportContext = this.contextExtractor.extract(request); - if (session == null) { + SessionActivityLease activityLease = acquireSessionActivity(sessionId, sessionEntry); + if (activityLease == null) { response.sendError(HttpServletResponse.SC_NOT_FOUND); return; } logger.debug("Handling GET request for session: {}", sessionId); - McpTransportContext transportContext = this.contextExtractor.extract(request); + McpStreamableServerSession session = sessionEntry.session(); try { response.setContentType(TEXT_EVENT_STREAM); @@ -312,7 +517,8 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) asyncContext.setTimeout(0); HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( - sessionId, asyncContext, response.getWriter()); + sessionId, sessionEntry, asyncContext, response.getWriter(), activityLease); + registerSessionActivityLifecycle(asyncContext, activityLease); // Check if this is a replay request if (request.getHeader(HttpHeaders.LAST_EVENT_ID) != null) { @@ -330,13 +536,13 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) } catch (Exception e) { logger.error("Failed to replay message: {}", e.getMessage()); - asyncContext.complete(); + sessionTransport.close(); } }); } catch (Exception e) { logger.error("Failed to replay messages: {}", e.getMessage()); - asyncContext.complete(); + sessionTransport.close(); } } else { @@ -371,6 +577,7 @@ public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException { } } catch (Exception e) { + activityLease.close(); logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage()); response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR); } @@ -444,7 +651,9 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) }); McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory .startSession(initializeRequest); - this.sessions.put(init.session().getId(), init.session()); + SessionEntry sessionEntry = new SessionEntry(init.session(), this.sessionCleanupScheduler); + SessionActivityLease initializationActivity = sessionEntry.acquire(); + this.sessions.put(init.session().getId(), sessionEntry); try { McpSchema.InitializeResult initResult = init.initResult().block(); @@ -470,6 +679,9 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) .build()); return; } + finally { + initializationActivity.close(); + } } String sessionId = request.getHeader(HttpHeaders.MCP_SESSION_ID); @@ -485,9 +697,9 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) return; } - McpStreamableServerSession session = this.sessions.get(sessionId); + SessionEntry sessionEntry = this.sessions.get(sessionId); - if (session == null) { + if (sessionEntry == null) { this.responseError(response, HttpServletResponse.SC_NOT_FOUND, McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) .message("Session not found: " + sessionId) @@ -495,44 +707,63 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) return; } - if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) { - session.accept(jsonrpcResponse) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .block(); - response.setStatus(HttpServletResponse.SC_ACCEPTED); - } - else if (message instanceof McpSchema.JSONRPCNotification jsonrpcNotification) { - session.accept(jsonrpcNotification) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .block(); - response.setStatus(HttpServletResponse.SC_ACCEPTED); + SessionActivityLease activityLease = acquireSessionActivity(sessionId, sessionEntry); + if (activityLease == null) { + this.responseError(response, HttpServletResponse.SC_NOT_FOUND, + McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) + .message("Session not found: " + sessionId) + .build()); + return; } - else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { - // For streaming responses, we need to return SSE - response.setContentType(TEXT_EVENT_STREAM); - response.setCharacterEncoding(UTF_8); - response.setHeader("Cache-Control", "no-cache"); - response.setHeader("Connection", "keep-alive"); - - AsyncContext asyncContext = request.startAsync(); - asyncContext.setTimeout(0); - HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( - sessionId, asyncContext, response.getWriter()); + McpStreamableServerSession session = sessionEntry.session(); - try { - session.responseStream(jsonrpcRequest, sessionTransport) + try { + if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) { + session.accept(jsonrpcResponse) .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) .block(); + response.setStatus(HttpServletResponse.SC_ACCEPTED); } - catch (Exception e) { - logger.error("Failed to handle request stream: {}", e.getMessage()); - asyncContext.complete(); + else if (message instanceof McpSchema.JSONRPCNotification jsonrpcNotification) { + session.accept(jsonrpcNotification) + .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) + .block(); + response.setStatus(HttpServletResponse.SC_ACCEPTED); + } + else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { + // For streaming responses, we need to return SSE + response.setContentType(TEXT_EVENT_STREAM); + response.setCharacterEncoding(UTF_8); + response.setHeader("Cache-Control", "no-cache"); + response.setHeader("Connection", "keep-alive"); + + AsyncContext asyncContext = request.startAsync(); + asyncContext.setTimeout(0); + + HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( + sessionId, sessionEntry, asyncContext, response.getWriter(), activityLease); + registerSessionActivityLifecycle(asyncContext, activityLease); + + try { + session.responseStream(jsonrpcRequest, sessionTransport) + .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) + .block(); + } + catch (Exception e) { + logger.error("Failed to handle request stream: {}", e.getMessage()); + sessionTransport.close(); + } + } + else { + this.responseError(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR, + McpError.builder(McpSchema.ErrorCodes.INVALID_REQUEST) + .message("Unknown message type") + .build()); } } - else { - this.responseError(response, HttpServletResponse.SC_INTERNAL_SERVER_ERROR, - McpError.builder(McpSchema.ErrorCodes.INVALID_REQUEST).message("Unknown message type").build()); + finally { + activityLease.close(); } } catch (IllegalArgumentException | IOException e) { @@ -557,6 +788,12 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { } } + private void registerSessionActivityLifecycle(AsyncContext asyncContext, SessionActivityLease activityLease) { + if (this.sessionTimeout != null) { + asyncContext.addListener(activityLease); + } + } + /** * Handles DELETE requests for session deletion. * @param request The HTTP servlet request @@ -604,16 +841,26 @@ protected void doDelete(HttpServletRequest request, HttpServletResponse response } String sessionId = request.getHeader(HttpHeaders.MCP_SESSION_ID); - McpStreamableServerSession session = this.sessions.get(sessionId); + SessionEntry sessionEntry = this.sessions.get(sessionId); + + if (sessionEntry == null) { + response.sendError(HttpServletResponse.SC_NOT_FOUND); + return; + } - if (session == null) { + SessionActivityLease activityLease = acquireSessionActivity(sessionId, sessionEntry); + if (activityLease == null) { response.sendError(HttpServletResponse.SC_NOT_FOUND); return; } + McpStreamableServerSession session = sessionEntry.session(); + try { session.delete().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); - this.sessions.remove(sessionId); + if (this.sessions.remove(sessionId, sessionEntry)) { + sessionEntry.close(); + } response.setStatus(HttpServletResponse.SC_OK); } catch (Exception e) { @@ -627,6 +874,9 @@ protected void doDelete(HttpServletRequest request, HttpServletResponse response response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR, "Error deleting session"); } } + finally { + activityLease.close(); + } } public void responseError(HttpServletResponse response, int httpCode, McpError mcpError) throws IOException { @@ -687,10 +937,14 @@ private class HttpServletStreamableMcpSessionTransport implements McpStreamableS private final String sessionId; + private final SessionEntry sessionEntry; + private final AsyncContext asyncContext; private final PrintWriter writer; + private final SessionActivityLease activityLease; + private volatile boolean closed = false; private final ReentrantLock lock = new ReentrantLock(); @@ -698,13 +952,18 @@ private class HttpServletStreamableMcpSessionTransport implements McpStreamableS /** * Creates a new session transport with the specified ID and SSE writer. * @param sessionId The unique identifier for this session + * @param sessionEntry The session state associated with this transport * @param asyncContext The async context for the session * @param writer The writer for sending server events to the client + * @param activityLease The activity lease held while this transport is open */ - HttpServletStreamableMcpSessionTransport(String sessionId, AsyncContext asyncContext, PrintWriter writer) { + HttpServletStreamableMcpSessionTransport(String sessionId, SessionEntry sessionEntry, AsyncContext asyncContext, + PrintWriter writer, SessionActivityLease activityLease) { this.sessionId = sessionId; + this.sessionEntry = sessionEntry; this.asyncContext = asyncContext; this.writer = writer; + this.activityLease = activityLease; logger.debug("Streamable session transport {} initialized with SSE writer", sessionId); } @@ -747,7 +1006,11 @@ public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId } catch (Exception e) { logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage()); - HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId); + if (HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId, + this.sessionEntry)) { + this.sessionEntry.close(); + } + this.activityLease.close(); this.asyncContext.complete(); } finally { @@ -801,6 +1064,7 @@ public void close() { logger.warn("Failed to complete async context for session {}: {}", sessionId, e.getMessage()); } finally { + this.activityLease.close(); lock.unlock(); } } @@ -828,6 +1092,10 @@ public static class Builder { private Duration keepAliveInterval; + private Duration sessionTimeout; + + private Scheduler sessionCleanupScheduler; + private ServerTransportSecurityValidator securityValidator = ServerTransportSecurityValidator.NOOP; /** @@ -889,6 +1157,29 @@ public Builder keepAliveInterval(Duration keepAliveInterval) { return this; } + /** + * Sets the idle timeout after which a session becomes eligible for periodic + * cleanup. Idle session cleanup is disabled by default. + * @param sessionTimeout The session idle timeout. Must be greater than zero. If + * null, idle session cleanup is disabled. + * @return this builder instance + * @throws IllegalArgumentException if the timeout is non-null and less than one + * millisecond + */ + public Builder sessionTimeout(Duration sessionTimeout) { + if (sessionTimeout != null) { + validateSessionTimeout(sessionTimeout); + } + this.sessionTimeout = sessionTimeout; + return this; + } + + Builder sessionCleanupScheduler(Scheduler sessionCleanupScheduler) { + Assert.notNull(sessionCleanupScheduler, "Session cleanup scheduler must not be null"); + this.sessionCleanupScheduler = sessionCleanupScheduler; + return this; + } + /** * Sets the security validator for validating HTTP requests. * @param securityValidator The security validator to use. Must not be null. @@ -911,7 +1202,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, securityValidator); + contextExtractor, keepAliveInterval, sessionTimeout, sessionCleanupScheduler, securityValidator); } } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerSessionTimeoutTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerSessionTimeoutTests.java new file mode 100644 index 000000000..cdd2ddd63 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerSessionTimeoutTests.java @@ -0,0 +1,321 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.server.transport; + +import java.io.BufferedReader; +import java.io.PrintWriter; +import java.io.StringReader; +import java.io.StringWriter; +import java.time.Duration; +import java.util.Collections; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import jakarta.servlet.AsyncContext; +import jakarta.servlet.AsyncEvent; +import jakarta.servlet.AsyncListener; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; +import reactor.test.scheduler.VirtualTimeScheduler; + +import io.modelcontextprotocol.common.McpTransportContext; +import io.modelcontextprotocol.json.McpJsonMapper; +import io.modelcontextprotocol.server.McpNotificationHandler; +import io.modelcontextprotocol.spec.HttpHeaders; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpStreamableServerSession; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class HttpServletStreamableServerSessionTimeoutTests { + + private static final Duration SESSION_TIMEOUT = Duration.ofMinutes(5); + + private final McpJsonMapper jsonMapper = new GsonMcpJsonMapper(); + + @Test + void sessionCleanupIsDisabledByDefault() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + SessionFixture fixture = createInitializedSession(HttpServletStreamableServerTransportProvider.builder(), + scheduler, "default-session", Map.of()); + scheduler.advanceTimeBy(Duration.ofDays(1)); + + assertThat(fixture.closed()).isFalse(); + } + + @Test + void rejectsSessionTimeoutShorterThanOneMillisecond() { + assertThatIllegalArgumentException() + .isThrownBy( + () -> HttpServletStreamableServerTransportProvider.builder().sessionTimeout(Duration.ofNanos(1))) + .withMessage("Session timeout must be at least 1 millisecond"); + } + + @Test + void rejectsNonPositiveSessionTimeout() { + assertThatIllegalArgumentException() + .isThrownBy(() -> HttpServletStreamableServerTransportProvider.builder().sessionTimeout(Duration.ZERO)) + .withMessage("Session timeout must be greater than zero"); + } + + @Test + void idleSessionIsClosedAfterConfiguredTimeout() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + String sessionId = "idle-session"; + SessionFixture fixture = createInitializedSession(scheduler, sessionId); + scheduler.advanceTimeBy(SESSION_TIMEOUT); + + assertThat(fixture.closed()).isTrue(); + + HttpServletRequest deleteRequest = mock(HttpServletRequest.class); + HttpServletResponse deleteResponse = mock(HttpServletResponse.class); + when(deleteRequest.getRequestURI()).thenReturn("/mcp"); + when(deleteRequest.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + when(deleteRequest.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + + fixture.provider().doDelete(deleteRequest, deleteResponse); + + verify(deleteResponse).sendError(HttpServletResponse.SC_NOT_FOUND); + } + + @Test + void openGetStreamKeepsSessionAliveUntilStreamCloses() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + String sessionId = "active-get-session"; + SessionFixture fixture = createInitializedSession(scheduler, sessionId); + + HttpServletRequest getRequest = mock(HttpServletRequest.class); + HttpServletResponse getResponse = mock(HttpServletResponse.class); + AsyncContext asyncContext = mock(AsyncContext.class); + when(getRequest.getRequestURI()).thenReturn("/mcp"); + when(getRequest.getHeader("Accept")).thenReturn("text/event-stream"); + when(getRequest.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + when(getRequest.getHeader(HttpHeaders.LAST_EVENT_ID)).thenReturn(null); + when(getRequest.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + when(getRequest.startAsync()).thenReturn(asyncContext); + when(getResponse.getWriter()).thenReturn(new PrintWriter(new StringWriter(), true)); + + fixture.provider().doGet(getRequest, getResponse); + ArgumentCaptor listenerCaptor = ArgumentCaptor.forClass(AsyncListener.class); + verify(asyncContext, times(2)).addListener(listenerCaptor.capture()); + + scheduler.advanceTimeBy(Duration.ofMinutes(10)); + + assertThat(fixture.closed()).isFalse(); + + for (AsyncListener listener : listenerCaptor.getAllValues()) { + listener.onComplete(new AsyncEvent(asyncContext)); + } + scheduler.advanceTimeBy(SESSION_TIMEOUT); + + assertThat(fixture.closed()).isTrue(); + } + + @Test + void replayGetKeepsSessionAliveUntilRequestCompletes() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + String sessionId = "replay-get-session"; + SessionFixture fixture = createInitializedSession(scheduler, sessionId); + + HttpServletRequest getRequest = mock(HttpServletRequest.class); + HttpServletResponse getResponse = mock(HttpServletResponse.class); + AsyncContext asyncContext = mock(AsyncContext.class); + when(getRequest.getRequestURI()).thenReturn("/mcp"); + when(getRequest.getHeader("Accept")).thenReturn("text/event-stream"); + when(getRequest.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + when(getRequest.getHeader(HttpHeaders.LAST_EVENT_ID)).thenReturn("last-event"); + when(getRequest.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + when(getRequest.startAsync()).thenReturn(asyncContext); + when(getResponse.getWriter()).thenReturn(new PrintWriter(new StringWriter(), true)); + + fixture.provider().doGet(getRequest, getResponse); + ArgumentCaptor listenerCaptor = ArgumentCaptor.forClass(AsyncListener.class); + verify(asyncContext).addListener(listenerCaptor.capture()); + + scheduler.advanceTimeBy(Duration.ofMinutes(10)); + + assertThat(fixture.closed()).isFalse(); + + listenerCaptor.getValue().onComplete(new AsyncEvent(asyncContext)); + scheduler.advanceTimeBy(SESSION_TIMEOUT); + + assertThat(fixture.closed()).isTrue(); + } + + @Test + void getContextExtractionFailureDoesNotKeepSessionActive() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + HttpServletStreamableServerTransportProvider.Builder builder = HttpServletStreamableServerTransportProvider + .builder() + .contextExtractor(request -> { + if ("GET".equals(request.getMethod())) { + throw new IllegalStateException("context extraction failed"); + } + return McpTransportContext.EMPTY; + }) + .sessionTimeout(SESSION_TIMEOUT); + String sessionId = "context-failure-session"; + SessionFixture fixture = createInitializedSession(builder, scheduler, sessionId, Map.of()); + + HttpServletRequest getRequest = mock(HttpServletRequest.class); + HttpServletResponse getResponse = mock(HttpServletResponse.class); + when(getRequest.getMethod()).thenReturn("GET"); + when(getRequest.getRequestURI()).thenReturn("/mcp"); + when(getRequest.getHeader("Accept")).thenReturn("text/event-stream"); + when(getRequest.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + when(getRequest.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + + assertThatThrownBy(() -> fixture.provider().doGet(getRequest, getResponse)) + .isInstanceOf(IllegalStateException.class) + .hasMessage("context extraction failed"); + + scheduler.advanceTimeBy(SESSION_TIMEOUT); + + assertThat(fixture.closed()).isTrue(); + } + + @Test + void postActivityResetsSessionTimeout() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + String sessionId = "active-post-session"; + SessionFixture fixture = createInitializedSession(scheduler, sessionId); + + scheduler.advanceTimeBy(Duration.ofMinutes(4)); + postNotification(fixture.provider(), sessionId); + scheduler.advanceTimeBy(Duration.ofMinutes(1)); + + assertThat(fixture.closed()).isFalse(); + + scheduler.advanceTimeBy(Duration.ofMinutes(4)); + + assertThat(fixture.closed()).isTrue(); + } + + @Test + void inProgressPostKeepsSessionAliveUntilRequestCompletes() throws Exception { + VirtualTimeScheduler scheduler = VirtualTimeScheduler.create(); + CountDownLatch requestStarted = new CountDownLatch(1); + Sinks.Empty requestCompletion = Sinks.empty(); + McpNotificationHandler handler = (exchange, params) -> { + requestStarted.countDown(); + return requestCompletion.asMono(); + }; + String sessionId = "in-progress-post-session"; + SessionFixture fixture = createInitializedSession( + HttpServletStreamableServerTransportProvider.builder().sessionTimeout(SESSION_TIMEOUT), scheduler, + sessionId, Map.of("notifications/test", handler)); + + CompletableFuture post = CompletableFuture.runAsync(() -> { + try { + postNotification(fixture.provider(), sessionId); + } + catch (Exception ex) { + throw new CompletionException(ex); + } + }); + assertThat(requestStarted.await(5, TimeUnit.SECONDS)).isTrue(); + + scheduler.advanceTimeBy(Duration.ofMinutes(10)); + + assertThat(fixture.closed()).isFalse(); + + assertThat(requestCompletion.tryEmitEmpty()).isEqualTo(Sinks.EmitResult.OK); + post.get(5, TimeUnit.SECONDS); + scheduler.advanceTimeBy(SESSION_TIMEOUT); + + assertThat(fixture.closed()).isTrue(); + } + + private SessionFixture createInitializedSession(VirtualTimeScheduler scheduler, String sessionId) throws Exception { + return createInitializedSession( + HttpServletStreamableServerTransportProvider.builder().sessionTimeout(SESSION_TIMEOUT), scheduler, + sessionId, Map.of()); + } + + private SessionFixture createInitializedSession(HttpServletStreamableServerTransportProvider.Builder builder, + VirtualTimeScheduler scheduler, String sessionId, Map notificationHandlers) + throws Exception { + HttpServletStreamableServerTransportProvider provider = builder.jsonMapper(this.jsonMapper) + .sessionCleanupScheduler(scheduler) + .build(); + AtomicBoolean closed = new AtomicBoolean(); + McpStreamableServerSession session = createSession(sessionId, closed, notificationHandlers); + provider.setSessionFactory(request -> new McpStreamableServerSession.McpStreamableServerSessionInit(session, + Mono.just(testInitializeResult()))); + initializeSession(provider); + return new SessionFixture(provider, closed); + } + + private void initializeSession(HttpServletStreamableServerTransportProvider provider) throws Exception { + HttpServletRequest request = mock(HttpServletRequest.class); + HttpServletResponse response = mock(HttpServletResponse.class); + when(request.getRequestURI()).thenReturn("/mcp"); + when(request.getHeader("Accept")).thenReturn("text/event-stream, application/json"); + when(request.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + when(request.getReader()).thenReturn(new BufferedReader(new StringReader(this.jsonMapper.writeValueAsString( + new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "init-1", testInitializeRequest()))))); + when(response.getWriter()).thenReturn(new PrintWriter(new StringWriter(), true)); + + provider.doPost(request, response); + + verify(response).setStatus(HttpServletResponse.SC_OK); + } + + private void postNotification(HttpServletStreamableServerTransportProvider provider, String sessionId) + throws Exception { + HttpServletRequest request = mock(HttpServletRequest.class); + HttpServletResponse response = mock(HttpServletResponse.class); + when(request.getRequestURI()).thenReturn("/mcp"); + when(request.getHeader("Accept")).thenReturn("text/event-stream, application/json"); + when(request.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + when(request.getHeaderNames()).thenReturn(Collections.emptyEnumeration()); + when(request.getReader()).thenReturn(new BufferedReader(new StringReader(this.jsonMapper + .writeValueAsString(new McpSchema.JSONRPCNotification("notifications/test", Map.of("value", "test")))))); + + provider.doPost(request, response); + + verify(response).setStatus(HttpServletResponse.SC_ACCEPTED); + } + + private McpStreamableServerSession createSession(String sessionId, AtomicBoolean closed, + Map notificationHandlers) { + return new McpStreamableServerSession(sessionId, testInitializeRequest().capabilities(), + testInitializeRequest().clientInfo(), Duration.ofSeconds(2), Map.of(), notificationHandlers, + () -> Mono.fromRunnable(() -> closed.set(true))); + } + + private McpSchema.InitializeRequest testInitializeRequest() { + return McpSchema.InitializeRequest + .builder("2025-11-25", new McpSchema.ClientCapabilities(null, null, null, null), + new McpSchema.Implementation("test-client", "1.0.0")) + .build(); + } + + private McpSchema.InitializeResult testInitializeResult() { + return McpSchema.InitializeResult + .builder("2025-11-25", new McpSchema.ServerCapabilities(null, null, null, null, null, null), + new McpSchema.Implementation("test-server", "1.0.0")) + .build(); + } + + private record SessionFixture(HttpServletStreamableServerTransportProvider provider, AtomicBoolean closed) { + } + +}