Fix hanging disconnected client during POST

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java

@@ -26,6 +26,7 @@
 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 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) {
 
 				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 Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String 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
+					// 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();

mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java

@@ -250,45 +250,35 @@ public Flux<McpSchema.JSONRPCMessage> 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<Void> 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 <T> Mono<T> sendRequest(String method, Object requestParams, TypeRef<T> t
 			});
 		}
 
+		/**
+		 * Runs the request this stream was created for, sending its messages to the
+		 * client and finalizing the stream with the response.
+		 * <p>
+		 * 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<Void> 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<Void> sendNotification(String method, Object params) {
 			McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(method, params);

mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java

@@ -10,6 +10,7 @@
 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<String, McpRequestHandler<?>> 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 @@ void closingAStreamFailsOnlyItsOwnPendingRequests() {
 		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<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId) {
+			return Mono.error(new RuntimeException("connection gone"));
+		}
+
+	}
+
 	static class RecordingTransport implements McpStreamableServerTransport {
 
 		final Queue<McpSchema.JSONRPCMessage> sent = new ConcurrentLinkedQueue<>();

mcp-test/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java

@@ -49,9 +49,11 @@
 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 @@ void sessionIsNotEvictedWithoutSweepInterval() throws Exception {
 		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.<McpSchema.CallToolResult>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<String> events = new ConcurrentLinkedQueue<>();
+		var streamRef = new AtomicReference<InputStream>();
+		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 @@ record StreamResponse(CompletableFuture<Void> streamFuture, Queue<String> events
 		void closeStream() {
 			// Close listening stream. We retry a few times in case the stream was not
 			// established on the first try
-			await().atMost(Duration.ofSeconds(1)).until(() -> {
+			await().pollDelay(Duration.ZERO).atMost(Duration.ofSeconds(1)).until(() -> {
 				var stream = streamRef.get();
 				if (stream != null) {
 					stream.close();