diff --git a/docs/server.md b/docs/server.md index 3c05aee30..9d654def1 100644 --- a/docs/server.md +++ b/docs/server.md @@ -161,7 +161,10 @@ The SDK includes several built-in transport provider implementations: - Efficient bidirectional HTTP communication - Session management for multiple client connections - - Configurable keep-alive intervals + - Keep-alive pings on sessions with an open stream, enabled by default every 30 minutes + (`keepAliveInterval`, `null` to disable) + - Eviction of idle sessions — no open stream and no request for a full interval — every + 30 minutes by default (`sessionSweepInterval`, `null` to keep sessions until deleted) - Security validation support - Graceful shutdown support 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 c9dbb2439..c339c22a6 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 @@ -10,7 +10,9 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; +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; @@ -21,13 +23,17 @@ import io.modelcontextprotocol.spec.HttpHeaders; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpSession; import io.modelcontextprotocol.spec.McpStreamableServerSession; import io.modelcontextprotocol.spec.McpStreamableServerTransport; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; +import io.modelcontextprotocol.spec.McpTransportException; 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 +41,10 @@ import jakarta.servlet.http.HttpServletResponse; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; /** * Server-side implementation of the Model Context Protocol (MCP) streamable transport @@ -53,6 +61,7 @@ * @author Zachary German * @author Christian Tzolov * @author Dariusz Jędrzejczyk + * @author Daniel Garnier-Moiroux * @see McpStreamableServerTransportProvider * @see HttpServlet */ @@ -115,6 +124,14 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private final ConcurrentHashMap 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. Only populated when a sweeper runs, as nothing would ever swap it + * otherwise. + */ + private final AtomicReference> activeSessions = new AtomicReference<>(ConcurrentHashMap.newKeySet()); + private McpTransportContextExtractor contextExtractor; /** @@ -128,6 +145,14 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet */ private KeepAliveScheduler keepAliveScheduler; + /** + * Periodic eviction of the sessions no client came back to. {@code null} if no + * sessionSweepInterval is set, in which case sessions are never reclaimed. + */ + private final Disposable sessionSweeper; + + private static Duration SESSION_SWEEP_TIMEOUT = Duration.ofSeconds(30); + /** * Security validator for validating HTTP requests. */ @@ -146,11 +171,14 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet * @param securityValidator The security 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 contextExtractor, - Duration keepAliveInterval, ServerTransportSecurityValidator securityValidator, int requestMaxSize) { + Duration keepAliveInterval, ServerTransportSecurityValidator securityValidator, 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"); @@ -166,15 +194,96 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S if (keepAliveInterval != null) { - this.keepAliveScheduler = KeepAliveScheduler - .builder(() -> (isClosing) ? Flux.empty() : Flux.fromIterable(sessions.values())) + this.keepAliveScheduler = KeepAliveScheduler.builder(this::sessionsToPing) .initialDelay(keepAliveInterval) .interval(keepAliveInterval) + .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(); } + this.sessionSweeper = sessionSweepInterval == null ? null : startSessionSweeper(sessionSweepInterval); + } + + /** + * Schedules the periodic eviction of idle sessions. + * @param sessionSweepInterval the interval between two sweeps + * @return the handle on the scheduled sweeps, to be disposed when the transport shuts + * down + */ + private Disposable startSessionSweeper(Duration sessionSweepInterval) { + return Flux + .interval(sessionSweepInterval, sessionSweepInterval, + Schedulers.newSingle("streamable-http-server-transport-session-sweeper")) + .concatMap( + tick -> sweepSessions().doOnError(e -> logger.error("Session sweep failed", e)).onErrorComplete()) + .subscribe(next -> { + }, error -> logger.error("Session sweeper error", error)); + } + + /** + * 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. + * @return a Mono completing once every evicted session has been closed + */ + private Mono sweepSessions() { + if (this.isClosing) { + return Mono.empty(); + } + Set active = this.activeSessions.getAndSet(ConcurrentHashMap.newKeySet()); + return Flux.fromIterable(this.sessions.values()).filter(session -> { + if (session.hasOpenStream() || active.contains(session.getId())) { + return false; + } + return this.sessions.remove(session.getId(), session); + }).flatMap(session -> { + logger.debug("Evicting idle session {}", session.getId()); + return session.closeGracefully() + .timeout(SESSION_SWEEP_TIMEOUT) + .doOnError(e -> logger.warn("Failed to close idle session {}: {}", session.getId(), e.getMessage())) + .onErrorComplete(); + }).then(); + } + + /** + * 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) { + if (this.sessionSweeper == null) { + // No sweep ever swaps the set of active sessions, so recording activity would + // only accumulate the ID of every session which ever issued a request + return; + } + this.activeSessions.get().add(sessionId); + } + + /** + * Returns the sessions a keep-alive ping can be sent to, that is the sessions having + * a listening stream. A session without one, e.g. a client which only ever issues + * POST requests, has nothing to write a ping to: pinging it would fail on every + * interval without ever telling us anything about the client being alive. + * @return the sessions to ping + */ + private Flux sessionsToPing() { + if (this.isClosing) { + return Flux.empty(); + } + return Flux.fromIterable(this.sessions.values()) + .filter(McpStreamableServerSession::hasListeningStream) + .cast(McpSession.class); } @Override @@ -244,6 +353,9 @@ public Mono closeGracefully() { if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.sessionSweeper != null) { + this.sessionSweeper.dispose(); + } }); } @@ -303,6 +415,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); @@ -320,61 +433,21 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( sessionId, asyncContext, response.getWriter()); - // Check if this is a replay request - if (request.getHeader(HttpHeaders.LAST_EVENT_ID) != null) { - String lastId = request.getHeader(HttpHeaders.LAST_EVENT_ID); - - try { - session.replay(lastId) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .toIterable() - .forEach(message -> { - try { - sessionTransport.sendMessage(message) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .block(); - } - catch (Exception e) { - logger.error("Failed to replay message: {}", e.getMessage()); - asyncContext.complete(); - } - }); - } - catch (Exception e) { - logger.error("Failed to replay messages: {}", e.getMessage()); - asyncContext.complete(); - } + // Replay the messages the client missed while its stream was broken + String lastEventId = request.getHeader(HttpHeaders.LAST_EVENT_ID); + if (lastEventId != null + && !this.tryReplayMissedMessages(session, lastEventId, sessionTransport, transportContext)) { + // The replay failed and already closed the transport + return; } - else { - // Establish new listening stream - 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(); - } + // Establish the listening stream. Resumed streams are registered too, so + // that the session keeps delivering messages to the reconnected client and + // the async context is completed once the client goes away. + McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session + .listeningStream(sessionTransport); - @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()); @@ -456,6 +529,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 { @@ -502,6 +576,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)) @@ -527,8 +603,16 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( sessionId, asyncContext, response.getWriter()); + // The listener is given the stream rather than its transport, so that the + // end of the connection detaches the stream from the session instead of + // only dropping the socket: a stream outliving its connection keeps the + // session looking busy and spares it from the sweeper + McpStreamableServerSession.McpStreamableServerSessionStream responseStream = session + .responseStream(sessionTransport); + registerAsyncLifecycle(asyncContext, sessionId, responseStream::releaseTransport); + try { - session.responseStream(jsonrpcRequest, sessionTransport) + responseStream.handle(jsonrpcRequest) .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) .block(); } @@ -677,6 +761,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. @@ -686,7 +834,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; @@ -750,9 +897,17 @@ public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId logger.debug("Message sent to session {} with ID {}", this.sessionId, messageId); } catch (Exception e) { + // The connection is gone, the session is not: the client may come + // back for it, and the sweeper 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(); + // Surfaced to the caller rather than swallowed: whoever is writing to + // this stream has to learn that it no longer leads anywhere, or it + // keeps producing messages for a client which is gone. A request + // being streamed a response would never finish, holding on to the + // container thread which has to be given back before the end of the + // connection can be acted upon. + throw new McpTransportException("Failed to send message to session " + this.sessionId, e); } finally { lock.unlock(); @@ -796,8 +951,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); } @@ -830,12 +983,14 @@ public static class Builder { private McpTransportContextExtractor contextExtractor = ( serverRequest) -> McpTransportContext.EMPTY; - private Duration keepAliveInterval; + private Duration keepAliveInterval = Duration.ofMinutes(30); private ServerTransportSecurityValidator securityValidator = ServerTransportSecurityValidator.NOOP; 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. @@ -887,7 +1042,7 @@ public Builder contextExtractor(McpTransportContextExtractor * Sets the keep-alive interval for the transport. If set, a keep-alive scheduler * will be activated to periodically ping active sessions. * @param keepAliveInterval The interval for keep-alive pings. If null, no - * keep-alive will be scheduled. + * keep-alive will be scheduled. Defaults to 30 minutes. * @return this builder instance */ public Builder keepAliveInterval(Duration keepAliveInterval) { @@ -920,6 +1075,22 @@ public Builder maxRequestSize(int requestMaxSize) { 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. @@ -930,7 +1101,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, requestMaxSize); + contextExtractor, keepAliveInterval, securityValidator, requestMaxSize, sessionSweepInterval); } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java index fe2b6d798..420e21972 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -1,16 +1,20 @@ /* - * Copyright 2024-2025 the original author or authors. + * Copyright 2024-2026 the original author or authors. */ package io.modelcontextprotocol.spec; import java.time.Duration; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Predicate; import java.util.function.Supplier; import org.slf4j.Logger; @@ -36,12 +40,26 @@ * * @author Dariusz Jędrzejczyk * @author Yanming Zhou + * @author Daniel Garnier-Moiroux */ public class McpStreamableServerSession implements McpLoggableSession { private static final Logger logger = LoggerFactory.getLogger(McpStreamableServerSession.class); - private final ConcurrentHashMap requestIdToStream = new ConcurrentHashMap<>(); + /** + * Every server-initiated request still awaiting its response, keyed by request ID. + * Tracked per session rather than per stream, because the client answers over a + * separate HTTP POST which outlives the stream the request was sent on: a stream + * going away must not lose the requests it carried, but the end of the session must + * resolve all of them. + */ + private final ConcurrentHashMap pendingRequests = new ConcurrentHashMap<>(); + + /** + * Every stream with a connection currently attached, whether the listening stream or + * a POST response stream, so that they can all be released when the session ends. + */ + private final Set openStreams = ConcurrentHashMap.newKeySet(); private final String id; @@ -134,72 +152,91 @@ public Mono delete() { } /** - * Create a listening stream (the generic HTTP GET request without Last-Event-ID - * header). + * Whether the session currently has a listening stream, that is a stream the server + * can send its own requests and notifications to. Sessions have none until the client + * issues the GET request establishing one, and clients are not required to ever issue + * it. + * @return {@code true} if the session has a listening stream + */ + public boolean hasListeningStream() { + return this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream; + } + + /** + * 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 + * the stream being replaced, if any, is closed: no message would ever be sent to it + * again, and leaving it open would leak the underlying connection. * @param transport The dedicated SSE transport stream * @return a stream representation */ public McpStreamableServerSessionStream listeningStream(McpStreamableServerTransport transport) { McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport); - this.listeningStreamRef.set(listeningStream); + McpLoggableSession replaced = this.listeningStreamRef.getAndSet(listeningStream); + if (replaced instanceof McpStreamableServerSessionStream replacedStream) { + logger.debug("Releasing the connection of the listening stream replaced in session {}", this.id); + replacedStream.releaseTransport(); + } return listeningStream; } + /** + * Releases the connection of the listening stream, if one is attached, leaving the + * session without one until the client establishes a new stream. Used when the + * connection turns out to be dead, typically because a keep-alive ping went + * unanswered, so that the socket behind it is not held on to for nothing. + */ + public void releaseListeningStream() { + if (this.listeningStreamRef.get() instanceof McpStreamableServerSessionStream stream) { + stream.releaseTransport(); + } + } + // TODO: keep track of history by keeping a map from eventId to stream and then // iterate over the events using the lastEventId public Flux replay(Object lastEventId) { return Flux.empty(); } + /** + * Create a response stream (the SSE stream of a single HTTP POST request, finalized + * with the response to the request it carries). The caller owns the returned stream + * and is responsible for releasing it once the connection behind it ends, the same + * way it does for a {@link #listeningStream(McpStreamableServerTransport)}: a stream + * outliving its connection keeps the session looking busy, see + * {@link #hasOpenStream()}. + * @param transport the SSE transport stream to send messages to + * @return a stream representation, on which + * {@link McpStreamableServerSessionStream#handle(McpSchema.JSONRPCRequest)} runs the + * request + */ + public McpStreamableServerSessionStream responseStream(McpStreamableServerTransport transport) { + return new McpStreamableServerSessionStream(transport); + } + /** * Provide the SSE stream of MCP messages finalized with a Response. * @param jsonrpcRequest the MCP request triggering the stream creation * @param transport the SSE transport stream to send messages to * @return Mono which completes once the processing is done + * @deprecated the stream created for the request is not exposed, which leaves the + * caller unable to release it when the connection carrying it ends. Use + * {@link #responseStream(McpStreamableServerTransport)} and + * {@link McpStreamableServerSessionStream#handle(McpSchema.JSONRPCRequest)} instead. */ + @Deprecated public Mono responseStream(McpSchema.JSONRPCRequest jsonrpcRequest, McpStreamableServerTransport transport) { - return Mono.deferContextual(ctx -> { - McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); - - McpStreamableServerSessionStream stream = new McpStreamableServerSessionStream(transport); - McpRequestHandler requestHandler = McpStreamableServerSession.this.requestHandlers - .get(jsonrpcRequest.method()); - // TODO: delegate to stream, which upon successful response should close - // remove itself from the registry and also close the underlying transport - // (sink) - if (requestHandler == null) { - MethodNotFoundError error = getMethodNotFoundError(jsonrpcRequest.method()); - return transport - .sendMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), null, - new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND, - error.message(), error.data()))); - } - return requestHandler - .handle(new McpAsyncServerExchange(this.id, stream, clientCapabilities.get(), clientInfo.get(), - transportContext), jsonrpcRequest.params()) - .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), result, - null)) - .onErrorResume(e -> { - McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError; - if (e instanceof McpError mcpError && mcpError.getJsonRpcError() != null) { - jsonRpcError = mcpError.getJsonRpcError(); - } - else { - // The cause chain may contain internal details: log it, but only - // send - // the message to the client - logger.warn("Error handling request {} [{}]", jsonrpcRequest.id(), jsonrpcRequest.method(), e); - jsonRpcError = new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, - Objects.requireNonNullElse(e.getMessage(), "Internal error"), null); - } - - var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), - null, jsonRpcError); - return Mono.just(errorResponse); - }) - .flatMap(transport::sendMessage) - .then(transport.closeGracefully()); - }); + return Mono.defer(() -> this.responseStream(transport).handle(jsonrpcRequest)); } /** @@ -232,22 +269,13 @@ public Mono accept(McpSchema.JSONRPCResponse response) { logger.debug("Received response: {}", response); if (response.id() != null) { - var stream = this.requestIdToStream.get(response.id()); - if (stream == null) { + var pendingRequest = this.pendingRequests.remove(response.id()); + if (pendingRequest == null) { return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR) .message("Unexpected response for unknown id " + response.id()) .build()); } - // TODO: encapsulate this inside the stream itself - var sink = stream.pendingResponses.remove(response.id()); - if (sink == null) { - return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR) - .message("Unexpected response for unknown id " + response.id()) - .build()); - } - else { - sink.success(response); - } + pendingRequest.sink().success(response); } else { logger.error("Discarded MCP request response without session id. " @@ -265,22 +293,48 @@ private MethodNotFoundError getMethodNotFoundError(String method) { return new MethodNotFoundError(method, "Method not found: " + method, null); } + /** + * Fails the matching pending requests with the given reason. Each request is removed + * before its sink is failed, so that the cleanup the failure triggers on the + * requesting side finds nothing left to remove. + * @param match selects the requests to fail + * @param reason the message of the error the sinks are failed with + */ + private void failPendingRequests(Predicate match, String reason) { + List requestIds = new ArrayList<>(); + this.pendingRequests.forEach((requestId, pendingRequest) -> { + if (match.test(pendingRequest)) { + requestIds.add(requestId); + } + }); + for (String requestId : requestIds) { + PendingRequest pendingRequest = this.pendingRequests.remove(requestId); + if (pendingRequest != null) { + pendingRequest.sink().error(new RuntimeException(reason)); + } + } + } + @Override public Mono closeGracefully() { return Mono.defer(() -> { - McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); - return listeningStream.closeGracefully(); - // TODO: Also close all the open streams + this.listeningStreamRef.set(this.missingMcpTransportSession); + return Flux.fromIterable(List.copyOf(this.openStreams)) + .flatMap(McpStreamableServerSessionStream::closeGracefully) + .then() + // Also the requests of the streams already released: the session ending + // is the point at which no client can answer them anymore + .then(Mono.fromRunnable(() -> this.failPendingRequests(request -> true, "Session closed"))); }); } @Override public void close() { - McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession); - if (listeningStream != null) { - listeningStream.close(); - } - // TODO: Also close all open streams + this.listeningStreamRef.set(this.missingMcpTransportSession); + List.copyOf(this.openStreams).forEach(McpStreamableServerSessionStream::close); + // Also the requests of the streams already released: the session ending is the + // point at which no client can answer them anymore + this.failPendingRequests(request -> true, "Session closed"); } /** @@ -322,15 +376,23 @@ public record McpStreamableServerSessionInit(McpStreamableServerSession session, Mono initResult) { } + /** + * A server-initiated request awaiting its response, along with the stream it was sent + * on. + * + * @param stream the stream the request was sent on + * @param sink the sink to resolve once the response arrives + */ + private record PendingRequest(McpStreamableServerSessionStream stream, MonoSink sink) { + } + /** * An individual SSE stream within a Streamable HTTP context. Can be either the * listening GET SSE stream or a request-specific POST SSE stream. */ public final class McpStreamableServerSessionStream implements McpLoggableSession { - private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); - - private final McpStreamableServerTransport transport; + private final McpStreamableServerTransport connection; private final String transportId; @@ -338,10 +400,11 @@ public final class McpStreamableServerSessionStream implements McpLoggableSessio /** * Constructor accepting the dedicated transport representing the SSE stream. - * @param transport request-specific SSE transport stream + * @param connection request-specific SSE transport stream */ - public McpStreamableServerSessionStream(McpStreamableServerTransport transport) { - this.transport = transport; + public McpStreamableServerSessionStream(McpStreamableServerTransport connection) { + this.connection = connection; + McpStreamableServerSession.this.openStreams.add(this); this.transportId = UUID.randomUUID().toString(); // This ID design allows for a constant-time extraction of the history by // precisely identifying the SSE stream using the first component @@ -363,19 +426,20 @@ public boolean isNotificationForLevelAllowed(McpSchema.LoggingLevel loggingLevel public Mono sendRequest(String method, Object requestParams, TypeRef typeRef) { String requestId = McpStreamableServerSession.this.generateRequestId(); - McpStreamableServerSession.this.requestIdToStream.put(requestId, this); - return Mono.create(sink -> { - this.pendingResponses.put(requestId, sink); + // Registered on subscription rather than on assembly, so that a request + // which is never sent does not leave the session tracking it forever + McpStreamableServerSession.this.pendingRequests.put(requestId, new PendingRequest(this, sink)); McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams); String messageId = this.uuidGenerator.get(); // TODO: store message in history - this.transport.sendMessage(jsonrpcRequest, messageId).subscribe(v -> { + this.connection.sendMessage(jsonrpcRequest, messageId).subscribe(v -> { }, sink::error); - }).timeout(requestTimeout).doOnError(e -> { - this.pendingResponses.remove(requestId); - McpStreamableServerSession.this.requestIdToStream.remove(requestId); + }).timeout(requestTimeout).doFinally(signal -> { + // Also on completion and cancellation: a resolved request keeps no state, + // and a deadline imposed by the caller cancels rather than errors + McpStreamableServerSession.this.pendingRequests.remove(requestId); }).handle((jsonRpcResponse, sink) -> { if (jsonRpcResponse.error() != null) { sink.error(new McpError(jsonRpcResponse.error())); @@ -385,43 +449,115 @@ public Mono sendRequest(String method, Object requestParams, TypeRef t sink.complete(); } else { - sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef)); + sink.next(this.connection.unmarshalFrom(jsonRpcResponse.result(), typeRef)); } } }); } + /** + * Runs the request this stream was created for, sending its messages to the + * client and finalizing the stream with the response. + *

+ * The stream is detached from the session once the request is done with it, + * whatever the outcome: a response, an error, or the caller giving up. + * @param jsonrpcRequest the MCP request this stream carries the response of + * @return Mono which completes once the processing is done + */ + public Mono handle(McpSchema.JSONRPCRequest jsonrpcRequest) { + // The stream is released whichever way the request ends: with a response, on + // an error, or by the caller giving up on it + return Mono.usingWhen(Mono.just(this), stream -> Mono.deferContextual(ctx -> { + McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, + McpTransportContext.EMPTY); + + McpRequestHandler requestHandler = McpStreamableServerSession.this.requestHandlers + .get(jsonrpcRequest.method()); + if (requestHandler == null) { + MethodNotFoundError error = getMethodNotFoundError(jsonrpcRequest.method()); + return this.connection + .sendMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), null, + new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.METHOD_NOT_FOUND, + error.message(), error.data()))); + } + return requestHandler + .handle(new McpAsyncServerExchange(McpStreamableServerSession.this.id, this, + clientCapabilities.get(), clientInfo.get(), transportContext), jsonrpcRequest.params()) + .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), result, + null)) + .onErrorResume(e -> { + McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError; + if (e instanceof McpError mcpError && mcpError.getJsonRpcError() != null) { + jsonRpcError = mcpError.getJsonRpcError(); + } + else { + // The cause chain may contain internal details: log it, but + // only + // send the message to the client + logger.warn("Error handling request {} [{}]", jsonrpcRequest.id(), jsonrpcRequest.method(), + e); + jsonRpcError = new McpSchema.JSONRPCResponse.JSONRPCError( + McpSchema.ErrorCodes.INTERNAL_ERROR, + Objects.requireNonNullElse(e.getMessage(), "Internal error"), null); + } + + var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + jsonrpcRequest.id(), null, jsonRpcError); + return Mono.just(errorResponse); + }) + .flatMap(this.connection::sendMessage); + }), McpStreamableServerSessionStream::closeGracefully); + } + @Override public Mono sendNotification(String method, Object params) { McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification( McpSchema.JSONRPC_VERSION, method, params); String messageId = this.uuidGenerator.get(); // TODO: store message in history - return this.transport.sendMessage(jsonrpcNotification, messageId); + return this.connection.sendMessage(jsonrpcNotification, messageId); } @Override public Mono closeGracefully() { return Mono.defer(() -> { - this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); - this.pendingResponses.clear(); + McpStreamableServerSession.this.openStreams.remove(this); + McpStreamableServerSession.this.failPendingRequests(request -> this.equals(request.stream()), + "Stream closed"); // If this was the generic stream, reset it McpStreamableServerSession.this.listeningStreamRef.compareAndExchange(this, McpStreamableServerSession.this.missingMcpTransportSession); - McpStreamableServerSession.this.requestIdToStream.values().removeIf(this::equals); - return this.transport.closeGracefully(); + return this.connection.closeGracefully(); }); } @Override public void close() { - this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); - this.pendingResponses.clear(); + McpStreamableServerSession.this.openStreams.remove(this); + McpStreamableServerSession.this.failPendingRequests(request -> this.equals(request.stream()), + "Stream closed"); + // If this was the generic stream, reset it + McpStreamableServerSession.this.listeningStreamRef.compareAndExchange(this, + McpStreamableServerSession.this.missingMcpTransportSession); + this.connection.close(); + } + + /** + * Releases the connection carrying this stream, detaching the stream from the + * session, but keeps its pending server-initiated requests resolvable: the client + * answers those with a separate HTTP POST request, which outlives the SSE stream + * the request was sent on. + *

+ * This is the counterpart of {@link #close()} for the end of a connection rather + * than the end of the session: an SSE stream going away, whether replaced, + * disconnected or timed out, does not invalidate the requests sent on it. + */ + public void releaseTransport() { + McpStreamableServerSession.this.openStreams.remove(this); // If this was the generic stream, reset it McpStreamableServerSession.this.listeningStreamRef.compareAndExchange(this, McpStreamableServerSession.this.missingMcpTransportSession); - McpStreamableServerSession.this.requestIdToStream.values().removeIf(this::equals); - this.transport.close(); + this.connection.close(); } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java b/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java index 6d53ed516..0199262ba 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java @@ -1,11 +1,12 @@ /** - * Copyright 2025 - 2025 the original author or authors. + * Copyright 2025 - 2026 the original author or authors. */ 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; @@ -28,6 +29,7 @@ * The pings are sent to all active mcp sessions at regular intervals. * * @author Christian Tzolov + * @author Daniel Garnier-Moiroux */ public class KeepAliveScheduler { @@ -57,6 +59,9 @@ public class KeepAliveScheduler { /** Supplier for reactive McpSession instances */ private final Supplier> mcpSessions; + /** Invoked with the session whose keep-alive ping went unanswered */ + private final Consumer onPingFailure; + /** * Creates a KeepAliveScheduler with a custom scheduler, initial delay, interval and a * supplier for McpSession instances. @@ -64,13 +69,15 @@ public class KeepAliveScheduler { * @param initialDelay Initial delay before the first keepAlive call * @param interval Interval between subsequent keepAlive calls * @param mcpSessions Supplier for McpSession instances + * @param onPingFailure Callback invoked with the session whose ping went unanswered */ KeepAliveScheduler(Scheduler scheduler, Duration initialDelay, Duration interval, - Supplier> mcpSessions) { + Supplier> mcpSessions, Consumer onPingFailure) { this.scheduler = scheduler; this.initialDelay = initialDelay; this.interval = interval; this.mcpSessions = mcpSessions; + this.onPingFailure = onPingFailure; } /** @@ -92,8 +99,14 @@ public Disposable start() { .doOnNext(tick -> { this.mcpSessions.get() .flatMap(session -> session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF) - .doOnError(e -> logger.warn("Failed to send keep-alive ping to session {}: {}", session, - e.getMessage())) + // A ping has to be answered before the next one is due. The + // request timeout of the session is unrelated to keeping the + // connection alive, and is measured in hours by default. + .timeout(this.interval) + .doOnError(e -> { + logger.warn("Keep-alive ping to session {} failed: {}", session, e.getMessage()); + this.onPingFailure.accept(session); + }) .onErrorComplete()) .subscribe(); }) @@ -148,12 +161,15 @@ public static class Builder { private Scheduler scheduler = Schedulers.boundedElastic(); - private Duration initialDelay = Duration.ofSeconds(0); + private Duration initialDelay = Duration.ofSeconds(30); private Duration interval = Duration.ofSeconds(30); private Supplier> mcpSessions; + private Consumer onPingFailure = session -> { + }; + /** * Creates a new Builder instance with a supplier for McpSession instances. * @param mcpSessions The supplier for McpSession instances @@ -204,12 +220,28 @@ public Builder interval(Duration interval) { return this; } + /** + * Sets the callback invoked when a session does not answer a keep-alive ping + * within the keep-alive interval. An unanswered ping means the connection the + * ping was written to is dead, which the operating system does not necessarily + * report: writing to a connection whose peer is gone keeps succeeding until it + * resets. It does not mean the session itself is over, as the client is free to + * reconnect to it. + * @param onPingFailure The callback receiving the unresponsive session + * @return This builder instance for method chaining + */ + public Builder onPingFailure(Consumer onPingFailure) { + Assert.notNull(onPingFailure, "onPingFailure must not be null"); + this.onPingFailure = onPingFailure; + return this; + } + /** * Builds and returns a new KeepAliveScheduler instance. * @return A new KeepAliveScheduler configured with the builder's settings */ public KeepAliveScheduler build() { - return new KeepAliveScheduler(scheduler, initialDelay, interval, mcpSessions); + return new KeepAliveScheduler(scheduler, initialDelay, interval, mcpSessions, onPingFailure); } } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java new file mode 100644 index 000000000..ffc37f1d0 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java @@ -0,0 +1,264 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.spec; + +import java.time.Duration; +import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; + +import io.modelcontextprotocol.json.TypeRef; +import io.modelcontextprotocol.server.McpRequestHandler; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link McpStreamableServerSession}. + */ +class McpStreamableServerSessionTests { + + private static final Duration TIMEOUT = Duration.ofSeconds(5); + + private McpStreamableServerSession session() { + return session(Map.of()); + } + + private McpStreamableServerSession session(Map> requestHandlers) { + return new McpStreamableServerSession("session-1", McpSchema.ClientCapabilities.builder().build(), + new McpSchema.Implementation("test-client", "1.0.0"), TIMEOUT, requestHandlers, Map.of()); + } + + @Test + void replacedListeningStreamHasItsConnectionReleased() { + var session = session(); + var firstTransport = new RecordingTransport(); + + session.listeningStream(firstTransport); + assertThat(firstTransport.closed).isFalse(); + + session.listeningStream(new RecordingTransport()); + assertThat(firstTransport.closed).isTrue(); + } + + @Test + void replacingListeningStreamKeepsItsPendingRequestsResolvable() { + var session = session(); + var firstTransport = new RecordingTransport(); + session.listeningStream(firstTransport); + + // Server-initiated requests (sampling, elicitation, roots/list) are sent on the + // listening SSE stream, but the client answers them with a separate HTTP POST + // which outlives that stream. + var pending = session.sendRequest("sampling/createMessage", null, new TypeRef() { + }).toFuture(); + assertThat(firstTransport.sent).hasSize(1); + var requestId = ((McpSchema.JSONRPCRequest) firstTransport.sent.peek()).id(); + + // The client reconnects with a Last-Event-ID header, replacing the listening + // stream. The request sent on the replaced stream must stay pending. + session.listeningStream(new RecordingTransport()); + assertThat(pending).isNotDone(); + + session.accept(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, requestId, "response-value", null)) + .block(TIMEOUT); + + assertThat(pending).succeedsWithin(TIMEOUT).isEqualTo("response-value"); + } + + @Test + void closingTheSessionReleasesTheConnectionOfItsStreams() { + var session = session(); + var transport = new RecordingTransport(); + session.listeningStream(transport); + + session.close(); + + assertThat(transport.closed).isTrue(); + } + + @Test + void closingTheSessionFailsThePendingRequestsOfReleasedStreams() { + var session = session(); + session.listeningStream(new RecordingTransport()); + + var pending = session.sendRequest("sampling/createMessage", null, new TypeRef() { + }).toFuture(); + + // The stream the request was sent on is released, which leaves the request + // pending: a reconnecting client can still answer it over a separate POST + session.releaseListeningStream(); + assertThat(pending).isNotDone(); + + // The client cannot respond on a closed session, the pending request must be + // failed. + session.close(); + + assertThat(pending).failsWithin(TIMEOUT).withThrowableThat().havingCause().withMessage("Session closed"); + } + + @Test + void closingTheSessionGracefullyFailsThePendingRequestsOfReleasedStreams() { + var session = session(); + session.listeningStream(new RecordingTransport()); + + var pending = session.sendRequest("sampling/createMessage", null, new TypeRef() { + }).toFuture(); + + session.releaseListeningStream(); + assertThat(pending).isNotDone(); + + session.closeGracefully().block(TIMEOUT); + + assertThat(pending).failsWithin(TIMEOUT).withThrowableThat().havingCause().withMessage("Session closed"); + } + + @Test + void closingAStreamFailsOnlyItsOwnPendingRequests() { + var session = session(); + var listeningTransport = new RecordingTransport(); + session.listeningStream(listeningTransport); + + var onListeningStream = session.sendRequest("sampling/createMessage", null, new TypeRef() { + }).toFuture(); + + // A POST response stream carries its own server-initiated requests, which can be + // closed independently + var responseTransport = new RecordingTransport(); + var responseStream = session.new McpStreamableServerSessionStream(responseTransport); + var onResponseStream = responseStream.sendRequest("elicitation/create", null, new TypeRef() { + }).toFuture(); + + responseStream.close(); + + assertThat(onResponseStream).failsWithin(TIMEOUT) + .withThrowableThat() + .havingCause() + .withMessage("Stream closed"); + assertThat(onListeningStream).isNotDone(); + + var requestId = ((McpSchema.JSONRPCRequest) listeningTransport.sent.peek()).id(); + session.accept(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, requestId, "response-value", null)) + .block(TIMEOUT); + assertThat(onListeningStream).succeedsWithin(TIMEOUT).isEqualTo("response-value"); + } + + @Test + void endOfTheConnectionCarryingAResponseStreamDetachesItFromTheSession() { + // A request whose handler never completes, as seen when a client gives up and + // disconnects while the server is still working on its tool call + var session = session(Map.of("tools/call", (exchange, params) -> Mono.never())); + var transport = new RecordingTransport(); + + // The caller owns the stream, so it can detach it from the session once the + // container tells it the connection carrying it is gone + var stream = session.responseStream(transport); + stream.handle(new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "tools/call", "request-1", null)) + .subscribe(); + assertThat(session.hasOpenStream()).isTrue(); + + stream.releaseTransport(); + + // The session must stop believing it holds a live connection: hasOpenStream() is + // what tells the session sweeper that a client is still around, so a stream which + // outlives its connection makes the session impossible to reclaim + assertThat(session.hasOpenStream()).isFalse(); + assertThat(transport.closed).isTrue(); + } + + @Test + void responseStreamIsDetachedFromTheSessionOnceItsRequestIsAnswered() { + var session = session(Map.of("tools/call", (exchange, params) -> Mono.just("result"))); + var transport = new RecordingTransport(); + + var stream = session.responseStream(transport); + assertThat(session.hasOpenStream()).isTrue(); + + stream.handle(new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "tools/call", "request-1", null)) + .block(TIMEOUT); + + assertThat(session.hasOpenStream()).isFalse(); + } + + @Test + void responseStreamIsDetachedFromTheSessionWhenItsResponseCannotBeSent() { + var session = session(Map.of("tools/call", (exchange, params) -> Mono.just("result"))); + + var stream = session.responseStream(new FailingTransport()); + var handling = stream + .handle(new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "tools/call", "request-1", null)) + .toFuture(); + + // The connection which was to carry the response failed. The caller gets to see + // it, and the stream must not be left attached to the session. + assertThat(handling).failsWithin(TIMEOUT).withThrowableThat().havingCause().withMessage("connection gone"); + assertThat(session.hasOpenStream()).isFalse(); + } + + @Test + void abandonedResponseStreamIsDetachedFromTheSession() { + var session = session(Map.of("tools/call", (exchange, params) -> Mono.never())); + + var stream = session.responseStream(new RecordingTransport()); + var subscription = stream + .handle(new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "tools/call", "request-1", null)) + .subscribe(); + assertThat(session.hasOpenStream()).isTrue(); + + // The caller gives up on a request which would never terminate on its own, so + // nothing sends the response the stream was created to carry + subscription.dispose(); + + assertThat(session.hasOpenStream()).isFalse(); + } + + /** + * A transport whose connection is gone, so that nothing can be written to it. + */ + static class FailingTransport extends RecordingTransport { + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId) { + return Mono.error(new RuntimeException("connection gone")); + } + + } + + static class RecordingTransport implements McpStreamableServerTransport { + + final Queue sent = new ConcurrentLinkedQueue<>(); + + volatile boolean closed; + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId) { + return Mono.fromRunnable(() -> this.sent.add(message)); + } + + @Override + public Mono sendMessage(McpSchema.JSONRPCMessage message) { + return sendMessage(message, null); + } + + @Override + public Mono closeGracefully() { + return Mono.fromRunnable(() -> this.closed = true); + } + + @Override + public void close() { + this.closed = true; + } + + @SuppressWarnings("unchecked") + @Override + public T unmarshalFrom(Object data, TypeRef typeRef) { + return (T) data; + } + + } + +} diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java index 85e355968..b8b5431a8 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java @@ -36,6 +36,7 @@ import java.io.IOException; import java.time.Duration; +import java.time.Instant; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -44,6 +45,7 @@ * Implementation of a WebFlux based {@link McpStreamableServerTransportProvider}. * * @author Dariusz Jędrzejczyk + * @author Dimitar Proynov */ public class WebFluxStreamableServerTransportProvider implements McpStreamableServerTransportProvider { @@ -69,6 +71,25 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe private KeepAliveScheduler keepAliveScheduler; + /** + * Duration of inactivity after which an idle session is evicted. If {@code null}, + * idle sessions are never evicted. + */ + private final Duration sessionIdleTimeout; + + /** + * Tracks the last access time per session id, used to evict idle sessions. + */ + private final ConcurrentHashMap sessionLastAccessTimes = new ConcurrentHashMap<>(); + + private Disposable idleSessionScheduler; + + /** + * Max number of session this provider supports before discarding new initialization + * requests. + */ + private final long maxSessions; + /** * Security validator for validating HTTP requests. */ @@ -76,7 +97,8 @@ public class WebFluxStreamableServerTransportProvider implements McpStreamableSe private WebFluxStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, McpTransportContextExtractor contextExtractor, boolean disallowDelete, - Duration keepAliveInterval, ServerTransportSecurityValidator securityValidator) { + Duration keepAliveInterval, Duration sessionIdleTimeout, ServerTransportSecurityValidator securityValidator, + long maxSessions) { Assert.notNull(jsonMapper, "JsonMapper must not be null"); Assert.notNull(mcpEndpoint, "Message endpoint must not be null"); Assert.notNull(contextExtractor, "Context extractor must not be null"); @@ -87,6 +109,8 @@ private WebFluxStreamableServerTransportProvider(McpJsonMapper jsonMapper, Strin this.contextExtractor = contextExtractor; this.disallowDelete = disallowDelete; this.securityValidator = securityValidator; + this.maxSessions = maxSessions; + this.sessionIdleTimeout = sessionIdleTimeout; this.routerFunction = RouterFunctions.route() .GET(this.mcpEndpoint, this::handleGet) .POST(this.mcpEndpoint, this::handlePost) @@ -102,6 +126,11 @@ private WebFluxStreamableServerTransportProvider(McpJsonMapper jsonMapper, Strin this.keepAliveScheduler.start(); } + + if (sessionIdleTimeout != null) { + this.idleSessionScheduler = Flux.interval(sessionIdleTimeout, sessionIdleTimeout) + .subscribe(tick -> this.evictIdleSessions()); + } } @Override @@ -142,9 +171,13 @@ public Mono closeGracefully() { .then(); }).then().doOnSuccess(v -> { sessions.clear(); + this.sessionLastAccessTimes.clear(); if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.idleSessionScheduler != null) { + this.idleSessionScheduler.dispose(); + } }); } @@ -191,7 +224,7 @@ private Mono handleGet(ServerRequest request) { return ServerResponse.badRequest().build(); } - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().build(); // TODO: say we need a session // id } @@ -204,6 +237,8 @@ private Mono handleGet(ServerRequest request) { return ServerResponse.notFound().build(); } + touchSession(sessionId); + if (!request.headers().header(HttpHeaders.LAST_EVENT_ID).isEmpty()) { String lastId = request.headers().asHttpHeaders().getFirst(HttpHeaders.LAST_EVENT_ID); return ServerResponse.ok() @@ -266,31 +301,10 @@ private Mono handlePost(ServerRequest request) { McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(jsonMapper, body); if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest && jsonrpcRequest.method().equals(McpSchema.METHOD_INITIALIZE)) { - var typeReference = new TypeRef() { - }; - McpSchema.InitializeRequest initializeRequest = jsonMapper.convertValue(jsonrpcRequest.params(), - typeReference); - McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory - .startSession(initializeRequest); - sessions.put(init.session().getId(), init.session()); - return init.initResult().map(initializeResult -> { - McpSchema.JSONRPCResponse jsonrpcResponse = new McpSchema.JSONRPCResponse( - McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), initializeResult, null); - try { - return this.jsonMapper.writeValueAsString(jsonrpcResponse); - } - catch (IOException e) { - logger.warn("Failed to serialize initResponse", e); - throw Exceptions.propagate(e); - } - }) - .flatMap(initResult -> ServerResponse.ok() - .contentType(MediaType.APPLICATION_JSON) - .header(HttpHeaders.MCP_SESSION_ID, init.session().getId()) - .bodyValue(initResult)); + return handleInitRequest(jsonrpcRequest, hasMcpSessionIdHeader(request)); } - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().bodyValue(new McpError("Session ID missing")); } @@ -302,6 +316,8 @@ private Mono handlePost(ServerRequest request) { .bodyValue(new McpError("Session not found: " + sessionId)); } + touchSession(sessionId); + if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) { return session.accept(jsonrpcResponse).then(ServerResponse.accepted().build()); } @@ -353,7 +369,7 @@ private Mono handleDelete(ServerRequest request) { McpTransportContext transportContext = this.contextExtractor.extract(request); return Mono.defer(() -> { - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().build(); // TODO: say we need a session // id } @@ -370,10 +386,99 @@ private Mono handleDelete(ServerRequest request) { return ServerResponse.notFound().build(); } - return session.delete().then(ServerResponse.ok().build()); + return session.delete().doOnSuccess(unused -> removeSession(sessionId)).then(ServerResponse.ok().build()); }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)); } + private Mono handleInitRequest(McpSchema.JSONRPCRequest jsonrpcRequest, + boolean hasMcpSessionIdHeader) { + if (hasMcpSessionIdHeader) { + return ServerResponse.status(HttpStatus.BAD_REQUEST) + .bodyValue(McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) + .message("Session already initialized") + .build()); + } + if (this.sessions.size() >= this.maxSessions) { + return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE) + .bodyValue(McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) + .message("Max number of sessions reached") + .build()); + } + var typeReference = new TypeRef() { + }; + McpSchema.InitializeRequest initializeRequest = jsonMapper.convertValue(jsonrpcRequest.params(), typeReference); + McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory + .startSession(initializeRequest); + sessions.put(init.session().getId(), init.session()); + if (this.sessionIdleTimeout != null) { + this.sessionLastAccessTimes.put(init.session().getId(), Instant.now()); + } + return init.initResult().map(initializeResult -> { + McpSchema.JSONRPCResponse jsonrpcResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, + jsonrpcRequest.id(), initializeResult, null); + try { + return this.jsonMapper.writeValueAsString(jsonrpcResponse); + } + catch (IOException e) { + logger.warn("Failed to serialize initResponse", e); + throw Exceptions.propagate(e); + } + }) + .flatMap(initResult -> ServerResponse.ok() + .contentType(MediaType.APPLICATION_JSON) + .header(HttpHeaders.MCP_SESSION_ID, init.session().getId()) + .bodyValue(initResult)); + } + + /** + * Records the current time as the last access time for the given session, so that an + * active session is not evicted as idle. No-op when idle eviction is disabled or the + * session is no longer tracked. + * @param sessionId the id of the session that was just accessed + */ + private void touchSession(String sessionId) { + if (this.sessionIdleTimeout != null) { + this.sessionLastAccessTimes.computeIfPresent(sessionId, (id, lastAccess) -> Instant.now()); + } + } + + /** + * Evicts and closes sessions that have been idle for longer than the configured + * {@code sessionIdleTimeout}. This reclaims session slots held by clients that + * initialized a session but never interacted with it again, mitigating resource + * exhaustion from abandoned sessions. + */ + private void evictIdleSessions() { + if (this.sessionIdleTimeout == null || this.isClosing) { + return; + } + Instant now = Instant.now(); + for (var entry : this.sessionLastAccessTimes.entrySet()) { + if (Duration.between(entry.getValue(), now).compareTo(this.sessionIdleTimeout) > 0) { + String sessionId = entry.getKey(); + McpStreamableServerSession session = removeSession(sessionId); + logger.debug("Evicting idle session: {}", sessionId); + if (session != null) { + session.closeGracefully().onErrorComplete().subscribe(); + } + } + } + } + + /** + * Checks whether the given request carries an {@code Mcp-Session-Id} header. + * @param request the incoming server request + * @return {@code true} if the header is present, regardless of its value + */ + private static boolean hasMcpSessionIdHeader(ServerRequest request) { + return !request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty(); + } + + private McpStreamableServerSession removeSession(String sessionId) { + this.sessionLastAccessTimes.remove(sessionId); + return this.sessions.remove(sessionId); + } + private class WebFluxStreamableMcpSessionTransport implements McpStreamableServerTransport { private final FluxSink> sink; @@ -450,8 +555,12 @@ public static class Builder { private Duration keepAliveInterval; + private long maxSessions = 100_000L; + private ServerTransportSecurityValidator securityValidator = ServerTransportSecurityValidator.NOOP; + private Duration sessionIdleTimeout; + private Builder() { // used by a static method } @@ -531,6 +640,41 @@ public Builder securityValidator(ServerTransportSecurityValidator securityValida return this; } + /** + * Sets the maximum supported sessions for this provider. If not set defaults to + * 100_000. + * @param maxSessions - maximum allowed sessions for this provider + * @return this builder instance + * @throws IllegalArgumentException if maximum sessions is not positive number + */ + public Builder maxSessions(long maxSessions) { + Assert.isTrue(maxSessions >= 0, "Max sessions must be greater than 0"); + this.maxSessions = maxSessions; + return this; + } + + /** + * Sets the idle timeout after which sessions with no client activity are evicted + * and closed. Activity is any inbound request (GET, POST or DELETE) resolving to + * the session. This reclaims session slots held by clients that initialized a + * session but never interacted with it again, mitigating resource exhaustion from + * abandoned sessions. + *

+ * Keep-alive and idle eviction are complementary and may be enabled together: + * keep-alive pings preserve connected clients (a ping response counts as + * activity), while idle eviction reclaims abandoned sessions. When both are set, + * {@code sessionIdleTimeout} must be greater than the keep-alive interval so that + * a connected client answering pings is not evicted before it can respond; + * otherwise {@link #build()} throws. + * @param sessionIdleTimeout the maximum idle duration before eviction, or null to + * disable idle eviction (the default) + * @return this builder instance + */ + public Builder sessionIdleTimeout(Duration sessionIdleTimeout) { + this.sessionIdleTimeout = sessionIdleTimeout; + return this; + } + /** * Builds a new instance of {@link WebFluxStreamableServerTransportProvider} with * the configured settings. @@ -539,9 +683,14 @@ public Builder securityValidator(ServerTransportSecurityValidator securityValida */ public WebFluxStreamableServerTransportProvider build() { Assert.notNull(mcpEndpoint, "Message endpoint must be set"); + if (this.sessionIdleTimeout != null && this.keepAliveInterval != null) { + Assert.isTrue(this.sessionIdleTimeout.compareTo(this.keepAliveInterval) > 0, + "sessionIdleTimeout must be greater than keepAliveInterval, otherwise a connected client " + + "may be evicted before it can answer a keep-alive ping"); + } return new WebFluxStreamableServerTransportProvider( jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, mcpEndpoint, contextExtractor, - disallowDelete, keepAliveInterval, securityValidator); + disallowDelete, keepAliveInterval, sessionIdleTimeout, securityValidator, maxSessions); } } diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProviderTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProviderTests.java new file mode 100644 index 000000000..b26553f04 --- /dev/null +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProviderTests.java @@ -0,0 +1,142 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.server.transport; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; + +import io.modelcontextprotocol.spec.HttpHeaders; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpStreamableServerSession; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; + +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.test.web.reactive.server.WebTestClient; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link WebFluxStreamableServerTransportProvider}. + * + * @author Dimitar Proynov + */ +class WebFluxStreamableServerTransportProviderTests { + + private static final String INITIALIZE_REQUEST = """ + {"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2024-11-05",\ + "capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}"""; + + private static final McpSchema.InitializeResult INIT_RESULT = new McpSchema.InitializeResult("2024-11-05", + McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null); + + @Test + void buildRejectsSessionIdleTimeoutNotGreaterThanKeepAliveInterval() { + assertThatThrownBy(() -> WebFluxStreamableServerTransportProvider.builder() + .keepAliveInterval(Duration.ofSeconds(2)) + .sessionIdleTimeout(Duration.ofSeconds(1)) + .build()).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("sessionIdleTimeout must be greater than keepAliveInterval"); + } + + @Test + void idleSessionIsEvictedAfterTimeout() { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session"); + when(session.closeGracefully()).thenReturn(Mono.empty()); + + WebFluxStreamableServerTransportProvider provider = WebFluxStreamableServerTransportProvider.builder() + .sessionIdleTimeout(Duration.ofMillis(25)) + .build(); + provider.setSessionFactory(initializeRequest -> new McpStreamableServerSession.McpStreamableServerSessionInit( + session, Mono.just(INIT_RESULT))); + + WebTestClient client = WebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isOk(); + + // No request touches the session while it sits idle, so the scheduler should + // evict it and close it within a few eviction intervals. + // 20x compared to sessionIdleTimeout should be safe for CI + verify(session, timeout(500)).closeGracefully(); + } + + @Test + void initializeRequestIsRejectedWhenMaxSessionsReached() { + AtomicInteger sessionCounter = new AtomicInteger(); + WebFluxStreamableServerTransportProvider provider = WebFluxStreamableServerTransportProvider.builder() + .maxSessions(1) + .build(); + provider.setSessionFactory(initializeRequest -> { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session-" + sessionCounter.incrementAndGet()); + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, Mono.just(INIT_RESULT)); + }); + + WebTestClient client = WebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + // The first initialize request fills the single available session slot. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isOk(); + + // The second initialize request is rejected because maxSessions is reached. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isEqualTo(HttpStatus.SERVICE_UNAVAILABLE) + .expectBody(String.class) + .value(body -> assertThat(body).contains("Max number of sessions reached")); + } + + @Test + void initializeRequestIsRejectedWhenSessionIdHeaderIsPresent() { + WebFluxStreamableServerTransportProvider provider = WebFluxStreamableServerTransportProvider.builder().build(); + provider.setSessionFactory(initializeRequest -> { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session"); + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, Mono.just(INIT_RESULT)); + }); + + WebTestClient client = WebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + // An initialize request must not carry a session id yet; the server rejects it + // rather than silently starting a fresh session. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .header(HttpHeaders.MCP_SESSION_ID, "some-existing-session-id") + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isEqualTo(HttpStatus.BAD_REQUEST) + .expectBody(String.class) + .value(body -> assertThat(body).contains("Session already initialized")); + } + +} diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java index f282d54b6..d3e240416 100644 --- a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java @@ -6,6 +6,7 @@ import java.io.IOException; import java.time.Duration; +import java.time.Instant; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -36,6 +37,7 @@ import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; +import reactor.core.Disposable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -51,6 +53,7 @@ * * @author Christian Tzolov * @author Dariusz Jędrzejczyk + * @author Dimitar Proynov * @see McpStreamableServerTransportProvider * @see RouterFunction */ @@ -109,6 +112,25 @@ public class WebMvcStreamableServerTransportProvider implements McpStreamableSer */ private final ServerTransportSecurityValidator securityValidator; + /** + * Max number of session this provider supports before discarding new initialization + * requests. + */ + private final long maxSessions; + + /** + * Duration of inactivity after which an idle session is evicted. If {@code null}, + * idle sessions are never evicted. + */ + private final Duration sessionIdleTimeout; + + /** + * Tracks the last access time per session id, used to evict idle sessions. + */ + private final ConcurrentHashMap sessionLastAccessTimes = new ConcurrentHashMap<>(); + + private Disposable idleSessionScheduler; + /** * Constructs a new WebMvcStreamableServerTransportProvider instance. * @param jsonMapper The McpJsonMapper to use for JSON serialization/deserialization @@ -121,11 +143,15 @@ public class WebMvcStreamableServerTransportProvider implements McpStreamableSer * @param keepAliveInterval The interval for keep-alive pings. If null, no keep-alive * will be scheduled. * @param securityValidator The security validator for validating HTTP requests. + * @param maxSessions The maximum number of concurrent sessions supported. + * @param sessionIdleTimeout The idle timeout after which sessions with no client + * activity are evicted. If null, idle sessions are never evicted. * @throws IllegalArgumentException if any parameter is null */ private WebMvcStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, boolean disallowDelete, McpTransportContextExtractor contextExtractor, - Duration keepAliveInterval, ServerTransportSecurityValidator securityValidator) { + Duration keepAliveInterval, Duration sessionIdleTimeout, ServerTransportSecurityValidator securityValidator, + long maxSessions) { Assert.notNull(jsonMapper, "McpJsonMapper must not be null"); Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); Assert.notNull(contextExtractor, "McpTransportContextExtractor must not be null"); @@ -136,6 +162,8 @@ private WebMvcStreamableServerTransportProvider(McpJsonMapper jsonMapper, String this.disallowDelete = disallowDelete; this.contextExtractor = contextExtractor; this.securityValidator = securityValidator; + this.maxSessions = maxSessions; + this.sessionIdleTimeout = sessionIdleTimeout; this.routerFunction = RouterFunctions.route() .GET(this.mcpEndpoint, this::handleGet) .POST(this.mcpEndpoint, this::handlePost) @@ -151,6 +179,11 @@ private WebMvcStreamableServerTransportProvider(McpJsonMapper jsonMapper, String this.keepAliveScheduler.start(); } + + if (sessionIdleTimeout != null) { + this.idleSessionScheduler = Flux.interval(sessionIdleTimeout, sessionIdleTimeout) + .subscribe(tick -> this.evictIdleSessions()); + } } @Override @@ -213,11 +246,15 @@ public Mono closeGracefully() { }); this.sessions.clear(); + this.sessionLastAccessTimes.clear(); logger.debug("Graceful shutdown completed"); }).then().doOnSuccess(v -> { if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.idleSessionScheduler != null) { + this.idleSessionScheduler.dispose(); + } }); } @@ -260,7 +297,7 @@ private ServerResponse handleGet(ServerRequest request) { McpTransportContext transportContext = this.contextExtractor.extract(request); - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().body("Session ID required in mcp-session-id header"); } @@ -271,6 +308,8 @@ private ServerResponse handleGet(ServerRequest request) { return ServerResponse.notFound().build(); } + touchSession(sessionId); + logger.debug("Handling GET request for session: {}", sessionId); try { @@ -366,30 +405,11 @@ private ServerResponse handlePost(ServerRequest request) { // Handle initialization request if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest && jsonrpcRequest.method().equals(McpSchema.METHOD_INITIALIZE)) { - McpSchema.InitializeRequest initializeRequest = jsonMapper.convertValue(jsonrpcRequest.params(), - new TypeRef() { - }); - McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory - .startSession(initializeRequest); - this.sessions.put(init.session().getId(), init.session()); - - try { - McpSchema.InitializeResult initResult = init.initResult().block(); - - return ServerResponse.ok() - .contentType(MediaType.APPLICATION_JSON) - .header(HttpHeaders.MCP_SESSION_ID, init.session().getId()) - .body(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), initResult, - null)); - } - catch (Exception e) { - logger.error("Failed to initialize session: {}", e.getMessage()); - return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).body(new McpError(e.getMessage())); - } + return handleInitRequest(jsonrpcRequest, hasMcpSessionIdHeader(request)); } // Handle other messages that require a session - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().body(new McpError("Session ID missing")); } @@ -401,6 +421,8 @@ private ServerResponse handlePost(ServerRequest request) { .body(new McpError("Session not found: " + sessionId)); } + touchSession(sessionId); + if (message instanceof McpSchema.JSONRPCResponse jsonrpcResponse) { session.accept(jsonrpcResponse) .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) @@ -476,7 +498,7 @@ private ServerResponse handleDelete(ServerRequest request) { McpTransportContext transportContext = this.contextExtractor.extract(request); - if (request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty()) { + if (!hasMcpSessionIdHeader(request)) { return ServerResponse.badRequest().body("Session ID required in mcp-session-id header"); } @@ -489,7 +511,7 @@ private ServerResponse handleDelete(ServerRequest request) { try { session.delete().contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); - this.sessions.remove(sessionId); + removeSession(sessionId); return ServerResponse.ok().build(); } catch (Exception e) { @@ -498,6 +520,92 @@ private ServerResponse handleDelete(ServerRequest request) { } } + private ServerResponse handleInitRequest(McpSchema.JSONRPCRequest jsonrpcRequest, boolean hasMcpSessionIdHeader) { + if (hasMcpSessionIdHeader) { + return ServerResponse.status(HttpStatus.BAD_REQUEST) + .body(McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) + .message("Session already initialized") + .build()); + } + if (this.sessions.size() >= this.maxSessions) { + return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE) + .body(McpError.builder(McpSchema.ErrorCodes.INTERNAL_ERROR) + .message("Max number of sessions reached") + .build()); + } + McpSchema.InitializeRequest initializeRequest = jsonMapper.convertValue(jsonrpcRequest.params(), + new TypeRef() { + }); + McpStreamableServerSession.McpStreamableServerSessionInit init = this.sessionFactory + .startSession(initializeRequest); + this.sessions.put(init.session().getId(), init.session()); + if (this.sessionIdleTimeout != null) { + this.sessionLastAccessTimes.put(init.session().getId(), Instant.now()); + } + + try { + McpSchema.InitializeResult initResult = init.initResult().block(); + + return ServerResponse.ok() + .contentType(MediaType.APPLICATION_JSON) + .header(HttpHeaders.MCP_SESSION_ID, init.session().getId()) + .body(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, jsonrpcRequest.id(), initResult, null)); + } + catch (Exception e) { + logger.error("Failed to initialize session: {}", e.getMessage()); + return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).body(new McpError(e.getMessage())); + } + } + + /** + * Records the current time as the last access time for the given session, so that an + * active session is not evicted as idle. No-op when idle eviction is disabled or the + * session is no longer tracked. + * @param sessionId the id of the session that was just accessed + */ + private void touchSession(String sessionId) { + if (this.sessionIdleTimeout != null) { + this.sessionLastAccessTimes.computeIfPresent(sessionId, (id, lastAccess) -> Instant.now()); + } + } + + /** + * Evicts and closes sessions that have been idle for longer than the configured + * {@code sessionIdleTimeout}. This reclaims session slots held by clients that + * initialized a session but never interacted with it again, mitigating resource + * exhaustion from abandoned sessions. + */ + private void evictIdleSessions() { + if (this.sessionIdleTimeout == null || this.isClosing) { + return; + } + Instant now = Instant.now(); + for (var entry : this.sessionLastAccessTimes.entrySet()) { + if (Duration.between(entry.getValue(), now).compareTo(this.sessionIdleTimeout) > 0) { + String sessionId = entry.getKey(); + McpStreamableServerSession session = removeSession(sessionId); + logger.debug("Evicting idle session: {}", sessionId); + if (session != null) { + session.closeGracefully().onErrorComplete().subscribe(); + } + } + } + } + + /** + * Checks whether the given request carries an {@code Mcp-Session-Id} header. + * @param request the incoming server request + * @return {@code true} if the header is present, regardless of its value + */ + private static boolean hasMcpSessionIdHeader(ServerRequest request) { + return !request.headers().header(HttpHeaders.MCP_SESSION_ID).isEmpty(); + } + + private McpStreamableServerSession removeSession(String sessionId) { + this.sessionLastAccessTimes.remove(sessionId); + return this.sessions.remove(sessionId); + } + /** * Implementation of McpStreamableServerTransport for WebMVC SSE sessions. This class * handles the transport-level communication for a specific client session. @@ -654,6 +762,10 @@ public static class Builder { private ServerTransportSecurityValidator securityValidator = ServerTransportSecurityValidator.NOOP; + private long maxSessions = 100_000L; + + private Duration sessionIdleTimeout; + /** * Sets the McpJsonMapper to use for JSON serialization/deserialization of MCP * messages. @@ -729,6 +841,41 @@ public Builder securityValidator(ServerTransportSecurityValidator securityValida return this; } + /** + * Sets the maximum supported sessions for this provider. If not set defaults to + * 100_000. + * @param maxSessions - maximum allowed sessions for this provider + * @return this builder instance + * @throws IllegalArgumentException if maximum sessions is not positive number + */ + public Builder maxSessions(long maxSessions) { + Assert.isTrue(maxSessions >= 0, "Max sessions must be greater than 0"); + this.maxSessions = maxSessions; + return this; + } + + /** + * Sets the idle timeout after which sessions with no client activity are evicted + * and closed. Activity is any inbound request (GET, POST or DELETE) resolving to + * the session. This reclaims session slots held by clients that initialized a + * session but never interacted with it again, mitigating resource exhaustion from + * abandoned sessions. + *

+ * Keep-alive and idle eviction are complementary and may be enabled together: + * keep-alive pings preserve connected clients (a ping response counts as + * activity), while idle eviction reclaims abandoned sessions. When both are set, + * {@code sessionIdleTimeout} must be greater than the keep-alive interval so that + * a connected client answering pings is not evicted before it can respond; + * otherwise {@link #build()} throws. + * @param sessionIdleTimeout the maximum idle duration before eviction, or null to + * disable idle eviction (the default) + * @return this builder instance + */ + public Builder sessionIdleTimeout(Duration sessionIdleTimeout) { + this.sessionIdleTimeout = sessionIdleTimeout; + return this; + } + /** * Builds a new instance of {@link WebMvcStreamableServerTransportProvider} with * the configured settings. @@ -737,9 +884,14 @@ public Builder securityValidator(ServerTransportSecurityValidator securityValida */ public WebMvcStreamableServerTransportProvider build() { Assert.notNull(this.mcpEndpoint, "MCP endpoint must be set"); + if (this.sessionIdleTimeout != null && this.keepAliveInterval != null) { + Assert.isTrue(this.sessionIdleTimeout.compareTo(this.keepAliveInterval) > 0, + "sessionIdleTimeout must be greater than keepAliveInterval, otherwise a connected client " + + "may be evicted before it can answer a keep-alive ping"); + } return new WebMvcStreamableServerTransportProvider( jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, mcpEndpoint, disallowDelete, - contextExtractor, keepAliveInterval, securityValidator); + contextExtractor, keepAliveInterval, sessionIdleTimeout, securityValidator, maxSessions); } } diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProviderTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProviderTests.java new file mode 100644 index 000000000..53422c51a --- /dev/null +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProviderTests.java @@ -0,0 +1,143 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.server.transport; + +import java.time.Duration; +import java.util.concurrent.atomic.AtomicInteger; + +import io.modelcontextprotocol.spec.HttpHeaders; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpStreamableServerSession; +import org.junit.jupiter.api.Test; +import reactor.core.publisher.Mono; + +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.test.web.reactive.server.WebTestClient; +import org.springframework.test.web.servlet.client.MockMvcWebTestClient; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link WebMvcStreamableServerTransportProvider}. + * + * @author Dimitar Proynov + */ +class WebMvcStreamableServerTransportProviderTests { + + private static final String INITIALIZE_REQUEST = """ + {"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2024-11-05",\ + "capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}"""; + + private static final McpSchema.InitializeResult INIT_RESULT = new McpSchema.InitializeResult("2024-11-05", + McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"), null); + + @Test + void buildRejectsSessionIdleTimeoutNotGreaterThanKeepAliveInterval() { + assertThatThrownBy(() -> WebMvcStreamableServerTransportProvider.builder() + .keepAliveInterval(Duration.ofSeconds(2)) + .sessionIdleTimeout(Duration.ofSeconds(1)) + .build()).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("sessionIdleTimeout must be greater than keepAliveInterval"); + } + + @Test + void idleSessionIsEvictedAfterTimeout() { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session"); + when(session.closeGracefully()).thenReturn(Mono.empty()); + + WebMvcStreamableServerTransportProvider provider = WebMvcStreamableServerTransportProvider.builder() + .sessionIdleTimeout(Duration.ofMillis(25)) + .build(); + provider.setSessionFactory(initializeRequest -> new McpStreamableServerSession.McpStreamableServerSessionInit( + session, Mono.just(INIT_RESULT))); + + WebTestClient client = MockMvcWebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isOk(); + + // No request touches the session while it sits idle, so the scheduler should + // evict it and close it within a few eviction intervals. + // 20x compared to sessionIdleTimeout should be safe for CI + verify(session, timeout(500)).closeGracefully(); + } + + @Test + void initializeRequestIsRejectedWhenMaxSessionsReached() { + AtomicInteger sessionCounter = new AtomicInteger(); + WebMvcStreamableServerTransportProvider provider = WebMvcStreamableServerTransportProvider.builder() + .maxSessions(1) + .build(); + provider.setSessionFactory(initializeRequest -> { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session-" + sessionCounter.incrementAndGet()); + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, Mono.just(INIT_RESULT)); + }); + + WebTestClient client = MockMvcWebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + // The first initialize request fills the single available session slot. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isOk(); + + // The second initialize request is rejected because maxSessions is reached. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isEqualTo(HttpStatus.SERVICE_UNAVAILABLE) + .expectBody(String.class) + .value(body -> assertThat(body).contains("Max number of sessions reached")); + } + + @Test + void initializeRequestIsRejectedWhenSessionIdHeaderIsPresent() { + WebMvcStreamableServerTransportProvider provider = WebMvcStreamableServerTransportProvider.builder().build(); + provider.setSessionFactory(initializeRequest -> { + McpStreamableServerSession session = mock(McpStreamableServerSession.class); + when(session.getId()).thenReturn("test-session"); + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, Mono.just(INIT_RESULT)); + }); + + WebTestClient client = MockMvcWebTestClient.bindToRouterFunction(provider.getRouterFunction()).build(); + + // An initialize request must not carry a session id yet; the server rejects it + // rather than silently starting a fresh session. + client.post() + .uri("/mcp") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM) + .header(HttpHeaders.MCP_SESSION_ID, "some-existing-session-id") + .bodyValue(INITIALIZE_REQUEST) + .exchange() + .expectStatus() + .isEqualTo(HttpStatus.BAD_REQUEST) + .expectBody(String.class) + .value(body -> assertThat(body).contains("Session already initialized")); + } + +} diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java index 6cb532ffc..8d46dd41e 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -1,21 +1,31 @@ /* - * Copyright 2024 - 2024 the original author or authors. + * Copyright 2024 - 2026 the original author or authors. */ 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; +import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.read.ListAppender; import io.modelcontextprotocol.AbstractMcpClientServerIntegrationTests; import io.modelcontextprotocol.client.McpClient; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; @@ -26,12 +36,15 @@ import io.modelcontextprotocol.server.transport.TomcatTestUtil; import io.modelcontextprotocol.spec.HttpHeaders; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.util.KeepAliveScheduler; import jakarta.servlet.http.HttpServletRequest; import jakarta.servlet.http.HttpServletResponse; import org.apache.catalina.LifecycleException; import org.apache.catalina.LifecycleState; import org.apache.catalina.startup.Tomcat; +import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; @@ -39,12 +52,24 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; +import reactor.test.StepVerifier; +import reactor.test.scheduler.VirtualTimeScheduler; import static io.modelcontextprotocol.util.ToolsUtils.EMPTY_JSON_SCHEMA; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.awaitility.Awaitility.await; +/** + * Based on {@link AbstractMcpClientServerIntegrationTests} for basic client <> server + * integration tests. Also contains some tests specific to streamable HTTP, around + * resumability, connection lifecycle and session management. + */ @Timeout(15) class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests { @@ -52,32 +77,72 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg private static final String MESSAGE_ENDPOINT = "/mcp/message"; + // Tomcat is started once for the whole class; each test swaps in its own transport + private static final TomcatTestUtil.DelegatingServlet MCP_SERVLET = new TomcatTestUtil.DelegatingServlet(); + + private static Tomcat tomcat; + private HttpServletStreamableServerTransportProvider mcpServerTransportProvider; - private Tomcat tomcat; + // Keep alive is fast. A ping failure releases the stream, so listening + // steams are released quickly. + private final Duration KEEP_ALIVE_INTERVAL = Duration.ofMillis(150); + + // Sweeping is slower than keep-alive, so that a failed ping doesn't immediately + // result in a session sweep + private final Duration SESSION_SWEEP_INTERVAL = KEEP_ALIVE_INTERVAL.multipliedBy(2); + + private final HttpClient httpClient = HttpClient.newHttpClient(); + + private void awaitClientStreamEstablished() { + var timeout = Duration.ofSeconds(5); + await().atMost(timeout).untilAsserted(() -> { + assertThat(MCP_SERVLET.isStreamEstablished()) + .withFailMessage("[Failed to observe MCP Client connection within %s]", timeout) + .isTrue(); + }); + } static Stream clientsForTesting() { return Stream.of(Arguments.of("httpclient")); } + @BeforeAll + public static void beforeAll() { + tomcat = TomcatTestUtil.createTomcatServer("", PORT, MCP_SERVLET); + try { + tomcat.start(); + assertThat(tomcat.getServer().getState()).isEqualTo(LifecycleState.STARTED); + } + catch (Exception e) { + throw new RuntimeException("Failed to start Tomcat", e); + } + } + + @AfterAll + public static void afterAll() { + if (tomcat != null) { + try { + tomcat.stop(); + tomcat.destroy(); + } + catch (LifecycleException e) { + throw new RuntimeException("Failed to stop Tomcat", e); + } + } + } + @BeforeEach public void before() { // Create and configure the transport provider mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) - .keepAliveInterval(Duration.ofSeconds(1)) + .keepAliveInterval(KEEP_ALIVE_INTERVAL) .maxRequestSize(MAX_REQUEST_SIZE) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) .build(); - - tomcat = TomcatTestUtil.createTomcatServer("", PORT, mcpServerTransportProvider); - try { - tomcat.start(); - assertThat(tomcat.getServer().getState()).isEqualTo(LifecycleState.STARTED); - } - catch (Exception e) { - throw new RuntimeException("Failed to start Tomcat", e); - } + MCP_SERVLET.setDelegate(mcpServerTransportProvider); clientBuilders .put("httpclient", @@ -101,14 +166,46 @@ public void after() { if (mcpServerTransportProvider != null) { mcpServerTransportProvider.closeGracefully().block(); } - if (tomcat != null) { - try { - tomcat.stop(); - tomcat.destroy(); - } - catch (LifecycleException e) { - throw new RuntimeException("Failed to stop Tomcat", e); - } + } + + @Test + void testMissingHandlerReturnsMethodNotFoundError() { + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) + .build(); + var clientTransport = HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) + .endpoint(MESSAGE_ENDPOINT) + .build(); + + try (var mcpClient = McpClient.sync(clientTransport).build()) { + // Create a session using an MCP client + McpSchema.InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Override the response handler in the client to capture responses + AtomicReference response = new AtomicReference<>(); + var handler = (Function, Mono>) ( + message) -> message.doOnNext(r -> { + if (r instanceof McpSchema.JSONRPCResponse resp) { + response.set(resp); + } + }); + StepVerifier.create(clientTransport.connect(handler)).verifyComplete(); + + // Send an incorrect request through the transport + StepVerifier + .create(clientTransport.sendMessage( + new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "foo/bar", "test-request-123", null))) + .verifyComplete(); + + // Wait until we've received the response + await().atMost(Duration.ofSeconds(1)).until(() -> response.get() != null); + + assertThat(response.get().error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND); + assertThat(response.get().error().message()).isEqualTo("Method not found: foo/bar"); + } + finally { + mcpServer.close(); } } @@ -158,12 +255,11 @@ void testRejectsWhenContentLengthHeaderExceedsLimit(String clientType) throws Ex } @Test - public void rejectsWhenBodyBytesExceedLimitWithoutContentLengthHeader() throws Exception { - var httpClient = HttpClient.newHttpClient(); + void rejectsWhenBodyBytesExceedLimitWithoutContentLengthHeader() throws Exception { // 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() { @@ -200,7 +296,6 @@ public void cancel() { @ParameterizedTest @ValueSource(strings = { "text/plain;charset=UTF-8", "application/x-www-form-urlencoded", "multipart/form-data" }) void rejectsInitializeWithNonJsonContentType(String contentType) throws Exception { - var httpClient = HttpClient.newHttpClient(); prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); // CORS-safelisted content types can be sent cross-origin by a browser without a @@ -224,7 +319,6 @@ void rejectsInitializeWithNonJsonContentType(String contentType) throws Exceptio @Test void rejectsToolCallWithNonJsonContentType() throws Exception { - var httpClient = HttpClient.newHttpClient(); var toolCalled = new AtomicBoolean(); prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") .capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) @@ -254,6 +348,316 @@ void rejectsToolCallWithNonJsonContentType() throws Exception { assertThat(toolCalled).isFalse(); } + @Test + void resumedStreamReceivesServerNotifications() { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + + var sessionId = initializeSession(httpClient); + + // Resume the stream the way a client does once its SSE connection broke. The + // resumed stream must become the session listening stream, otherwise the + // reconnected client never receives anything again. + var stream = openListeningStream(httpClient, sessionId, sessionId + "_0"); + + awaitClientStreamEstablished(); + awaitNotification(stream.events()); + } + + @Test + void replacedListeningStreamIsClosed() { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + + var sessionId = initializeSession(httpClient); + + var firstStream = openListeningStream(httpClient, sessionId, null); + 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. 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)); + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block(); + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { + assertThat(secondStream.events()).anyMatch(line -> line.contains("notifications/resources/list_changed")); + assertThat(firstStream.events()).noneMatch(line -> line.contains("notifications/resources/list_changed")); + }); + } + + @Test + void keepAliveSkipsSessionsWithoutListeningStream() { + var keepAliveLogger = (ch.qos.logback.classic.Logger) LoggerFactory.getLogger(KeepAliveScheduler.class); + ListAppender logAppender = new ListAppender<>(); + logAppender.start(); + keepAliveLogger.addAppender(logAppender); + + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(KEEP_ALIVE_INTERVAL, null); + initializeSession(httpClient); + + // go through multiple keep-alive-rounds + virtualTime.advanceTimeBy(KEEP_ALIVE_INTERVAL.multipliedBy(3)); + + assertThat(logAppender.list).as("Should not display any Keep-Alive warning log") + .noneMatch(event -> event.getLevel() == Level.WARN); + } + finally { + keepAliveLogger.detachAppender(logAppender); + logAppender.stop(); + VirtualTimeScheduler.reset(); + } + + } + + @Test + void keepAlivePingsSessionsWithListeningStream() { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + + var sessionId = initializeSession(httpClient); + var stream = openListeningStream(httpClient, sessionId, null); + 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() { + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + + 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() { + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(null, SESSION_SWEEP_INTERVAL); + var sessionId = initializeSession(httpClient); + + // One client request every sweep interval/2, the session remains active + // We do 5 total requests so the total time is > 2 sweep intervals, verifying + // that the sweeper does not reclaim a session + for (int i = 1; i < 6; i++) { + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.dividedBy(2)); + assertThat(postNotification(httpClient, sessionId)).as("Session should be active for request #%s", i) + .isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); + assertThat(postNotification(httpClient, sessionId)).as("Session should be sweeped after SWEEP_INTERVAL * 2") + .isEqualTo(HttpServletResponse.SC_NOT_FOUND); + } + finally { + VirtualTimeScheduler.reset(); + } + } + + @Test + void sessionHoldingAnOpenStreamIsNotEvicted() { + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(null, SESSION_SWEEP_INTERVAL); + var sessionId = initializeSession(httpClient); + openListeningStream(httpClient, sessionId, null); + awaitClientStreamEstablished(); + + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); + assertThat(postNotification(httpClient, sessionId)) + .as("Session should not be sweeped when a GET SSE stream is open") + .isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + finally { + VirtualTimeScheduler.reset(); + } + } + + @Test + void sessionEvictionAfterReleasingStream() { + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(null, SESSION_SWEEP_INTERVAL); + 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()); + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); + assertThat(postNotification(httpClient, sessionId)).as("Session should be active when a stream is open") + .isEqualTo(HttpServletResponse.SC_ACCEPTED); + + // Close the client stream + clientStream.closeStream(); + + await().atMost(Duration.ofSeconds(5)) + .pollDelay(Duration.ZERO) + .pollInterval(Duration.ZERO) + .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 + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); + assertThat(postNotification(httpClient, sessionId)) + .as("Session should be sweeped after a stream is released") + .isEqualTo(HttpServletResponse.SC_NOT_FOUND); + }); + } + finally { + VirtualTimeScheduler.reset(); + } + } + + @Test + void sessionIsNotEvictedWithoutSweepInterval() { + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(KEEP_ALIVE_INTERVAL, null); + var sessionId = initializeSession(httpClient); + + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(10)); + assertThat(postNotification(httpClient, sessionId)) + .as("Session should never be sweeped when there is no sessionSweepInterval") + .isEqualTo(HttpServletResponse.SC_ACCEPTED); + } + finally { + VirtualTimeScheduler.reset(); + } + } + + /** + * A client which aborts a tool call must not leave its session behind. + */ + @Test + void sessionIsEvictedWhenTheClientAbortsAResponseStream() { + var virtualTime = VirtualTimeScheduler.getOrSet(); + try { + withTransportProvider(null, SESSION_SWEEP_INTERVAL); + var toolEmissionInterval = Duration.ofMillis(10); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(McpSchema.ServerCapabilities.builder().tools(true).build()) + .tools(McpServerFeatures.AsyncToolSpecification.builder() + .tool(McpSchema.Tool.builder() + .name("hangs") + .inputSchema(EMPTY_JSON_SCHEMA) + .description("never returns") + .build()) + .callHandler((exchange, request) -> Flux + .interval(Duration.ZERO, toolEmissionInterval, Schedulers.newSingle("test-tool")) + .flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.INFO) + .data("x".repeat(64 * 1024)) + .build())) + .then(Mono.never())) + .build()) + .build(); + var sessionId = initializeSession(httpClient); + + // The POST opens a response SSE stream, which the session counts as an open + // stream for as long as the call is in flight + var responseStream = postToolCall(httpClient, sessionId, "hangs"); + + await().atMost(Duration.ofSeconds(5)) + .pollDelay(Duration.ZERO) + .pollInterval(Duration.ofMillis(10)) + .untilAsserted(() -> { + // We do this in a loop for the wire-side of the connection - data has + // to be buffered then flushed + + // Advance time so the tool emits some data + virtualTime.advanceTimeBy(toolEmissionInterval); + // Verify the data is received + assertThat(responseStream.events()).anyMatch(line -> line.contains("notifications/message")); + }); + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(2)); + assertThat(postNotification(httpClient, sessionId)).as("Session with POST stream is active") + .isEqualTo(HttpServletResponse.SC_ACCEPTED); + + // The client gives up on the call and disconnects + responseStream.closeStream(); + + // Nothing is connected to the session anymore, so the sweeper must reclaim + // it. + // Probing only once: any request would count as activity and reset the clock. + virtualTime.advanceTimeBy(SESSION_SWEEP_INTERVAL.multipliedBy(5)); + + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_NOT_FOUND); + } + finally { + VirtualTimeScheduler.reset(); + } + } + + /** + * Calls a tool with a POST request, returning a handle on the SSE response stream it + * opens. + */ + private StreamResponse postToolCall(HttpClient httpClient, String sessionId, String toolName) { + var post = 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\",\"id\":\"call-1\"," + + "\"method\":\"tools/call\",\"params\":{\"name\":\"" + toolName + "\",\"arguments\":{}}}")) + .build(); + Queue events = new ConcurrentLinkedQueue<>(); + var streamRef = new AtomicReference(); + var clientFuture = httpClient.sendAsync(post, HttpResponse.BodyHandlers.ofInputStream()) + .thenAccept(response -> { + 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 closeStream(), not a failure + } + }); + return new StreamResponse(clientFuture, events, streamRef); + } + + private int postNotification(HttpClient httpClient, String sessionId) { + 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(); + try { + return httpClient.send(notification, HttpResponse.BodyHandlers.ofString()).statusCode(); + } + catch (IOException | InterruptedException e) { + return -1; + } + } + private String initializeSession(HttpClient httpClient) { var initialize = HttpRequest.newBuilder() .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) @@ -276,4 +680,72 @@ private String initializeSession(HttpClient httpClient) { return response.headers().firstValue(HttpHeaders.MCP_SESSION_ID).orElse(null); } + /** + * Opens an SSE listening stream with a GET request, collecting the received lines. + * @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() + .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) + .header("Accept", "text/event-stream") + .header(HttpHeaders.MCP_SESSION_ID, sessionId); + if (lastEventId != null) { + get.header(HttpHeaders.LAST_EVENT_ID, lastEventId); + } + Queue events = new ConcurrentLinkedQueue<>(); + var streamRef = new AtomicReference(); + var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofInputStream()) + .thenAccept(response -> { + 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, events, streamRef); + } + + private void awaitNotification(Queue events) { + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block(); + await().atMost(Duration.ofSeconds(5)).pollDelay(Duration.ofMillis(100)).untilAsserted(() -> { + assertThat(events).anyMatch(line -> line.contains("notifications/resources/list_changed")); + }); + } + + record StreamResponse(CompletableFuture streamFuture, Queue events, + AtomicReference streamRef) { + + void closeStream() { + // Close listening stream. We retry a few times in case the stream was not + // established on the first try + await().pollDelay(Duration.ZERO).atMost(Duration.ofSeconds(1)).until(() -> { + var stream = streamRef.get(); + if (stream != null) { + stream.close(); + return true; + } + return false; + + }); + } + } + + private void withTransportProvider(Duration keepAliveInterval, Duration sessionSweepInterval) { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + .keepAliveInterval(keepAliveInterval) + // no sweeping: the session must survive every keep-alive round + .sessionSweepInterval(sessionSweepInterval) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); + } + } diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java index 490e29838..84747017d 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java @@ -7,10 +7,17 @@ import java.io.IOException; import java.net.InetSocketAddress; import java.net.ServerSocket; +import java.util.concurrent.atomic.AtomicBoolean; import jakarta.servlet.Filter; import jakarta.servlet.Servlet; +import jakarta.servlet.ServletConfig; +import jakarta.servlet.ServletException; +import jakarta.servlet.ServletRequest; +import jakarta.servlet.ServletResponse; +import jakarta.servlet.http.HttpServletRequest; import org.apache.catalina.Context; +import org.apache.catalina.Wrapper; import org.apache.catalina.startup.Tomcat; import org.apache.tomcat.util.descriptor.web.FilterDef; import org.apache.tomcat.util.descriptor.web.FilterMap; @@ -37,7 +44,7 @@ public static Tomcat createTomcatServer(String contextPath, int port, Servlet se Context context = tomcat.addContext(contextPath, baseDir); // Add transport servlet to Tomcat - org.apache.catalina.Wrapper wrapper = context.createWrapper(); + Wrapper wrapper = context.createWrapper(); wrapper.setName("mcpServlet"); wrapper.setServlet(servlet); wrapper.setLoadOnStartup(1); @@ -63,6 +70,71 @@ public static Tomcat createTomcatServer(String contextPath, int port, Servlet se return tomcat; } + /** + * A servlet forwarding all requests to a delegate that can be swapped between tests. + * Register one with {@link #createTomcatServer} to start Tomcat once per test class + * while still handing each test a freshly built transport. + */ + public static class DelegatingServlet implements Servlet { + + private volatile ServletConfig servletConfig; + + private volatile Servlet delegate; + + // Crude way of tracking whether a GET SSE stream has been established, to ensure + // a Streamable HTTP MCP Client is connected. + // This is an approximation: it switches to "true" whenever the handler is done + // servicing a GET request - even if the request is rejected or if it's a replay. + private final AtomicBoolean sseStreamEstablished = new AtomicBoolean(false); + + /** + * Sets the servlet handling subsequent requests. The delegate is not + * {@link Servlet#init(ServletConfig) initialized}, since the MCP servlet + * transports do not rely on their {@link ServletConfig}. + */ + public void setDelegate(Servlet delegate) { + this.sseStreamEstablished.set(false); + this.delegate = delegate; + } + + @Override + public void init(ServletConfig config) { + this.servletConfig = config; + } + + @Override + public ServletConfig getServletConfig() { + return this.servletConfig; + } + + @Override + public void service(ServletRequest request, ServletResponse response) throws ServletException, IOException { + var current = this.delegate; + if (current == null) { + throw new IllegalStateException("No delegate servlet has been set"); + } + current.service(request, response); + if (request instanceof HttpServletRequest req && req.getMethod().equals("GET")) { + sseStreamEstablished.set(true); + } + } + + @Override + public String getServletInfo() { + return DelegatingServlet.class.getSimpleName(); + } + + @Override + public void destroy() { + this.delegate = null; + } + + public boolean isStreamEstablished() { + return this.sseStreamEstablished.get(); + } + + } + /** * Finds an available port on the local machine. * @return an available port number