Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,27 @@
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.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;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

/**
* Server-side implementation of the Model Context Protocol (MCP) streamable transport
Expand Down Expand Up @@ -89,6 +94,8 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet

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

private static final Duration MAX_IDLE_SWEEP_INTERVAL = Duration.ofMinutes(1);

/**
* The endpoint URI where clients should send their JSON-RPC messages. Defaults to
* "/mcp".
Expand Down Expand Up @@ -127,6 +134,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
*/
private KeepAliveScheduler keepAliveScheduler;

/**
* How long a session may stay idle before being evicted. Disabled when {@code null}.
*/
private final Duration sessionIdleTimeout;

/**
* Subscription of the periodic sweep evicting idle sessions. Only set when
* {@link #sessionIdleTimeout} is configured.
*/
private Disposable idleSessionSweep;

/**
* Security validator for validating HTTP requests.
*/
Expand All @@ -142,14 +160,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
* @param contextExtractor The extractor for transport context from the request.
* @param keepAliveInterval The interval for keep-alive pings. If null, no keep-alive
* will be scheduled.
* @param sessionIdleTimeout How long a session may stay idle before being evicted. If
* null, idle sessions are never evicted.
* @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.
* @throws IllegalArgumentException if any parameter is null
*/
private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, String mcpEndpoint,
boolean disallowDelete, McpTransportContextExtractor<HttpServletRequest> contextExtractor,
Duration keepAliveInterval, ServerHttpHeaderValidator httpHeaderValidator, int requestMaxSize) {
Duration keepAliveInterval, Duration sessionIdleTimeout, ServerHttpHeaderValidator httpHeaderValidator,
int requestMaxSize) {
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");
Expand All @@ -162,20 +183,92 @@ private HttpServletStreamableServerTransportProvider(McpJsonMapper jsonMapper, S
this.contextExtractor = contextExtractor;
this.httpHeaderValidator = httpHeaderValidator;
this.requestMaxSize = requestMaxSize;
this.sessionIdleTimeout = sessionIdleTimeout;

if (sessionIdleTimeout != null) {
Assert.isTrue(!sessionIdleTimeout.isNegative() && !sessionIdleTimeout.isZero(),
"Session idle timeout must be positive");
// Sweeping more often than the timeout bounds how long an idle session
// outlives it, without making the sweep itself expensive
Duration sweepInterval = sessionIdleTimeout.compareTo(MAX_IDLE_SWEEP_INTERVAL) < 0 ? sessionIdleTimeout
: MAX_IDLE_SWEEP_INTERVAL;
// Evicting closes connections, which is not work for the parallel scheduler
this.idleSessionSweep = Flux.interval(sweepInterval, Schedulers.boundedElastic())
.doOnNext(tick -> this.evictIdleSessions())
.onErrorContinue((error, tick) -> logger.warn("Idle session sweep failed: {}", error.getMessage()))
.subscribe();
}

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();
}

}

/**
* 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<McpSession> sessionsToPing() {
if (this.isClosing) {
return Flux.empty();
}
return Flux.fromIterable(this.sessions.values())
.filter(McpStreamableServerSession::hasListeningStream)
.cast(McpSession.class);
}

/**
* Removes a session and releases everything it holds, in particular the connections
* of its streams. Together with the client asking for it through a DELETE request,
* this is the only way a session leaves this transport.
* @param sessionId the session to evict
* @param reason why the session is being evicted, for logging
*/
private void evictSession(String sessionId, String reason) {
McpStreamableServerSession session = this.sessions.remove(sessionId);
if (session == null) {
return;
}
logger.info("Evicting session {}: {}", sessionId, reason);
session.close();
}

/**
* Evicts the sessions idle for longer than the configured session idle timeout.
* Clients are not required to send a DELETE request when they go away, so a client
* which simply vanishes, whether killed, disconnected or suspended, would otherwise
* keep its session, and the resources behind it, until the server restarts.
*/
private void evictIdleSessions() {
if (this.isClosing) {
return;
}
this.sessions.values()
.stream()
.filter(session -> session.isIdleFor(this.sessionIdleTimeout))
.map(McpStreamableServerSession::getId)
.toList()
.forEach(sessionId -> this.evictSession(sessionId, "idle for more than " + this.sessionIdleTimeout));
}

@Override
public void setSessionFactory(McpStreamableServerSession.Factory sessionFactory) {
this.sessionFactory = sessionFactory;
Expand Down Expand Up @@ -248,6 +341,9 @@ public Mono<Void> closeGracefully() {
if (this.keepAliveScheduler != null) {
this.keepAliveScheduler.shutdown();
}
if (this.idleSessionSweep != null) {
this.idleSessionSweep.dispose();
}
});
}

Expand Down Expand Up @@ -338,65 +434,14 @@ protected void doGet(HttpServletRequest request, HttpServletResponse response)
McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session
.listeningStream(sessionTransport);

asyncContext.addListener(new jakarta.servlet.AsyncListener() {
@Override
public void onComplete(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection completed for session: {}", sessionId);
listeningStream.close();
}

@Override
public void onTimeout(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection timed out for session: {}", sessionId);
listeningStream.close();
}

@Override
public void onError(jakarta.servlet.AsyncEvent event) throws IOException {
logger.debug("SSE connection error for session: {}", sessionId);
listeningStream.close();
}

@Override
public void onStartAsync(jakarta.servlet.AsyncEvent event) throws IOException {
// No action needed
}
});
registerAsyncLifecycle(asyncContext, sessionId, listeningStream::releaseTransport);
}
catch (Exception e) {
logger.error("Failed to handle GET request for session {}: {}", sessionId, e.getMessage());
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
}
}

/**
* Replays the messages the client missed while its SSE stream was broken.
* @param session the session the client is resuming
* @param lastEventId the ID of the last event received by the client
* @param sessionTransport the transport of the resumed SSE stream
* @param transportContext the context extracted from the request
* @return {@code true} if the replay completed, {@code false} if it failed, in which
* case the transport has been closed
*/
private boolean tryReplayMissedMessages(McpStreamableServerSession session, String lastEventId,
McpStreamableServerTransport sessionTransport, McpTransportContext transportContext) {
try {
for (McpSchema.JSONRPCMessage message : session.replay(lastEventId)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.toIterable()) {
sessionTransport.sendMessage(message)
.contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext))
.block();
}
return true;
}
catch (Exception e) {
logger.error("Failed to replay messages for session {}: {}", session.getId(), e.getMessage());
sessionTransport.close();
return false;
}
}

/**
* Handles POST requests for incoming JSON-RPC messages from clients.
* @param request The HTTP servlet request containing the JSON-RPC message
Expand Down Expand Up @@ -538,6 +583,7 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {

HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport(
sessionId, asyncContext, response.getWriter());
registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close);

try {
session.responseStream(jsonrpcRequest, sessionTransport)
Expand Down Expand Up @@ -694,6 +740,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.
Expand Down Expand Up @@ -767,9 +877,10 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId
logger.debug("Message sent to session {} with ID {}", this.sessionId, messageId);
}
catch (Exception e) {
// The connection is gone, the session is not: the client may come
// back for it, and the idle timeout reclaims it if it never does
logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage());
HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId);
this.asyncContext.complete();
this.close();
}
finally {
lock.unlock();
Expand Down Expand Up @@ -813,8 +924,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);
}
Expand Down Expand Up @@ -855,6 +964,8 @@ public static class Builder {

private Duration keepAliveInterval;

private Duration sessionIdleTimeout;

private ServerHttpHeaderValidator httpHeaderValidator = ServerHttpHeaderValidator.NOOP;

private int requestMaxSize = DEFAULT_REQUEST_MAX_SIZE;
Expand Down Expand Up @@ -913,6 +1024,21 @@ public Builder contextExtractor(McpTransportContextExtractor<HttpServletRequest>
* keep-alive will be scheduled.
* @return this builder instance
*/
/**
* Sets how long a session may stay idle, that is hold no connection and serve no
* request, before the server evicts it. Clients are expected to end their session
* with a DELETE request, but nothing guarantees they get the chance to, so
* without this the sessions of vanished clients are kept until the server
* restarts.
* @param sessionIdleTimeout the idle timeout. If null, idle sessions are never
* evicted.
* @return this builder instance
*/
public Builder sessionIdleTimeout(Duration sessionIdleTimeout) {
this.sessionIdleTimeout = sessionIdleTimeout;
return this;
}

public Builder keepAliveInterval(Duration keepAliveInterval) {
this.keepAliveInterval = keepAliveInterval;
return this;
Expand Down Expand Up @@ -968,7 +1094,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, sessionIdleTimeout, httpHeaderValidator, requestMaxSize);
}

}
Expand Down
Loading
Loading