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 90e42ed..8fa3220 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 @@ -26,6 +26,7 @@ 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; @@ -593,10 +594,17 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet HttpServletStreamableMcpSessionTransport sessionTransport = new HttpServletStreamableMcpSessionTransport( sessionId, asyncContext, response.getWriter()); - registerAsyncLifecycle(asyncContext, sessionId, sessionTransport::close); + + // 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(); } @@ -887,9 +895,16 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet } 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 + // back for it, and the sweeper reclaims it if it never does logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage()); 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(); 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 aa9b15e..d9aac28 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -250,45 +250,35 @@ public class McpStreamableServerSession implements McpLoggableSession { 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()); - 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(stream.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(stream.closeGracefully()); - }); + return Mono.defer(() -> this.responseStream(transport).handle(jsonrpcRequest)); } /** @@ -509,6 +499,50 @@ public class McpStreamableServerSession implements McpLoggableSession { }); } + /** + * 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); diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java index 7103f26..48a7eb2 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java @@ -10,6 +10,7 @@ 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; @@ -23,8 +24,12 @@ 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, Map.of(), Map.of()); + new McpSchema.Implementation("test-client", "1.0.0"), TIMEOUT, requestHandlers, Map.of()); } @Test @@ -139,6 +144,81 @@ class McpStreamableServerSessionTests { 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<>(); 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 4b32027..7efff3d 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -49,9 +49,11 @@ 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.test.StepVerifier; +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; @@ -70,9 +72,13 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg private HttpServletStreamableServerTransportProvider mcpServerTransportProvider; - private final Duration KEEP_ALIVE_INTERVAL = Duration.ofMillis(200); + // 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); - private final Duration SESSION_SWEEP_INTERVAL = Duration.ofMillis(200); + // 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); @Override protected void awaitClientStreamEstablished() { @@ -463,6 +469,90 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_ACCEPTED); } + /** + * A client which aborts a tool call must not leave its session behind. + */ + @Test + void sessionIsEvictedWhenTheClientAbortsAResponseStream() throws Exception { + mcpServerTransportProvider.closeGracefully().block(); + mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .mcpEndpoint(MESSAGE_ENDPOINT) + // remove keepalive, so only the response stream can keep the session alive + .keepAliveInterval(null) + .sessionSweepInterval(SESSION_SWEEP_INTERVAL) + .build(); + MCP_SERVLET.setDelegate(mcpServerTransportProvider); + + // A tool which never returns but keeps writing to its response stream. The + // payloads are large and frequent on purpose: writes to a connection whose peer + // is gone keep succeeding until the socket buffer fills up, and that is the only + // thing which can surface the disconnect here. + 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.ofMillis(10)) + .flatMap(tick -> exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.INFO) + .data("x".repeat(64 * 1024)) + .build())) + .then(Mono.never())) + .build()) + .build(); + + var httpClient = HttpClient.newHttpClient(); + 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)) + .untilAsserted( + () -> assertThat(responseStream.events()).anyMatch(line -> line.contains("notifications/message"))); + assertThat(postNotification(httpClient, sessionId)).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. + Thread.sleep(SESSION_SWEEP_INTERVAL.multipliedBy(2).toMillis()); + + assertThat(postNotification(httpClient, sessionId)).isEqualTo(HttpServletResponse.SC_NOT_FOUND); + } + + /** + * 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) throws Exception { var notification = HttpRequest.newBuilder() .uri(URI.create("http://localhost:" + PORT + MESSAGE_ENDPOINT)) @@ -533,7 +623,7 @@ class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerInteg void closeStream() { // Close listening stream. We retry a few times in case the stream was not // established on the first try - await().atMost(Duration.ofSeconds(1)).until(() -> { + await().pollDelay(Duration.ZERO).atMost(Duration.ofSeconds(1)).until(() -> { var stream = streamRef.get(); if (stream != null) { stream.close();