Skip to content

Commit bf508ca

Browse files
committed
wip evict sessions
1 parent 1ba72bd commit bf508ca

5 files changed

Lines changed: 511 additions & 87 deletions

File tree

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

Lines changed: 168 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -28,15 +28,19 @@
2828
import io.modelcontextprotocol.util.Assert;
2929
import io.modelcontextprotocol.util.KeepAliveScheduler;
3030
import jakarta.servlet.AsyncContext;
31+
import jakarta.servlet.AsyncEvent;
32+
import jakarta.servlet.AsyncListener;
3133
import jakarta.servlet.ServletException;
3234
import jakarta.servlet.annotation.WebServlet;
3335
import jakarta.servlet.http.HttpServlet;
3436
import jakarta.servlet.http.HttpServletRequest;
3537
import jakarta.servlet.http.HttpServletResponse;
3638
import org.slf4j.Logger;
3739
import org.slf4j.LoggerFactory;
40+
import reactor.core.Disposable;
3841
import reactor.core.publisher.Flux;
3942
import reactor.core.publisher.Mono;
43+
import reactor.core.scheduler.Schedulers;
4044

4145
/**
4246
* Server-side implementation of the Model Context Protocol (MCP) streamable transport
@@ -90,6 +94,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
9094

9195
public static final String FAILED_TO_SEND_ERROR_RESPONSE = "Failed to send error response: {}";
9296

97+
private static final Duration MAX_IDLE_SWEEP_INTERVAL = Duration.ofMinutes(1);
98+
9399
/**
94100
* The endpoint URI where clients should send their JSON-RPC messages. Defaults to
95101
* "/mcp".
@@ -128,6 +134,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
128134
*/
129135
private KeepAliveScheduler keepAliveScheduler;
130136

137+
/**
138+
* How long a session may stay idle before being evicted. Disabled when {@code null}.
139+
*/
140+
private final Duration sessionIdleTimeout;
141+
142+
/**
143+
* Subscription of the periodic sweep evicting idle sessions. Only set when
144+
* {@link #sessionIdleTimeout} is configured.
145+
*/
146+
private Disposable idleSessionSweep;
147+
131148
/**
132149
* Security validator for validating HTTP requests.
133150
*/
@@ -143,14 +160,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
143160
* @param contextExtractor The extractor for transport context from the request.
144161
* @param keepAliveInterval The interval for keep-alive pings. If null, no keep-alive
145162
* will be scheduled.
163+
* @param sessionIdleTimeout How long a session may stay idle before being evicted. If
164+
* null, idle sessions are never evicted.
146165
* @param httpHeaderValidator The HTTP header validator for validating HTTP requests.
147166
* @param requestMaxSize The maximum size, in bytes, of a single request body. Must be
148167
* positive.
149168
* @throws IllegalArgumentException if any parameter is null
150169
*/
151170
private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint,
152171
boolean disallowDelete, McpTransportContextExtractor<HttpServletRequest> contextExtractor,
153-
Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize) {
172+
Duration keepAliveInterval, Duration sessionIdleTimeout, ServerHttpHeaderValidator httpHeaderValidator,
173+
int requestMaxSize) {
154174
Assert.notNull(jsonMapper, "JsonMapper must not be null");
155175
Assert.notNull(mcpEndpoint, "MCP endpoint must not be null");
156176
Assert.notNull(contextExtractor, "Context extractor must not be null");
@@ -163,12 +183,35 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S
163183
this.contextExtractor = contextExtractor;
164184
this.httpHeaderValidator = httpHeaderValidator;
165185
this.requestMaxSize = requestMaxSize;
186+
this.sessionIdleTimeout = sessionIdleTimeout;
187+
188+
if (sessionIdleTimeout != null) {
189+
Assert.isTrue(!sessionIdleTimeout.isNegative() && !sessionIdleTimeout.isZero(),
190+
"Session idle timeout must be positive");
191+
// Sweeping more often than the timeout bounds how long an idle session
192+
// outlives it, without making the sweep itself expensive
193+
Duration sweepInterval = sessionIdleTimeout.compareTo(MAX_IDLE_SWEEP_INTERVAL) < 0 ? sessionIdleTimeout
194+
: MAX_IDLE_SWEEP_INTERVAL;
195+
// Evicting closes connections, which is not work for the parallel scheduler
196+
this.idleSessionSweep = Flux.interval(sweepInterval, Schedulers.boundedElastic())
197+
.doOnNext(tick -> this.evictIdleSessions())
198+
.onErrorContinue((error, tick) -> logger.warn("Idle session sweep failed: {}", error.getMessage()))
199+
.subscribe();
200+
}
166201

167202
if (keepAliveInterval != null) {
168203

169204
this.keepAliveScheduler = KeepAliveScheduler.builder(this::sessionsToPing)
170205
.initialDelay(keepAliveInterval)
171206
.interval(keepAliveInterval)
207+
.onPingFailure(session -> {
208+
// The stream the ping was written to is dead. The session survives:
209+
// the client may reconnect to it, and the idle timeout reclaims it if
210+
// it never does.
211+
if (session instanceof McpStreamableServerSession streamableSession) {
212+
streamableSession.releaseListeningStream();
213+
}
214+
})
172215
.build();
173216

174217
this.keepAliveScheduler.start();
@@ -192,6 +235,40 @@ private Flux<McpSession> sessionsToPing() {
192235
.cast(McpSession.class);
193236
}
194237

238+
/**
239+
* Removes a session and releases everything it holds, in particular the connections
240+
* of its streams. Together with the client asking for it through a DELETE request,
241+
* this is the only way a session leaves this transport.
242+
* @param sessionId the session to evict
243+
* @param reason why the session is being evicted, for logging
244+
*/
245+
private void evictSession(String sessionId, String reason) {
246+
McpStreamableServerSession session = this.sessions.remove(sessionId);
247+
if (session == null) {
248+
return;
249+
}
250+
logger.info("Evicting session {}: {}", sessionId, reason);
251+
session.close();
252+
}
253+
254+
/**
255+
* Evicts the sessions idle for longer than the configured session idle timeout.
256+
* Clients are not required to send a DELETE request when they go away, so a client
257+
* which simply vanishes, whether killed, disconnected or suspended, would otherwise
258+
* keep its session, and the resources behind it, until the server restarts.
259+
*/
260+
private void evictIdleSessions() {
261+
if (this.isClosing) {
262+
return;
263+
}
264+
this.sessions.values()
265+
.stream()
266+
.filter(session -> session.isIdleFor(this.sessionIdleTimeout))
267+
.map(McpStreamableServerSession::getId)
268+
.toList()
269+
.forEach(sessionId -> this.evictSession(sessionId, "idle for more than " + this.sessionIdleTimeout));
270+
}
271+
195272
@Override
196273
public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) {
197274
this.sessionFactory = sessionFactory;
@@ -264,6 +341,9 @@ public Mono<Void> closeGracefully() {
264341
if (this.keepAliveScheduler != null) {
265342
this.keepAliveScheduler.shutdown();
266343
}
344+
if (this.idleSessionSweep != null) {
345+
this.idleSessionSweep.dispose();
346+
}
267347
});
268348
}
269349

@@ -354,65 +434,14 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
354434
McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
355435
.listeningStream(sessionTransport);
356436

357-
asyncContext.addListener(new jakarta.servlet.AsyncListener() {
358-
@Override
359-
public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
360-
logger.debug("SSE connection completed for session: {}", sessionId);
361-
listeningStream.close();
362-
}
363-
364-
@Override
365-
public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
366-
logger.debug("SSE connection timed out for session: {}", sessionId);
367-
listeningStream.close();
368-
}
369-
370-
@Override
371-
public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
372-
logger.debug("SSE connection error for session: {}", sessionId);
373-
listeningStream.close();
374-
}
375-
376-
@Override
377-
public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
378-
// No action needed
379-
}
380-
});
437+
registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport);
381438
}
382439
catch (Exception e) {
383440
logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage());
384441
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
385442
}
386443
}
387444

388-
/**
389-
* Replays the messages the client missed while its SSE stream was broken.
390-
* @param session the session the client is resuming
391-
* @param lastEventId the ID of the last event received by the client
392-
* @param sessionTransport the transport of the resumed SSE stream
393-
* @param transportContext the context extracted from the request
394-
* @return {@code true} if the replay completed, {@code false} if it failed, in which
395-
* case the transport has been closed
396-
*/
397-
private boolean tryReplayMissedMessages(McpStreamableServerSession session, String lastEventId,
398-
McpStreamableServerTransport sessionTransport, McpTransportContext transportContext) {
399-
try {
400-
for (McpSchema.JSONRPCMessage message : session.replay(lastEventId)
401-
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
402-
.toIterable()) {
403-
sessionTransport.sendMessage(message)
404-
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
405-
.block();
406-
}
407-
return true;
408-
}
409-
catch (Exception e) {
410-
logger.error("Failed to replay messages for session {}: {}", session.getId(), e.getMessage());
411-
sessionTransport.close();
412-
return false;
413-
}
414-
}
415-
416445
/**
417446
* Handles POST requests for incoming JSON-RPC messages from clients.
418447
* @param request The HTTP servlet request containing the JSON-RPC message
@@ -554,6 +583,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
554583

555584
HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
556585
sessionId, asyncContext, response.getWriter());
586+
registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close);
557587

558588
try {
559589
session.responseStream(jsonrpcRequest, sessionTransport)
@@ -710,6 +740,70 @@ public void destroy() {
710740
super.destroy();
711741
}
712742

743+
/**
744+
* Replays the messages the client missed while its SSE stream was broken.
745+
* @param session the session the client is resuming
746+
* @param lastEventId the ID of the last event received by the client
747+
* @param sessionTransport the transport of the resumed SSE stream
748+
* @param transportContext the context extracted from the request
749+
* @return {@code true} if the replay completed, {@code false} if it failed, in which
750+
* case the transport has been closed
751+
*/
752+
private static boolean tryReplayMissedMessages(McpStreamableServerSession session, String lastEventId,
753+
McpStreamableServerTransport sessionTransport, McpTransportContext transportContext) {
754+
try {
755+
for (McpSchema.JSONRPCMessage message : session.replay(lastEventId)
756+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
757+
.toIterable()) {
758+
sessionTransport.sendMessage(message)
759+
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
760+
.block();
761+
}
762+
return true;
763+
}
764+
catch (Exception e) {
765+
logger.error("Failed to replay messages for session {}: {}", session.getId(), e.getMessage());
766+
sessionTransport.close();
767+
return false;
768+
}
769+
}
770+
771+
/**
772+
* Registers a listener releasing the SSE stream carried by the given asynchronous
773+
* request once the container is done with it, whether the client went away, the
774+
* request timed out or it errored. Without this, the connection is left open, holding
775+
* on to a socket and a container thread, until the process restarts.
776+
* @param asyncContext the asynchronous context of the SSE request
777+
* @param sessionId the session the stream belongs to
778+
* @param onConnectionEnd the action releasing the stream
779+
*/
780+
private static void registerAsyncLifecycle(AsyncContext asyncContext, String sessionId, Runnable onConnectionEnd) {
781+
asyncContext.addListener(new AsyncListener() {
782+
@Override
783+
public void onComplete(AsyncEvent event) throws IOException {
784+
logger.debug("SSE connection completed for session: {}", sessionId);
785+
onConnectionEnd.run();
786+
}
787+
788+
@Override
789+
public void onTimeout(AsyncEvent event) throws IOException {
790+
logger.debug("SSE connection timed out for session: {}", sessionId);
791+
onConnectionEnd.run();
792+
}
793+
794+
@Override
795+
public void onError(AsyncEvent event) throws IOException {
796+
logger.debug("SSE connection error for session: {}", sessionId);
797+
onConnectionEnd.run();
798+
}
799+
800+
@Override
801+
public void onStartAsync(AsyncEvent event) throws IOException {
802+
// No action needed
803+
}
804+
});
805+
}
806+
713807
/**
714808
* Implementation of McpStreamableServerTransport for HttpServlet SSE sessions. This
715809
* class handles the transport-level communication for a specific client session.
@@ -783,9 +877,10 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId
783877
logger.debug("Message sent to session {} with ID {}", this.sessionId, messageId);
784878
}
785879
catch (Exception e) {
880+
// The connection is gone, the session is not: the client may come
881+
// back for it, and the idle timeout reclaims it if it never does
786882
logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage());
787-
HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId);
788-
this.asyncContext.complete();
883+
this.close();
789884
}
790885
finally {
791886
lock.unlock();
@@ -829,8 +924,6 @@ public void close() {
829924
}
830925

831926
this.closed = true;
832-
833-
// HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId);
834927
this.asyncContext.complete();
835928
logger.debug("Successfully completed async context for session {}", sessionId);
836929
}
@@ -871,6 +964,8 @@ public static class Builder {
871964

872965
private Duration keepAliveInterval;
873966

967+
private Duration sessionIdleTimeout;
968+
874969
private ServerHttpHeaderValidator httpHeaderValidator = ServerHttpHeaderValidator.NOOP;
875970

876971
private int requestMaxSize = DEFAULT_REQUEST_MAX_SIZE;
@@ -929,6 +1024,21 @@ public Builder contextExtractor(McpTransportContextExtractor<HttpServletRequest>
9291024
* keep-alive will be scheduled.
9301025
* @return this builder instance
9311026
*/
1027+
/**
1028+
* Sets how long a session may stay idle, that is hold no connection and serve no
1029+
* request, before the server evicts it. Clients are expected to end their session
1030+
* with a DELETE request, but nothing guarantees they get the chance to, so
1031+
* without this the sessions of vanished clients are kept until the server
1032+
* restarts.
1033+
* @param sessionIdleTimeout the idle timeout. If null, idle sessions are never
1034+
* evicted.
1035+
* @return this builder instance
1036+
*/
1037+
public Builder sessionIdleTimeout(Duration sessionIdleTimeout) {
1038+
this.sessionIdleTimeout = sessionIdleTimeout;
1039+
return this;
1040+
}
1041+
9321042
public Builder keepAliveInterval(Duration keepAliveInterval) {
9331043
this.keepAliveInterval = keepAliveInterval;
9341044
return this;
@@ -984,7 +1094,7 @@ public HttpServletStreamableServerTransportProvider build() {
9841094
Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set");
9851095
return new HttpServletStreamableServerTransportProvider(
9861096
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, mcpEndpoint, disallowDelete,
987-
contextExtractor, keepAliveInterval, httpHeaderValidator, requestMaxSize);
1097+
contextExtractor, keepAliveInterval, sessionIdleTimeout, httpHeaderValidator, requestMaxSize);
9881098
}
9891099

9901100
}

0 commit comments

Comments
 (0)