Fix hanging disconnected client during POST

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

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