Error pending requests on session close

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

Author: Kehrlann

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

@@ -5,13 +5,15 @@
 package io.modelcontextprotocol.spec;
 
 import java.time.Duration;
+import java.util.ArrayList;
 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.Predicate;
 import java.util.function.Supplier;
 
 import io.modelcontextprotocol.common.McpTransportContext;
@@ -41,7 +43,14 @@ public class McpStreamableServerSession implements McpLoggableSession {
 
 	private static final Logger logger = LoggerFactory.getLogger(McpStreamableServerSession.class);
 
-	private final ConcurrentHashMap<String, McpStreamableServerSessionStream> requestIdToStream = new ConcurrentHashMap<>();
+	/**
+	 * Every server-initiated request still awaiting its response, keyed by request ID.
+	 * Tracked per session rather than per stream, because the client answers over a
+	 * separate HTTP POST which outlives the stream the request was sent on: a stream
+	 * going away must not lose the requests it carried, but the end of the session must
+	 * resolve all of them.
+	 */
+	private final ConcurrentHashMap<String, PendingRequest> pendingRequests = new ConcurrentHashMap<>();
 
 	/**
 	 * Every stream with a connection currently attached, whether the listening stream or
@@ -313,22 +322,13 @@ public Mono<Void> accept(McpSchema.JSONRPCResponse response) {
 			logger.debug("Received response: {}", response);
 
 			if (response.id() != null) {
-				var stream = this.requestIdToStream.get(response.id());
-				if (stream == null) {
-					return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR)
-						.message("Unexpected response for unknown id " + response.id())
-						.build());
-				}
-				// TODO: encapsulate this inside the stream itself
-				var sink = stream.pendingResponses.remove(response.id());
-				if (sink == null) {
+				var pendingRequest = this.pendingRequests.remove(response.id());
+				if (pendingRequest == null) {
 					return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR)
 						.message("Unexpected response for unknown id " + response.id())
 						.build());
 				}
-				else {
-					sink.success(response);
-				}
+				pendingRequest.sink().success(response);
 			}
 			else {
 				logger.error("Discarded MCP request response without session id. "
@@ -346,13 +346,38 @@ private MethodNotFoundError getMethodNotFoundError(String method) {
 		return new MethodNotFoundError(method, "Method not found: " + method, null);
 	}
 
+	/**
+	 * Fails the matching pending requests with the given reason. Each request is removed
+	 * before its sink is failed, so that the cleanup the failure triggers on the
+	 * requesting side finds nothing left to remove.
+	 * @param match selects the requests to fail
+	 * @param reason the message of the error the sinks are failed with
+	 */
+	private void failPendingRequests(Predicate<PendingRequest> match, String reason) {
+		List<String> requestIds = new ArrayList<>();
+		this.pendingRequests.forEach((requestId, pendingRequest) -> {
+			if (match.test(pendingRequest)) {
+				requestIds.add(requestId);
+			}
+		});
+		for (String requestId : requestIds) {
+			PendingRequest pendingRequest = this.pendingRequests.remove(requestId);
+			if (pendingRequest != null) {
+				pendingRequest.sink().error(new RuntimeException(reason));
+			}
+		}
+	}
+
 	@Override
 	public Mono<Void> closeGracefully() {
 		return this.onClose.get().onErrorComplete().then(Mono.defer(() -> {
 			this.listeningStreamRef.set(this.missingMcpTransportSession);
 			return Flux.fromIterable(List.copyOf(this.openStreams))
 				.flatMap(McpStreamableServerSessionStream::closeGracefully)
-				.then();
+				.then()
+				// Also the requests of the streams already released: the session ending
+				// is the point at which no client can answer them anymore
+				.then(Mono.fromRunnable(() -> this.failPendingRequests(request -> true, "Session closed")));
 		}));
 	}
 
@@ -361,6 +386,9 @@ public void close() {
 		this.onClose.get().onErrorComplete().subscribe();
 		this.listeningStreamRef.set(this.missingMcpTransportSession);
 		List.copyOf(this.openStreams).forEach(McpStreamableServerSessionStream::close);
+		// Also the requests of the streams already released: the session ending is the
+		// point at which no client can answer them anymore
+		this.failPendingRequests(request -> true, "Session closed");
 	}
 
 	/**
@@ -402,14 +430,22 @@ public record McpStreamableServerSessionInit(McpStreamableServerSession session,
 			Mono<McpSchema.InitializeResult> initResult) {
 	}
 
+	/**
+	 * A server-initiated request awaiting its response, along with the stream it was sent
+	 * on.
+	 *
+	 * @param stream the stream the request was sent on
+	 * @param sink the sink to resolve once the response arrives
+	 */
+	private record PendingRequest(McpStreamableServerSessionStream stream, MonoSink<McpSchema.JSONRPCResponse> sink) {
+	}
+
 	/**
 	 * An individual SSE stream within a Streamable HTTP context. Can be either the
 	 * listening GET SSE stream or a request-specific POST SSE stream.
 	 */
 	public final class McpStreamableServerSessionStream implements McpLoggableSession {
 
-		private final ConcurrentHashMap<String, MonoSink<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
-
 		private final McpStreamableServerTransport connection;
 
 		private final String transportId;
@@ -444,10 +480,10 @@ public boolean isNotificationForLevelAllowed(McpSchema.LoggingLevel loggingLevel
 		public <T> Mono<T> sendRequest(String method, Object requestParams, TypeRef<T> typeRef) {
 			String requestId = McpStreamableServerSession.this.generateRequestId();
 
-			McpStreamableServerSession.this.requestIdToStream.put(requestId, this);
-
 			return Mono.<McpSchema.JSONRPCResponse>create(sink -> {
-				this.pendingResponses.put(requestId, sink);
+				// Registered on subscription rather than on assembly, so that a request
+				// which is never sent does not leave the session tracking it forever
+				McpStreamableServerSession.this.pendingRequests.put(requestId, new PendingRequest(this, sink));
 				McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(method, requestId,
 						requestParams);
 				String messageId = this.uuidGenerator.get();
@@ -457,8 +493,7 @@ public <T> Mono<T> sendRequest(String method, Object requestParams, TypeRef<T> t
 			}).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);
+				McpStreamableServerSession.this.pendingRequests.remove(requestId);
 			}).handle((jsonRpcResponse, sink) -> {
 				if (jsonRpcResponse.error() != null) {
 					sink.error(new McpError(jsonRpcResponse.error()));
@@ -486,25 +521,23 @@ public Mono<Void> sendNotification(String method, Object params) {
 		public Mono<Void> closeGracefully() {
 			return Mono.defer(() -> {
 				McpStreamableServerSession.this.openStreams.remove(this);
-				this.pendingResponses.values().forEach(s -> s.error(new RuntimeException("Stream closed")));
-				this.pendingResponses.clear();
+				McpStreamableServerSession.this.failPendingRequests(request -> this.equals(request.stream()),
+						"Stream closed");
 				// 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.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();
+			McpStreamableServerSession.this.failPendingRequests(request -> this.equals(request.stream()),
+					"Stream closed");
 			// If this was the generic stream, reset it
 			McpStreamableServerSession.this.listeningStreamRef.compareAndExchange(this,
 					McpStreamableServerSession.this.missingMcpTransportSession);
-			McpStreamableServerSession.this.requestIdToStream.values().removeIf(this::equals);
 			this.connection.close();
 		}
 

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

@@ -74,6 +74,71 @@ void closingTheSessionReleasesTheConnectionOfItsStreams() {
 		assertThat(transport.closed).isTrue();
 	}
 
+	@Test
+	void closingTheSessionFailsThePendingRequestsOfReleasedStreams() {
+		var session = session();
+		session.listeningStream(new RecordingTransport());
+
+		var pending = session.sendRequest("sampling/createMessage", null, new TypeRef<String>() {
+		}).toFuture();
+
+		// The stream the request was sent on is released, which leaves the request
+		// pending: a reconnecting client can still answer it over a separate POST
+		session.releaseListeningStream();
+		assertThat(pending).isNotDone();
+
+		// The client cannot respond on a closed session, the pending request must be
+		// failed.
+		session.close();
+
+		assertThat(pending).failsWithin(TIMEOUT).withThrowableThat().havingCause().withMessage("Session closed");
+	}
+
+	@Test
+	void closingTheSessionGracefullyFailsThePendingRequestsOfReleasedStreams() {
+		var session = session();
+		session.listeningStream(new RecordingTransport());
+
+		var pending = session.sendRequest("sampling/createMessage", null, new TypeRef<String>() {
+		}).toFuture();
+
+		session.releaseListeningStream();
+		assertThat(pending).isNotDone();
+
+		session.closeGracefully().block(TIMEOUT);
+
+		assertThat(pending).failsWithin(TIMEOUT).withThrowableThat().havingCause().withMessage("Session closed");
+	}
+
+	@Test
+	void closingAStreamFailsOnlyItsOwnPendingRequests() {
+		var session = session();
+		var listeningTransport = new RecordingTransport();
+		session.listeningStream(listeningTransport);
+
+		var onListeningStream = session.sendRequest("sampling/createMessage", null, new TypeRef<String>() {
+		}).toFuture();
+
+		// A POST response stream carries its own server-initiated requests, which can be
+		// closed independently
+		var responseTransport = new RecordingTransport();
+		var responseStream = session.new McpStreamableServerSessionStream(responseTransport);
+		var onResponseStream = responseStream.sendRequest("elicitation/create", null, new TypeRef<String>() {
+		}).toFuture();
+
+		responseStream.close();
+
+		assertThat(onResponseStream).failsWithin(TIMEOUT)
+			.withThrowableThat()
+			.havingCause()
+			.withMessage("Stream closed");
+		assertThat(onListeningStream).isNotDone();
+
+		var requestId = ((McpSchema.JSONRPCRequest) listeningTransport.sent.peek()).id();
+		session.accept(McpSchema.JSONRPCResponse.result(requestId, "response-value")).block(TIMEOUT);
+		assertThat(onListeningStream).succeedsWithin(TIMEOUT).isEqualTo("response-value");
+	}
+
 	static class RecordingTransport implements McpStreamableServerTransport {
 
 		final Queue<McpSchema.JSONRPCMessage> sent = new ConcurrentLinkedQueue<>();