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>

mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java

@@ -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();
 		});

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

@@ -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();
+	}
+
+}

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

@@ -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");
+	}
+
+}

mcp-core/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java

@@ -10,7 +10,7 @@ import eu.rekawek.toxiproxy.model.ToxicDirection;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransport;
-import org.junit.jupiter.api.Disabled;
+import io.modelcontextprotocol.spec.McpTransportSessionClosedException;
 import org.junit.jupiter.api.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -222,9 +222,10 @@ public abstract class AbstractMcpAsyncClientResiliencyTests {
 			// In case of Streamable HTTP this call should issue a HTTP DELETE request
 			// invalidating the session
 			StepVerifier.create(mcpAsyncClient.closeGracefully()).expectComplete().verify();
-			// The next use should immediately re-initialize with no issue and send the
+			// The next tries to use the closed session and fails
-			// request without any broken connections.
+			StepVerifier.create(mcpAsyncClient.ping())
-			StepVerifier.create(mcpAsyncClient.ping()).expectNextCount(1).verifyComplete();
+				.expectErrorMatches(err -> err.getCause() instanceof McpTransportSessionClosedException)
+				.verify();
 		});
 	}
 

mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportTest.java

@@ -125,4 +125,39 @@ class HttpClientStreamableHttpTransportTest {
 		});
 	}
 
+	@Test
+	void testCloseUninitialized() {
+		var transport = HttpClientStreamableHttpTransport.builder(host).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 = HttpClientStreamableHttpTransport.builder(host).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();
+	}
+
 }

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

@@ -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();
 		});

mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportTest.java

@@ -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();
+	}
+
+}

mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java

@@ -10,7 +10,7 @@ import eu.rekawek.toxiproxy.model.ToxicDirection;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransport;
-import org.junit.jupiter.api.Disabled;
+import io.modelcontextprotocol.spec.McpTransportSessionClosedException;
 import org.junit.jupiter.api.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -221,9 +221,10 @@ public abstract class AbstractMcpAsyncClientResiliencyTests {
 			// In case of Streamable HTTP this call should issue a HTTP DELETE request
 			// invalidating the session
 			StepVerifier.create(mcpAsyncClient.closeGracefully()).expectComplete().verify();
-			// The next use should immediately re-initialize with no issue and send the
+			// The next tries to use the closed session and fails
-			// request without any broken connections.
+			StepVerifier.create(mcpAsyncClient.ping())
-			StepVerifier.create(mcpAsyncClient.ping()).expectNextCount(1).verifyComplete();
+				.expectErrorMatches(err -> err.getCause() instanceof McpTransportSessionClosedException)
+				.verify();
 		});
 	}