fix: avoid dropped errors when transport is closed or uninitialized (#995)
Signed-off-by: Dariusz Jędrzejczyk <2554306+chemicL@users.noreply.github.com>
This commit is contained in:
parent
dbb9bdad51
commit
df75857fb1
5 changed files with 109 additions and 47 deletions
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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
|
||||
* {@link McpTransportSessionClosedException}.
|
||||
* Represents a closed MCP session, which may not be reused.
|
||||
*
|
||||
* @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) {
|
||||
this.sessionId = sessionId;
|
||||
private ClosedMcpTransportSession() {
|
||||
}
|
||||
|
||||
@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) {
|
||||
throw new McpTransportSessionClosedException(sessionId);
|
||||
public void addConnection(Disposable connection) {
|
||||
throw new IllegalStateException("MCP Session is already closed");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void removeConnection(CONNECTION connection) {
|
||||
throw new McpTransportSessionClosedException(sessionId);
|
||||
public void removeConnection(Disposable connection) {
|
||||
throw new IllegalStateException("MCP Session is already closed");
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
Loading…
Reference in a new issue