modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Close POST-related streams in McpStreamableServerSession

Browse source 
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux 2026-09-14 15:49:53 +02:00 • committed by Daniel Garnier-Moiroux
commit 04bc1e8710
2 changed files with 177 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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<Object, McpStreamableServerSessionStream> 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<McpStreamableServerSessionStream> 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<McpSchema.JSONRPCMessage> 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<Void> 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<Object, MonoSink<McpSchema.JSONRPCResponse>> 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<Void> 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.
* <p>
* 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();
}
}

111
mcp-core/src/test/java/io/modelcontextprotocol/spec/McpStreamableServerSessionTests.java Normal file
View file

@ -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<String>() {
}).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<McpSchema.JSONRPCMessage> sent = new ConcurrentLinkedQueue<>();
volatile boolean closed;
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message, String messageId) {
return Mono.fromRunnable(() -> this.sent.add(message));
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
return sendMessage(message, null);
}
@Override
public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> this.closed = true);
}
@Override
public void close() {
this.closed = true;
}
@SuppressWarnings("unchecked")
@Override
public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
return (T) data;
}
}
}