Close McpTransportSession on transport close (#632)

* Client HTTP transports: use McpTransportSession interface instead of concrete types
* Streamable HTTP tranports: .closeGracefully closes the session

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux 2025-10-28 14:26:21 +01:00 • committed by GitHub
commit 9bbfbe573a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 229 additions and 16 deletions

View file

@ -29,6 +29,7 @@ import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientReq
import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
import io.modelcontextprotocol.client.transport.ResponseSubscribers.ResponseEvent;
import io.modelcontextprotocol.common.McpTransportContext;
import io.modelcontextprotocol.spec.ClosedMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportStream;
import io.modelcontextprotocol.spec.HttpHeaders;
@ -118,7 +119,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
private final McpAsyncHttpClientRequestCustomizer httpRequestCustomizer;
private final AtomicReference<DefaultMcpTransportSession> activeSession = new AtomicReference<>();
private final AtomicReference<McpTransportSession<Disposable>> activeSession = new AtomicReference<>();
private final AtomicReference<Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>> handler = new AtomicReference<>();
@ -163,12 +164,20 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
});
}
private DefaultMcpTransportSession createTransportSession() {
private McpTransportSession<Disposable> createTransportSession() {
Function<String, Publisher<Void>> onClose = sessionId -> sessionId == null ? Mono.empty()
: createDelete(sessionId);
return new DefaultMcpTransportSession(onClose);
}
private McpTransportSession<Disposable> createClosedSession(McpTransportSession<Disposable> existingSession) {
var existingSessionId = Optional.ofNullable(existingSession)
.filter(session -> !(session instanceof ClosedMcpTransportSession<Disposable>))
.flatMap(McpTransportSession::sessionId)
.orElse(null);
return new ClosedMcpTransportSession<>(existingSessionId);
}
private Publisher<Void> createDelete(String sessionId) {
var uri = Utils.resolveUri(this.baseUri, this.endpoint);
@ -210,9 +219,9 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
public Mono<Void> closeGracefully() {
return Mono.defer(() -> {
logger.debug("Graceful close triggered");
DefaultMcpTransportSession currentSession = this.activeSession.getAndSet(createTransportSession());
McpTransportSession<Disposable> currentSession = this.activeSession.getAndUpdate(this::createClosedSession);
if (currentSession != null) {
return currentSession.closeGracefully();
return Mono.from(currentSession.closeGracefully());
}
return Mono.empty();
});

View file

@ -0,0 +1,58 @@
/*
* Copyright 2025-2025 the original author or authors.
*/
package io.modelcontextprotocol.spec;
import java.util.Optional;
import org.reactivestreams.Publisher;
import reactor.core.publisher.Mono;
import reactor.util.annotation.Nullable;
/**
* Represents a closed MCP session, which may not be reused. All calls will throw a
* {@link McpTransportSessionClosedException}.
*
* @param <CONNECTION> the resource representing the connection that the transport
* manages.
* @author Daniel Garnier-Moiroux
*/
public class ClosedMcpTransportSession<CONNECTION> implements McpTransportSession<CONNECTION> {
private final String sessionId;
public ClosedMcpTransportSession(@Nullable String sessionId) {
this.sessionId = sessionId;
}
@Override
public Optional<String> sessionId() {
throw new McpTransportSessionClosedException(sessionId);
}
@Override
public boolean markInitialized(String sessionId) {
throw new McpTransportSessionClosedException(sessionId);
}
@Override
public void addConnection(CONNECTION connection) {
throw new McpTransportSessionClosedException(sessionId);
}
@Override
public void removeConnection(CONNECTION connection) {
throw new McpTransportSessionClosedException(sessionId);
}
@Override
public void close() {
}
@Override
public Publisher<Void> closeGracefully() {
return Mono.empty();
}
}

View file

@ -0,0 +1,23 @@
/*
* Copyright 2025-2025 the original author or authors.
*/
package io.modelcontextprotocol.spec;
import reactor.util.annotation.Nullable;
/**
* Exception thrown when trying to use an {@link McpTransportSession} that has been
* closed.
*
* @see ClosedMcpTransportSession
* @author Daniel Garnier-Moiroux
*/
public class McpTransportSessionClosedException extends RuntimeException {
public McpTransportSessionClosedException(@Nullable String sessionId) {
super(sessionId != null ? "MCP session with ID %s has been closed".formatted(sessionId)
: "MCP session has been closed");
}
}