From 04bc1e87105d415d7c224f6f7b6aeb3f3b086f3c Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Mon, 14 Sep 2026 15:49:53 +0200 Subject: [PATCH] Close POST-related streams in McpStreamableServerSession Signed-off-by: Daniel Garnier-Moiroux --- .../spec/McpStreamableServerSession.java | 96 ++++++++++----- .../spec/McpStreamableServerSessionTests.java | 111 ++++++++++++++++++ 2 files changed, 177 insertions(+), 30 deletions(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java 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 f41fb8d..e015b10 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -5,25 +5,25 @@ package io.modelcontextprotocol.spec; import java.time.Duration; +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.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; @@ -43,6 +43,12 @@ public class McpStreamableServerSession implements McpLoggableSession { private final ConcurrentHashMap requestIdToStream = 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; private final Duration requestTimeout; @@ -201,12 +207,24 @@ public class McpStreamableServerSession implements McpLoggableSession { 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) { @@ -226,9 +244,6 @@ public class McpStreamableServerSession implements McpLoggableSession { 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 @@ -237,7 +252,7 @@ public class McpStreamableServerSession implements McpLoggableSession { .error(jsonrpcRequest.id(), new McpSchema.JSONRPCResponse.JSONRPCError( McpSchema.ErrorCodes.METHOD_NOT_FOUND, error.message(), error.data()))) - .then(transport.closeGracefully()); + .then(stream.closeGracefully()); } return requestHandler .handle(new McpAsyncServerExchange(this.id, stream, clientCapabilities.get(), clientInfo.get(), @@ -253,7 +268,7 @@ public class McpStreamableServerSession implements McpLoggableSession { return Mono.just(errorResponse); }) .flatMap(transport::sendMessage) - .then(transport.closeGracefully()); + .then(stream.closeGracefully()); }); } @@ -324,20 +339,18 @@ public class McpStreamableServerSession implements McpLoggableSession { @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(); })); } @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); } /** @@ -387,7 +400,7 @@ public class McpStreamableServerSession implements McpLoggableSession { private final ConcurrentHashMap> pendingResponses = new ConcurrentHashMap<>(); - private final McpStreamableServerTransport transport; + private final McpStreamableServerTransport connection; private final String transportId; @@ -395,10 +408,11 @@ public class McpStreamableServerSession implements McpLoggableSession { /** * 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 @@ -428,9 +442,11 @@ public class McpStreamableServerSession implements McpLoggableSession { 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 -> { + }).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 this.pendingResponses.remove(requestId); McpStreamableServerSession.this.requestIdToStream.remove(requestId); }).handle((jsonRpcResponse, sink) -> { @@ -442,7 +458,7 @@ public class McpStreamableServerSession implements McpLoggableSession { sink.complete(); } else { - sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef)); + sink.next(this.connection.unmarshalFrom(jsonRpcResponse.result(), typeRef)); } } }); @@ -453,31 +469,51 @@ public class McpStreamableServerSession implements McpLoggableSession { 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(() -> { + McpStreamableServerSession.this.openStreams.remove(this); this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); this.pendingResponses.clear(); // 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() { + McpStreamableServerSession.this.openStreams.remove(this); this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed"))); this.pendingResponses.clear(); // 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(); + } + + /** + * 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); + this.connection.close(); } } 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 0000000..ae52330 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java @@ -0,0 +1,111 @@ +/* + * Copyright 2024-2025 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 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 new McpStreamableServerSession("session-1", McpSchema.ClientCapabilities.builder().build(), + new McpSchema.Implementation("test-client", "1.0.0"), TIMEOUT, Map.of(), 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(); + } + + 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; + } + + } + +}