Handle HTTP DELETE upon client close

Author: chemicL

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

@@ -10,6 +10,7 @@
 import io.modelcontextprotocol.spec.McpSessionNotFoundException;
 import io.modelcontextprotocol.spec.McpTransportSession;
 import io.modelcontextprotocol.spec.McpTransportStream;
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.core.ParameterizedTypeReference;
@@ -31,6 +32,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 
 public class WebClientStreamableHttpTransport implements McpClientTransport {
 
@@ -69,7 +71,24 @@ public WebClientStreamableHttpTransport(ObjectMapper objectMapper, WebClient.Bui
 		this.endpoint = endpoint;
 		this.resumableStreams = resumableStreams;
 		this.openConnectionOnStartup = openConnectionOnStartup;
-		this.activeSession.set(new DefaultMcpTransportSession());
+		this.activeSession.set(createTransportSession());
+	}
+
+	private DefaultMcpTransportSession createTransportSession() {
+		Supplier<Publisher<Void>> onClose = () -> {
+			DefaultMcpTransportSession transportSession = this.activeSession.get();
+			return transportSession.sessionId().isEmpty() ? Mono.empty() : webClient
+				.delete()
+				.uri(this.endpoint)
+				.headers(httpHeaders -> {
+					httpHeaders.add("mcp-session-id", transportSession.sessionId().get());
+				})
+				.retrieve()
+				.toBodilessEntity()
+					.doOnError(e -> logger.info("Got response {}", e))
+				.then();
+		};
+		return new DefaultMcpTransportSession(onClose);
 	}
 
 	@Override
@@ -93,7 +112,7 @@ public void setExceptionHandler(Consumer<Throwable> handler) {
 	private void handleException(Throwable t) {
 		logger.debug("Handling exception for session {}", sessionIdOrPlaceholder(this.activeSession.get()), t);
 		if (t instanceof McpSessionNotFoundException) {
-			McpTransportSession<?> invalidSession = this.activeSession.getAndSet(new DefaultMcpTransportSession());
+			McpTransportSession<?> invalidSession = this.activeSession.getAndSet(createTransportSession());
 			logger.warn("Server does not recognize session {}. Invalidating.", invalidSession.sessionId());
 			invalidSession.close();
 		}
@@ -107,7 +126,7 @@ private void handleException(Throwable t) {
 	public Mono<Void> closeGracefully() {
 		return Mono.defer(() -> {
 			logger.debug("Graceful close triggered");
-			DefaultMcpTransportSession currentSession = this.activeSession.get();
+			DefaultMcpTransportSession currentSession = this.activeSession.getAndSet(createTransportSession());
 			if (currentSession != null) {
 				return currentSession.closeGracefully();
 			}

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

@@ -195,4 +195,15 @@ void testCallTool() {
 		});
 	}
 
+	@Test
+	void testSessionClose() {
+		withClient(createMcpTransport(), mcpAsyncClient -> {
+			StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
+			// 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 request without any broken connections.
+			StepVerifier.create(mcpAsyncClient.ping()).expectNextCount(1).verifyComplete();
+		});
+	}
+
 }

mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -321,9 +321,11 @@ public void close() {
 	 * @return A Mono that completes when the connection is closed
 	 */
 	public Mono<Void> closeGracefully() {
-		Initialization current = this.initialization.getAndSet(null);
-		Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
-		return sessionClose.then(transport.closeGracefully());
+		return Mono.defer(() -> {
+			Initialization current = this.initialization.getAndSet(null);
+			Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
+			return sessionClose.then(transport.closeGracefully());
+		});
 	}
 
 	// --------------------------

mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java

@@ -1,5 +1,6 @@
 package io.modelcontextprotocol.spec;
 
+import org.reactivestreams.Publisher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.Disposable;
@@ -9,6 +10,7 @@
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 public class DefaultMcpTransportSession implements McpTransportSession<Disposable> {
 
@@ -20,8 +22,11 @@ public class DefaultMcpTransportSession implements McpTransportSession<Disposabl
 
 	private final AtomicReference<String> sessionId = new AtomicReference<>();
 
-	public DefaultMcpTransportSession() {
-	}
+	private final Supplier<Publisher<Void>> onClose;
+
+	public DefaultMcpTransportSession(Supplier<Publisher<Void>> onClose) {
+        this.onClose = onClose;
+    }
 
 	@Override
 	public Optional<String> sessionId() {
@@ -61,7 +66,7 @@ public void close() {
 
 	@Override
 	public Mono<Void> closeGracefully() {
-		return Mono.fromRunnable(this.openConnections::dispose);
+		return Mono.from(this.onClose.get()).then(Mono.fromRunnable(this.openConnections::dispose));
 	}
 
 }