fix: avoid dropped errors when transport is closed or uninitialized (#995)

Signed-off-by: Dariusz Jędrzejczyk <2554306+chemicL@users.noreply.github.com>

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

@@ -38,6 +38,7 @@ import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransportException;
 import io.modelcontextprotocol.spec.McpTransportSession;
+import io.modelcontextprotocol.spec.McpTransportSessionClosedException;
 import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
 import io.modelcontextprotocol.spec.McpTransportStream;
 import io.modelcontextprotocol.spec.ProtocolVersions;
@@ -189,14 +190,6 @@ public class HttpClientStreamableHttpTransport 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);
-	}
-
 	private Publisher<Void> createDelete(String sessionId) {
 
 		var uri = Utils.resolveUri(this.baseUri, this.endpoint);
@@ -240,7 +233,8 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 	public Mono<Void> closeGracefully() {
 		return Mono.defer(() -> {
 			logger.debug("Graceful close triggered");
-			McpTransportSession<Disposable> currentSession = this.activeSession.getAndUpdate(this::createClosedSession);
+			McpTransportSession<Disposable> currentSession = this.activeSession
+				.getAndSet(ClosedMcpTransportSession.INSTANCE);
 			if (currentSession != null) {
 				return Mono.from(currentSession.closeGracefully());
 			}
@@ -250,6 +244,19 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 
 	private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
 		return Mono.deferContextual(ctx -> {
+			var rh = this.handler.get();
+			if (rh == null) {
+				logger.warn("Transport has no request handler registered. Remember to call connect!");
+			}
+
+			final Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> requestHandler = rh != null
+					? rh : msg -> Mono.error(new IllegalStateException("No request handler"));
+
+			final McpTransportSession<Disposable> transportSession = this.activeSession.get();
+
+			if (ClosedMcpTransportSession.INSTANCE.equals(transportSession)) {
+				throw new McpTransportSessionClosedException();
+			}
 
 			if (stream != null) {
 				logger.debug("Reconnecting stream {} with lastId {}", stream.streamId(), stream.lastId());
@@ -259,7 +266,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 			}
 
 			final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
-			final McpTransportSession<Disposable> transportSession = this.activeSession.get();
+
 			var uri = Utils.resolveUri(this.baseUri, this.endpoint);
 
 			Disposable connection = Mono.deferContextual(connectionCtx -> {
@@ -389,18 +396,18 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 								"Received unrecognized SSE event type: " + sseResponseEvent.sseEvent().event()));
 					})
 					.retryWhen(authorizationErrorRetrySpec())
-					.flatMap(jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
+					.flatMap(jsonrpcMessage -> requestHandler.apply(Mono.just(jsonrpcMessage)))
 					.onErrorMap(CompletionException.class, t -> t.getCause())
-					.onErrorComplete(t -> {
-						this.handleException(t);
-						return true;
-					})
 					.doFinally(s -> {
 						Disposable ref = disposableRef.getAndSet(null);
 						if (ref != null) {
 							transportSession.removeConnection(ref);
 						}
 					}))
+				.onErrorComplete(t -> {
+					this.handleException(t);
+					return true;
+				})
 				.contextWrite(ctx)
 				.subscribe();
 
@@ -467,10 +474,23 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 
 	public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
 		return Mono.create(deliveredSink -> {
+			var rh = this.handler.get();
+			if (rh == null) {
+				logger.warn("Transport has no request handler registered. Remember to call connect!");
+			}
+
+			final Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> requestHandler = rh != null
+					? rh : msg -> Mono.error(new IllegalStateException("No request handler"));
+
+			var transportSession = this.activeSession.get();
+
+			if (ClosedMcpTransportSession.INSTANCE.equals(transportSession)) {
+				throw new McpTransportSessionClosedException();
+			}
+
 			logger.debug("Sending message {}", sentMessage);
 
 			final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
-			final McpTransportSession<Disposable> transportSession = this.activeSession.get();
 
 			var uri = Utils.resolveUri(this.baseUri, this.endpoint);
 			String jsonBody = this.toString(sentMessage);
@@ -643,22 +663,26 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 						new RuntimeException("Failed to send message: " + responseEvent));
 			})
 				.retryWhen(authorizationErrorRetrySpec())
-				.flatMap(jsonRpcMessage -> this.handler.get().apply(Mono.just(jsonRpcMessage)))
+				.flatMap(jsonRpcMessage -> requestHandler.apply(Mono.just(jsonRpcMessage)))
 				.onErrorMap(CompletionException.class, t -> t.getCause())
-				.onErrorComplete(t -> {
-					// handle the error first
-					this.handleException(t);
-					// inform the caller of sendMessage
-					deliveredSink.error(t);
-					return true;
-				})
 				.doFinally(s -> {
 					logger.debug("SendMessage finally: {}", s);
 					Disposable ref = disposableRef.getAndSet(null);
 					if (ref != null) {
 						transportSession.removeConnection(ref);
 					}
-				})).contextWrite(deliveredSink.contextView()).subscribe();
+				})).onErrorComplete(t -> {
+					// handle the error first
+					try {
+						this.handleException(t);
+					}
+					catch (Exception e) {
+						logger.error("Error handling exception {}", t.getMessage(), e);
+					}
+					// inform the caller of sendMessage
+					deliveredSink.error(t);
+					return true;
+				}).contextWrite(deliveredSink.contextView()).subscribe();
 
 			disposableRef.set(connection);
 			transportSession.addConnection(connection);

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

@@ -6,43 +6,41 @@ package io.modelcontextprotocol.spec;
 import java.util.Optional;
 
 import org.reactivestreams.Publisher;
+import reactor.core.Disposable;
 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
+ * Represents a closed MCP session, which may not be reused.
- * {@link McpTransportSessionClosedException}.
  *
- * @param <CONNECTION> the resource representing the connection that the transport
- * manages.
  * @author Daniel Garnier-Moiroux
+ * @author Dariusz Jędrzejczyk
  */
-public class ClosedMcpTransportSession<CONNECTION> implements McpTransportSession<CONNECTION> {
+public final class ClosedMcpTransportSession implements McpTransportSession<Disposable> {
 
-	private final String sessionId;
+	public static final ClosedMcpTransportSession INSTANCE = new ClosedMcpTransportSession();
 
-	public ClosedMcpTransportSession(@Nullable String sessionId) {
+	private ClosedMcpTransportSession() {
-		this.sessionId = sessionId;
 	}
 
 	@Override
 	public Optional<String> sessionId() {
-		throw new McpTransportSessionClosedException(sessionId);
+		return Optional.empty();
 	}
 
 	@Override
 	public boolean markInitialized(String sessionId) {
-		throw new McpTransportSessionClosedException(sessionId);
+		throw new IllegalStateException("MCP Session is already closed");
 	}
 
 	@Override
-	public void addConnection(CONNECTION connection) {
+	public void addConnection(Disposable connection) {
-		throw new McpTransportSessionClosedException(sessionId);
+		throw new IllegalStateException("MCP Session is already closed");
 	}
 
 	@Override
-	public void removeConnection(CONNECTION connection) {
+	public void removeConnection(Disposable connection) {
-		throw new McpTransportSessionClosedException(sessionId);
+		throw new IllegalStateException("MCP Session is already closed");
 	}
 
 	@Override

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

@@ -119,7 +119,8 @@ public class McpClientSession implements McpSession {
 		this.requestHandlers.putAll(requestHandlers);
 		this.notificationHandlers.putAll(notificationHandlers);
 
-		this.transport.connect(mono -> mono.doOnNext(this::handle)).transform(connectHook).subscribe();
+		this.transport.connect(mono -> mono.doOnNext(this::handle)).transform(connectHook).subscribe(ignored -> {
+		}, error -> logger.warn("Client failed during connect", error));
 	}
 
 	private void dismissPendingResponses() {
@@ -160,7 +161,15 @@ public class McpClientSession implements McpSession {
 				var errorResponse = McpSchema.JSONRPCResponse.error(request.id(), jsonRpcError);
 				return Mono.just(errorResponse);
 			}).flatMap(this.transport::sendMessage).onErrorComplete(t -> {
-				logger.warn("Issue sending response to the client, ", t);
+				if (t instanceof McpTransportSessionClosedException) {
+					logger.debug("Can't send response to request {} when the transport is closed", request.id());
+				}
+				else if (McpTransport.isPeerClosed(t)) {
+					logger.debug("Can't send response to request {}: connection closed by peer", request.id(), t);
+				}
+				else {
+					logger.warn("Failed to send response to the server", t);
+				}
 				return true;
 			}).subscribe();
 		}

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

@@ -8,6 +8,8 @@ import java.util.List;
 
 import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
 import io.modelcontextprotocol.json.TypeRef;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Mono;
 
 /**
@@ -39,6 +41,8 @@ import reactor.core.publisher.Mono;
  */
 public interface McpTransport {
 
+	Logger logger = LoggerFactory.getLogger(McpTransport.class);
+
 	/**
 	 * Closes the transport connection and releases any associated resources.
 	 *
@@ -48,7 +52,24 @@ public interface McpTransport {
 	 * </p>
 	 */
 	default void close() {
-		this.closeGracefully().subscribe();
+		this.closeGracefully().subscribe(ignored -> {
+		}, error -> {
+			if (isPeerClosed(error)) {
+				logger.debug("Error during asynchronous close", error);
+			}
+			else {
+				logger.warn("Error during asynchronous close", error);
+			}
+		});
+	}
+
+	static boolean isPeerClosed(Throwable t) {
+		for (Throwable c = t; c != null; c = c.getCause()) {
+			if (c instanceof java.io.EOFException) {
+				return true;
+			}
+		}
+		return false;
 	}
 
 	/**

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

@@ -13,8 +13,14 @@ import reactor.util.annotation.Nullable;
  * @see ClosedMcpTransportSession
  * @author Daniel Garnier-Moiroux
  */
+
 public class McpTransportSessionClosedException extends RuntimeException {
 
+	public McpTransportSessionClosedException() {
+		super("Transport has already been closed.");
+	}
+
+	@Deprecated(forRemoval = true)
 	public McpTransportSessionClosedException(@Nullable String sessionId) {
 		super(sessionId != null ? "MCP session with ID %s has been closed".formatted(sessionId)
 				: "MCP session has been closed");

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

@@ -464,7 +464,7 @@ public abstract class AbstractMcpAsyncClientTests {
 	void testInitializeWithRootsListProviders() {
 		withClient(createMcpTransport(),
 				builder -> builder.roots(Root.builder("file:///test/path").name("test-root").build()), client -> {
-					StepVerifier.create(client.initialize().then(client.closeGracefully())).verifyComplete();
+					StepVerifier.create(client.initialize()).expectNextCount(1).verifyComplete();
 				});
 	}
 
@@ -728,8 +728,6 @@ public abstract class AbstractMcpAsyncClientTests {
 				builder -> builder.loggingConsumer(notification -> Mono.fromRunnable(() -> logReceived.set(true))),
 				client -> {
 					StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
-					StepVerifier.create(client.closeGracefully()).verifyComplete();
-
 				});
 
 	}

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

@@ -8,11 +8,14 @@ import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientReq
 import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.McpTransportSessionClosedException;
 import io.modelcontextprotocol.spec.ProtocolVersions;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Map;
 import java.util.function.Consumer;
+import java.util.function.Function;
+
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -139,13 +142,14 @@ class HttpClientStreamableHttpTransportTest {
 		var testMessage = new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "test-id", initializeRequest);
 
 		StepVerifier.create(transport.sendMessage(testMessage))
-			.expectErrorMessage("MCP session has been closed")
+			.expectErrorMessage("Transport has already been closed.")
 			.verify();
 	}
 
 	@Test
 	void testCloseInitialized() {
 		var transport = HttpClientStreamableHttpTransport.builder(host).build();
+		transport.connect(Function.identity()).block();
 
 		var initializeRequest = McpSchema.InitializeRequest
 			.builder(ProtocolVersions.MCP_2025_11_25, McpSchema.ClientCapabilities.builder().roots(true).build(),
@@ -157,7 +161,8 @@ class HttpClientStreamableHttpTransportTest {
 		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"))
+			.expectErrorMatches(err -> err instanceof McpTransportSessionClosedException
+					&& err.getMessage().contains("Transport has already been closed"))
 			.verify();
 	}
 

mcp-test/src/test/java/io/modelcontextprotocol/server/transport/ServerTransportSecurityIntegrationTests.java

@@ -81,6 +81,7 @@ class ServerTransportSecurityIntegrationTests {
 
 	@AfterEach
 	void tearDown() {
+		requestCustomizer.reset();
 		mcpClient.close();
 	}