diff --git a/docs/server.md b/docs/server.md index 93fcf68bc..5948f4ade 100644 --- a/docs/server.md +++ b/docs/server.md @@ -164,7 +164,10 @@ Key features: - 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 cacb30522..5f883cb98 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 @@ -9,7 +9,9 @@ import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.ReentrantLock; import io.modelcontextprotocol.common.McpTransportContext; @@ -20,13 +22,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; @@ -34,8 +40,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 @@ -52,6 +60,7 @@ * @author Zachary German * @author Christian Tzolov * @author Dariusz Jędrzejczyk + * @author Daniel Garnier-Moiroux * @see McpStreamableServerTransportProvider * @see HttpServlet */ @@ -114,6 +123,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; /** @@ -127,6 +144,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. */ @@ -145,11 +170,14 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet * @param httpHeaderValidator The HTTP header validator for validating HTTP requests. * @param requestMaxSize The maximum size, in bytes, of a single request body. Must be * positive. + * @param sessionSweepInterval The interval at which idle sessions are evicted. If + * null, no sweeping will be scheduled. * @throws IllegalArgumentException if any parameter is null */ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint, boolean disallowDelete, McpTransportContextExtractor contextExtractor, - Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize) { + Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize, + Duration sessionSweepInterval) { Assert.notNull(jsonMapper, "JsonMapper must not be null"); Assert.notNull(mcpEndpoint, "MCP endpoint must not be null"); Assert.notNull(contextExtractor, "Context extractor must not be null"); @@ -165,15 +193,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 @@ -248,6 +357,9 @@ public Mono closeGracefully() { if (this.keepAliveScheduler != null) { this.keepAliveScheduler.shutdown(); } + if (this.sessionSweeper != null) { + this.sessionSweeper.dispose(); + } }); } @@ -307,6 +419,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); @@ -338,30 +451,7 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response) McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session .listeningStream(sessionTransport); - asyncContext.addListener(new jakarta.servlet.AsyncListener() { - @Override - public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException { - logger.debug("SSE connection completed for session: {}", sessionId); - listeningStream.close(); - } - - @Override - public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException { - logger.debug("SSE connection timed out for session: {}", sessionId); - listeningStream.close(); - } - - @Override - public void onError(jakarta.servlet.AsyncEvent event) throws IOException { - logger.debug("SSE connection error for session: {}", sessionId); - listeningStream.close(); - } - - @Override - public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException { - // No action needed - } - }); + registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport); } catch (Exception e) { logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage()); @@ -369,34 +459,6 @@ public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException { } } - /** - * Replays the messages the client missed while its SSE stream was broken. - * @param session the session the client is resuming - * @param lastEventId the ID of the last event received by the client - * @param sessionTransport the transport of the resumed SSE stream - * @param transportContext the context extracted from the request - * @return {@code true} if the replay completed, {@code false} if it failed, in which - * case the transport has been closed - */ - private boolean tryReplayMissedMessages(McpStreamableServerSession session, String lastEventId, - McpStreamableServerTransport sessionTransport, McpTransportContext transportContext) { - try { - for (McpSchema.JSONRPCMessage message : session.replay(lastEventId) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .toIterable()) { - sessionTransport.sendMessage(message) - .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)) - .block(); - } - return true; - } - catch (Exception e) { - logger.error("Failed to replay messages for session {}: {}", session.getId(), e.getMessage()); - sessionTransport.close(); - return false; - } - } - /** * Handles POST requests for incoming JSON-RPC messages from clients. * @param request The HTTP servlet request containing the JSON-RPC message @@ -463,6 +525,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 { @@ -514,6 +577,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)) @@ -539,8 +604,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(); } @@ -694,6 +767,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. @@ -703,7 +840,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; @@ -767,9 +903,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(); @@ -813,8 +957,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); } @@ -853,12 +995,14 @@ public static class Builder { private McpTransportContextExtractor contextExtractor = ( serverRequest) -> McpTransportContext.EMPTY; - private Duration keepAliveInterval; + private Duration keepAliveInterval = Duration.ofMinutes(30); private ServerHttpHeaderValidator httpHeaderValidator = ServerHttpHeaderValidator.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. @@ -910,7 +1054,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) { @@ -958,6 +1102,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. @@ -968,7 +1128,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, httpHeaderValidator, requestMaxSize); + contextExtractor, keepAliveInterval, httpHeaderValidator, 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 7f892df09..5ac08635c 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -1,29 +1,31 @@ /* - * 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.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; -import org.slf4j.LoggerFactory; - -import io.modelcontextprotocol.json.TypeRef; - import io.modelcontextprotocol.common.McpTransportContext; +import io.modelcontextprotocol.json.TypeRef; import io.modelcontextprotocol.json.schema.JsonSchemaValidator; import io.modelcontextprotocol.server.McpAsyncServerExchange; import io.modelcontextprotocol.server.McpNotificationHandler; import io.modelcontextprotocol.server.McpRequestHandler; import io.modelcontextprotocol.spec.McpSchema.ErrorCodes; import io.modelcontextprotocol.util.Assert; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.publisher.MonoSink; @@ -36,12 +38,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; @@ -178,6 +194,27 @@ public Mono delete() { })); } + /** + * 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 @@ -190,60 +227,59 @@ public McpStreamableServerSessionStream listeningStream(McpStreamableServerTrans McpStreamableServerSessionStream listeningStream = new McpStreamableServerSessionStream(transport); McpLoggableSession replaced = this.listeningStreamRef.getAndSet(listeningStream); if (replaced instanceof McpStreamableServerSessionStream replacedStream) { - logger.debug("Closing the listening stream replaced in session {}", this.id); - replacedStream.close(); + 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( - McpSchema.JSONRPCResponse - .error(jsonrpcRequest.id(), - new McpSchema.JSONRPCResponse.JSONRPCError( - McpSchema.ErrorCodes.METHOD_NOT_FOUND, error.message(), error.data()))) - .then(transport.closeGracefully()); - } - return requestHandler - .handle(new McpAsyncServerExchange(this.id, stream, clientCapabilities.get(), clientInfo.get(), - transportContext, this.jsonSchemaValidator), jsonrpcRequest.params()) - .map(result -> McpSchema.JSONRPCResponse.result(jsonrpcRequest.id(), result)) - .onErrorResume(e -> { - McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError = (e instanceof McpError mcpError - && mcpError.getJsonRpcError() != null) ? mcpError.getJsonRpcError() - : new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, - e.getMessage(), McpError.aggregateExceptionMessages(e)); - - var errorResponse = McpSchema.JSONRPCResponse.error(jsonrpcRequest.id(), jsonRpcError); - return Mono.just(errorResponse); - }) - .flatMap(transport::sendMessage) - .then(transport.closeGracefully()); - }); + return Mono.defer(() -> this.responseStream(transport).handle(jsonrpcRequest)); } /** @@ -277,22 +313,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) { - 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) { + 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()); } - else { - sink.success(response); - } + pendingRequest.sink().success(response); } else { logger.error("Discarded MCP request response without session id. " @@ -310,23 +337,49 @@ 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 this.onClose.get().onErrorComplete().then(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() { this.onClose.get().onErrorComplete().subscribe(); - 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"); } /** @@ -368,15 +421,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; @@ -384,10 +445,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 @@ -409,19 +471,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(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())); @@ -431,42 +494,104 @@ 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(McpSchema.JSONRPCResponse.error(jsonrpcRequest.id(), + 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, + McpStreamableServerSession.this.jsonSchemaValidator), jsonrpcRequest.params()) + .map(result -> McpSchema.JSONRPCResponse.result(jsonrpcRequest.id(), result)) + .onErrorResume(e -> { + McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError = (e instanceof McpError mcpError + && mcpError.getJsonRpcError() != null) + ? mcpError.getJsonRpcError() + : new McpSchema.JSONRPCResponse.JSONRPCError( + McpSchema.ErrorCodes.INTERNAL_ERROR, e.getMessage(), + McpError.aggregateExceptionMessages(e)); + + var errorResponse = McpSchema.JSONRPCResponse.error(jsonrpcRequest.id(), 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(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..f53be412a --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java @@ -0,0 +1,256 @@ +/* + * 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(McpSchema.JSONRPCResponse.result(requestId, "response-value")).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(McpSchema.JSONRPCResponse.result(requestId, "response-value")).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("tools/call", "request-1")).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("tools/call", "request-1")).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("tools/call", "request-1")).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("tools/call", "request-1")).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-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java index 0a918b6df..c6796ce3f 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -1,25 +1,30 @@ /* - * 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; @@ -30,6 +35,7 @@ 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; @@ -42,12 +48,23 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.provider.Arguments; +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.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 { @@ -62,6 +79,16 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg private HttpServletStreamableServerTransportProvider mcpServerTransportProvider; + // 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(); + @Override protected void awaitClientStreamEstablished() { var timeout = Duration.ofSeconds(5); @@ -107,8 +134,9 @@ public void before() { 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(); MCP_SERVLET.setDelegate(mcpServerTransportProvider); @@ -185,11 +213,10 @@ protected void prepareClients(int port, String mcpEndpoint) { @Test void rejectsWhenBodyBytesExceedLimitWithoutContentLengthHeader() throws Exception { - var httpClient = HttpClient.newHttpClient(); // 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() { @@ -224,9 +251,8 @@ public void cancel() { } @Test - void resumedStreamReceivesServerNotifications() throws Exception { + void resumedStreamReceivesServerNotifications() { prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); - var httpClient = HttpClient.newHttpClient(); var sessionId = initializeSession(httpClient); @@ -235,39 +261,302 @@ void resumedStreamReceivesServerNotifications() throws Exception { // reconnected client never receives anything again. var stream = openListeningStream(httpClient, sessionId, sessionId + "_0"); - awaitStreamOpen(stream); + awaitClientStreamEstablished(); awaitNotification(stream.events()); } @Test - void replacedListeningStreamIsClosed() throws Exception { + void replacedListeningStreamIsClosed() { prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").build(); - var httpClient = HttpClient.newHttpClient(); var sessionId = initializeSession(httpClient); var firstStream = openListeningStream(httpClient, sessionId, null); - awaitStreamOpen(firstStream); + awaitClientStreamEstablished(); awaitNotification(firstStream.events()); // stream keeps receiving pings, so we just ensure we've removed the notification firstStream.events().clear(); assertThat(firstStream.events()).noneMatch(line -> line.contains("notifications/resources/list_changed")); - // Resuming installs a new listening stream. The session can no longer - // address the first one, so it must not be left open. + // Resuming installs a new listening stream. The session can no longer address the + // first one, so it must not be left open. the first stream is only closed when + // the second stream is turned on, so we don't need to wait for a client stream to + // be established var secondStream = openListeningStream(httpClient, sessionId, sessionId + "_0"); assertThat(firstStream.streamFuture()).succeedsWithin(Duration.ofSeconds(5)); - awaitStreamOpen(secondStream); + mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null).block(); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - mcpServerTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED, null) - .block(); assertThat(secondStream.events()).anyMatch(line -> line.contains("notifications/resources/list_changed")); assertThat(firstStream.events()).noneMatch(line -> line.contains("notifications/resources/list_changed")); }); } - private String initializeSession(HttpClient httpClient) throws Exception { + @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("hangs", 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)) .header("Content-Type", "application/json") @@ -278,14 +567,20 @@ private String initializeSession(HttpClient httpClient) throws Exception { "clientInfo":{"name":"test-client","version":"1.0.0"}}}""")) .build(); - var response = httpClient.send(initialize, HttpResponse.BodyHandlers.ofString()); + HttpResponse response = null; + try { + response = httpClient.send(initialize, HttpResponse.BodyHandlers.ofString()); + } + catch (IOException | InterruptedException e) { + return null; + } assertThat(response.statusCode()).isEqualTo(HttpServletResponse.SC_OK); - return response.headers().firstValue(HttpHeaders.MCP_SESSION_ID).orElseThrow(); + 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 future completing once the server closes the stream + * @return a handle on the stream, whose future completes once the server closes it */ private StreamResponse openListeningStream(HttpClient httpClient, String sessionId, String lastEventId) { var get = HttpRequest.newBuilder() @@ -296,13 +591,21 @@ private StreamResponse openListeningStream(HttpClient httpClient, String session get.header(HttpHeaders.LAST_EVENT_ID, lastEventId); } Queue events = new ConcurrentLinkedQueue<>(); - var eventsReceived = new AtomicBoolean(false); - var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofLines()) + var streamRef = new AtomicReference(); + var clientFuture = httpClient.sendAsync(get.GET().build(), HttpResponse.BodyHandlers.ofInputStream()) .thenAccept(response -> { - eventsReceived.set(true); - response.body().forEach(events::add); + streamRef.set(response.body()); + try (var r = new BufferedReader(new InputStreamReader(response.body(), UTF_8))) { + String l; + while ((l = r.readLine()) != null) { + events.add(l); + } + } + catch (IOException e) { + // "closed" here is our own stop(), not a failure + } }); - return new StreamResponse(clientFuture, eventsReceived, events); + return new StreamResponse(clientFuture, events, streamRef); } private void awaitNotification(Queue events) { @@ -312,11 +615,35 @@ private void awaitNotification(Queue events) { }); } - private static void awaitStreamOpen(StreamResponse stream) { - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> assertThat(stream.isOpen()).isTrue()); + 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; + + }); + } } - record StreamResponse(CompletableFuture streamFuture, AtomicBoolean isOpen, Queue events) { + 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(); } }