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 McpTransportSession on transport close (#632)

Browse source 
* 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
2 changed files with 229 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

17
mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java
View file

@ -25,6 +25,7 @@ import org.springframework.web.reactive.function.client.WebClientResponseExcepti
import io.modelcontextprotocol.json.TypeRef;
import io.modelcontextprotocol.json.McpJsonMapper;
import io.modelcontextprotocol.spec.ClosedMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportSession;
import io.modelcontextprotocol.spec.DefaultMcpTransportStream;
import io.modelcontextprotocol.spec.HttpHeaders;
@ -98,7 +99,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
private final boolean resumableStreams;
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<>();
@ -143,7 +144,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
});
}
private DefaultMcpTransportSession createTransportSession() {
private McpTransportSession<Disposable> createTransportSession() {
Function<String, Publisher<Void>> onClose = sessionId -> sessionId == null ? Mono.empty()
: webClient.delete()
.uri(this.endpoint)
@ -159,6 +160,14 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
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);
}
@Override
public void setExceptionHandler(Consumer<Throwable> handler) {
logger.debug("Exception handler registered");
@ -182,9 +191,9 @@ public class WebClientStreamableHttpTransport 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();
});

77
mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportTest.java Normal file
View file

@ -0,0 +1,77 @@
/*
* Copyright 2024-2025 the original author or authors.
*/
package io.modelcontextprotocol.client.transport;
import io.modelcontextprotocol.spec.McpSchema;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import reactor.test.StepVerifier;
import org.springframework.web.reactive.function.client.WebClient;
class WebClientStreamableHttpTransportTest {
static String host = "http://localhost:3001";
static WebClient.Builder builder;
@SuppressWarnings("resource")
static GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v3")
.withCommand("node dist/index.js streamableHttp")
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
.withExposedPorts(3001)
.waitingFor(Wait.forHttp("/").forStatusCode(404));
@BeforeAll
static void startContainer() {
container.start();
int port = container.getMappedPort(3001);
host = "http://" + container.getHost() + ":" + port;
builder = WebClient.builder().baseUrl(host);
}
@AfterAll
static void stopContainer() {
container.stop();
}
@Test
void testCloseUninitialized() {
var transport = WebClientStreamableHttpTransport.builder(builder).build();
StepVerifier.create(transport.closeGracefully()).verifyComplete();
var initializeRequest = new McpSchema.InitializeRequest(McpSchema.LATEST_PROTOCOL_VERSION,
McpSchema.ClientCapabilities.builder().roots(true).build(),
new McpSchema.Implementation("Spring AI MCP Client", "0.3.1"));
var testMessage = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE,
"test-id", initializeRequest);
StepVerifier.create(transport.sendMessage(testMessage))
.expectErrorMessage("MCP session has been closed")
.verify();
}
@Test
void testCloseInitialized() {
var transport = WebClientStreamableHttpTransport.builder(builder).build();
var initializeRequest = new McpSchema.InitializeRequest(McpSchema.LATEST_PROTOCOL_VERSION,
McpSchema.ClientCapabilities.builder().roots(true).build(),
new McpSchema.Implementation("Spring AI MCP Client", "0.3.1"));
var testMessage = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE,
"test-id", initializeRequest);
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
StepVerifier.create(transport.closeGracefully()).verifyComplete();
StepVerifier.create(transport.sendMessage(testMessage))
.expectErrorMatches(err -> err.getMessage().matches("MCP session with ID [a-zA-Z0-9-]* has been closed"))
.verify();
}
}