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 <git@garnier.wf>

Author: Kehrlann

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java

@@ -9,7 +9,9 @@
 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,15 +30,19 @@
 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;
 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<String, McpStreamableServerSession> 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<Set<String>> activeSessions = new AtomicReference<>(ConcurrentHashMap.newKeySet());
+
 	private McpTransportContextExtractor<HttpServletRequest> 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<HttpServletRequest> 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 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S
 			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<String> 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 Mono<Void> closeGracefully() {
 			if (this.keepAliveScheduler != null) {
 				this.keepAliveScheduler.shutdown();
 			}
+			if (this.sessionSweeper != null) {
+				this.sessionSweeper.dispose();
+			}
 		});
 	}
 
@@ -323,6 +409,7 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
 			response.sendError(HttpServletResponse.SC_NOT_FOUND);
 			return;
 		}
+		this.markSessionActive(sessionId);
 
 		logger.debug("Handling GET request for session: {}", sessionId);
 
@@ -354,65 +441,14 @@ 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());
 			response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
 		}
 	}
 
-	/**
-	 * 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 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 						});
 				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 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
 				return;
 			}
 
+			this.markSessionActive(sessionId);
+
 			if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) {
 				session.accept(jsonrpcResponse)
 					.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
@@ -554,6 +593,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 +750,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.
@@ -719,7 +823,6 @@ public void destroy() {
 	 * 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 Mono<Void> 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 +933,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);
 			}
@@ -875,6 +977,8 @@ public static class Builder {
 
 		private int requestMaxSize = DEFAULT_REQUEST_MAX_SIZE;
 
+		private Duration sessionSweepInterval;
+
 		/**
 		 * Sets the JsonMapper to use for JSON serialization/deserialization of MCP
 		 * messages.
@@ -984,7 +1088,23 @@ 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, httpHeaderValidator, requestMaxSize, sessionSweepInterval);
+		}
+
+		/**
+		 * 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.
+		 * @return this builder instance
+		 */
+		public Builder sessionSweepInterval(Duration sessionSweepInterval) {
+			this.sessionSweepInterval = sessionSweepInterval;
+			return this;
 		}
 
 	}