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