feat: WebClient Streamable HTTP support (#292)
An implementation of Streamable HTTP Client with WebFlux WebClient. Aside from implementing the specification, several improvements have been incorporated throughout the client-side of the architecture. The changes cover: - resilience tests using toxiproxy in testcontainers - integration tests using updated everything-server with streamableHttp support - improved logging - session invalidation handling (both transport session and JSON-RPC concept of session) - implicit initialization and burst protection (in case of concurrent `Mcp(Sync|Async)Client` use - more logging, e.g. stdio process lifecycle logs Related #72, #273, #253, #107, #105 Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Dariusz Jędrzejczyk
GitHub
parent
8a5a591d39
commit
2f944349cc
31 changed files with 1501 additions and 242 deletions
|
|
@ -99,6 +99,12 @@
|
|||
<version>${testcontainers.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>toxiproxy</artifactId>
|
||||
<version>${toxiproxy.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,520 @@
|
|||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.spec.DefaultMcpTransportSession;
|
||||
import io.modelcontextprotocol.spec.DefaultMcpTransportStream;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.spec.McpTransportSession;
|
||||
import io.modelcontextprotocol.spec.McpTransportStream;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import org.reactivestreams.Publisher;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.core.ParameterizedTypeReference;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.reactive.function.client.ClientResponse;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.springframework.web.reactive.function.client.WebClientResponseException;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.util.function.Tuple2;
|
||||
import reactor.util.function.Tuples;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.function.Function;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
/**
|
||||
* An implementation of the Streamable HTTP protocol as defined by the
|
||||
* <code>2025-03-26</code> version of the MCP specification.
|
||||
*
|
||||
* <p>
|
||||
* The transport is capable of resumability and reconnects. It reacts to transport-level
|
||||
* session invalidation and will propagate {@link McpTransportSessionNotFoundException
|
||||
* appropriate exceptions} to the higher level abstraction layer when needed in order to
|
||||
* allow proper state management. The implementation handles servers that are stateful and
|
||||
* provide session meta information, but can also communicate with stateless servers that
|
||||
* do not provide a session identifier and do not support SSE streams.
|
||||
* </p>
|
||||
* <p>
|
||||
* This implementation does not handle backwards compatibility with the <a href=
|
||||
* "https://modelcontextprotocol.io/specification/2024-11-05/basic/transports#http-with-sse">"HTTP
|
||||
* with SSE" transport</a>. In order to communicate over the phased-out
|
||||
* <code>2024-11-05</code> protocol, use {@link HttpClientSseClientTransport} or
|
||||
* {@link WebFluxSseClientTransport}.
|
||||
* </p>
|
||||
*
|
||||
* @author Dariusz Jędrzejczyk
|
||||
* @see <a href=
|
||||
* "https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#streamable-http">Streamable
|
||||
* HTTP transport specification</a>
|
||||
*/
|
||||
public class WebClientStreamableHttpTransport implements McpClientTransport {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(WebClientStreamableHttpTransport.class);
|
||||
|
||||
private static final String DEFAULT_ENDPOINT = "/mcp";
|
||||
|
||||
/**
|
||||
* Event type for JSON-RPC messages received through the SSE connection. The server
|
||||
* sends messages with this event type to transmit JSON-RPC protocol data.
|
||||
*/
|
||||
private static final String MESSAGE_EVENT_TYPE = "message";
|
||||
|
||||
private static final ParameterizedTypeReference<ServerSentEvent<String>> PARAMETERIZED_TYPE_REF = new ParameterizedTypeReference<>() {
|
||||
};
|
||||
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
private final WebClient webClient;
|
||||
|
||||
private final String endpoint;
|
||||
|
||||
private final boolean openConnectionOnStartup;
|
||||
|
||||
private final boolean resumableStreams;
|
||||
|
||||
private final AtomicReference<DefaultMcpTransportSession> activeSession = new AtomicReference<>();
|
||||
|
||||
private final AtomicReference<Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>> handler = new AtomicReference<>();
|
||||
|
||||
private final AtomicReference<Consumer<Throwable>> exceptionHandler = new AtomicReference<>();
|
||||
|
||||
private WebClientStreamableHttpTransport(ObjectMapper objectMapper, WebClient.Builder webClientBuilder,
|
||||
String endpoint, boolean resumableStreams, boolean openConnectionOnStartup) {
|
||||
this.objectMapper = objectMapper;
|
||||
this.webClient = webClientBuilder.build();
|
||||
this.endpoint = endpoint;
|
||||
this.resumableStreams = resumableStreams;
|
||||
this.openConnectionOnStartup = openConnectionOnStartup;
|
||||
this.activeSession.set(createTransportSession());
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a stateful builder for creating {@link WebClientStreamableHttpTransport}
|
||||
* instances.
|
||||
* @param webClientBuilder the {@link WebClient.Builder} to use
|
||||
* @return a builder which will create an instance of
|
||||
* {@link WebClientStreamableHttpTransport} once {@link Builder#build()} is called
|
||||
*/
|
||||
public static Builder builder(WebClient.Builder webClientBuilder) {
|
||||
return new Builder(webClientBuilder);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
return Mono.deferContextual(ctx -> {
|
||||
this.handler.set(handler);
|
||||
if (openConnectionOnStartup) {
|
||||
logger.debug("Eagerly opening connection on startup");
|
||||
return this.reconnect(null).then();
|
||||
}
|
||||
return Mono.empty();
|
||||
});
|
||||
}
|
||||
|
||||
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
|
||||
public void setExceptionHandler(Consumer<Throwable> handler) {
|
||||
logger.debug("Exception handler registered");
|
||||
this.exceptionHandler.set(handler);
|
||||
}
|
||||
|
||||
private void handleException(Throwable t) {
|
||||
logger.debug("Handling exception for session {}", sessionIdOrPlaceholder(this.activeSession.get()), t);
|
||||
if (t instanceof McpTransportSessionNotFoundException) {
|
||||
McpTransportSession<?> invalidSession = this.activeSession.getAndSet(createTransportSession());
|
||||
logger.warn("Server does not recognize session {}. Invalidating.", invalidSession.sessionId());
|
||||
invalidSession.close();
|
||||
}
|
||||
Consumer<Throwable> handler = this.exceptionHandler.get();
|
||||
if (handler != null) {
|
||||
handler.accept(t);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
logger.debug("Graceful close triggered");
|
||||
DefaultMcpTransportSession currentSession = this.activeSession.getAndSet(createTransportSession());
|
||||
if (currentSession != null) {
|
||||
return currentSession.closeGracefully();
|
||||
}
|
||||
return Mono.empty();
|
||||
});
|
||||
}
|
||||
|
||||
private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
|
||||
return Mono.deferContextual(ctx -> {
|
||||
if (stream != null) {
|
||||
logger.debug("Reconnecting stream {} with lastId {}", stream.streamId(), stream.lastId());
|
||||
}
|
||||
else {
|
||||
logger.debug("Reconnecting with no prior stream");
|
||||
}
|
||||
// Here we attempt to initialize the client. In case the server supports SSE,
|
||||
// we will establish a long-running
|
||||
// session here and listen for messages. If it doesn't, that's ok, the server
|
||||
// is a simple, stateless one.
|
||||
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
|
||||
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
|
||||
|
||||
Disposable connection = webClient.get()
|
||||
.uri(this.endpoint)
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.headers(httpHeaders -> {
|
||||
transportSession.sessionId().ifPresent(id -> httpHeaders.add("mcp-session-id", id));
|
||||
if (stream != null) {
|
||||
stream.lastId().ifPresent(id -> httpHeaders.add("last-event-id", id));
|
||||
}
|
||||
})
|
||||
.exchangeToFlux(response -> {
|
||||
if (isEventStream(response)) {
|
||||
return eventStream(stream, response);
|
||||
}
|
||||
else if (isNotAllowed(response)) {
|
||||
logger.debug("The server does not support SSE streams, using request-response mode.");
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (isNotFound(response)) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
return mcpSessionNotFoundError(sessionIdRepresentation);
|
||||
}
|
||||
else {
|
||||
return response.<McpSchema.JSONRPCMessage>createError().doOnError(e -> {
|
||||
logger.info("Opening an SSE stream failed. This can be safely ignored.", e);
|
||||
}).flux();
|
||||
}
|
||||
})
|
||||
.onErrorComplete(t -> {
|
||||
this.handleException(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
}
|
||||
})
|
||||
.contextWrite(ctx)
|
||||
.subscribe();
|
||||
|
||||
disposableRef.set(connection);
|
||||
transportSession.addConnection(connection);
|
||||
return Mono.just(connection);
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
return Mono.create(sink -> {
|
||||
logger.debug("Sending message {}", message);
|
||||
// Here we attempt to initialize the client.
|
||||
// In case the server supports SSE, we will establish a long-running session
|
||||
// here and
|
||||
// listen for messages.
|
||||
// If it doesn't, nothing actually happens here, that's just the way it is...
|
||||
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
|
||||
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
|
||||
|
||||
Disposable connection = webClient.post()
|
||||
.uri(this.endpoint)
|
||||
.accept(MediaType.TEXT_EVENT_STREAM, MediaType.APPLICATION_JSON)
|
||||
.headers(httpHeaders -> {
|
||||
transportSession.sessionId().ifPresent(id -> httpHeaders.add("mcp-session-id", id));
|
||||
})
|
||||
.bodyValue(message)
|
||||
.exchangeToFlux(response -> {
|
||||
if (transportSession
|
||||
.markInitialized(response.headers().asHttpHeaders().getFirst("mcp-session-id"))) {
|
||||
// Once we have a session, we try to open an async stream for
|
||||
// the server to send notifications and requests out-of-band.
|
||||
reconnect(null).contextWrite(sink.contextView()).subscribe();
|
||||
}
|
||||
|
||||
String sessionRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
|
||||
// The spec mentions only ACCEPTED, but the existing SDKs can return
|
||||
// 200 OK for notifications
|
||||
if (response.statusCode().is2xxSuccessful()) {
|
||||
Optional<MediaType> contentType = response.headers().contentType();
|
||||
// Existing SDKs consume notifications with no response body nor
|
||||
// content type
|
||||
if (contentType.isEmpty()) {
|
||||
logger.trace("Message was successfully sent via POST for session {}",
|
||||
sessionRepresentation);
|
||||
// signal the caller that the message was successfully
|
||||
// delivered
|
||||
sink.success();
|
||||
// communicate to downstream there is no streamed data coming
|
||||
return Flux.empty();
|
||||
}
|
||||
else {
|
||||
MediaType mediaType = contentType.get();
|
||||
if (mediaType.isCompatibleWith(MediaType.TEXT_EVENT_STREAM)) {
|
||||
// communicate to caller that the message was delivered
|
||||
sink.success();
|
||||
// starting a stream
|
||||
return newEventStream(response, sessionRepresentation);
|
||||
}
|
||||
else if (mediaType.isCompatibleWith(MediaType.APPLICATION_JSON)) {
|
||||
logger.trace("Received response to POST for session {}", sessionRepresentation);
|
||||
// communicate to caller the message was delivered
|
||||
sink.success();
|
||||
return responseFlux(response);
|
||||
}
|
||||
else {
|
||||
logger.warn("Unknown media type {} returned for POST in session {}", contentType,
|
||||
sessionRepresentation);
|
||||
return Flux.error(new RuntimeException("Unknown media type returned: " + contentType));
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
if (isNotFound(response)) {
|
||||
return mcpSessionNotFoundError(sessionRepresentation);
|
||||
}
|
||||
return extractError(response, sessionRepresentation);
|
||||
}
|
||||
})
|
||||
.flatMap(jsonRpcMessage -> this.handler.get().apply(Mono.just(jsonRpcMessage)))
|
||||
.onErrorResume(t -> {
|
||||
// handle the error first
|
||||
this.handleException(t);
|
||||
// inform the caller of sendMessage
|
||||
sink.error(t);
|
||||
return Flux.empty();
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
}
|
||||
})
|
||||
.contextWrite(sink.contextView())
|
||||
.subscribe();
|
||||
disposableRef.set(connection);
|
||||
transportSession.addConnection(connection);
|
||||
});
|
||||
}
|
||||
|
||||
private static Flux<McpSchema.JSONRPCMessage> mcpSessionNotFoundError(String sessionRepresentation) {
|
||||
logger.warn("Session {} was not found on the MCP server", sessionRepresentation);
|
||||
// inform the stream/connection subscriber
|
||||
return Flux.error(new McpTransportSessionNotFoundException(sessionRepresentation));
|
||||
}
|
||||
|
||||
private Flux<McpSchema.JSONRPCMessage> extractError(ClientResponse response, String sessionRepresentation) {
|
||||
return response.<McpSchema.JSONRPCMessage>createError().onErrorResume(e -> {
|
||||
WebClientResponseException responseException = (WebClientResponseException) e;
|
||||
byte[] body = responseException.getResponseBodyAsByteArray();
|
||||
McpSchema.JSONRPCResponse.JSONRPCError jsonRpcError = null;
|
||||
Exception toPropagate;
|
||||
try {
|
||||
McpSchema.JSONRPCResponse jsonRpcResponse = objectMapper.readValue(body,
|
||||
McpSchema.JSONRPCResponse.class);
|
||||
jsonRpcError = jsonRpcResponse.error();
|
||||
toPropagate = new McpError(jsonRpcError);
|
||||
}
|
||||
catch (IOException ex) {
|
||||
toPropagate = new RuntimeException("Sending request failed", e);
|
||||
logger.debug("Received content together with {} HTTP code response: {}", response.statusCode(), body);
|
||||
}
|
||||
|
||||
// Some implementations can return 400 when presented with a
|
||||
// session id that it doesn't know about, so we will
|
||||
// invalidate the session
|
||||
// https://github.com/modelcontextprotocol/typescript-sdk/issues/389
|
||||
if (responseException.getStatusCode().isSameCodeAs(HttpStatus.BAD_REQUEST)) {
|
||||
return Mono.error(new McpTransportSessionNotFoundException(sessionRepresentation, toPropagate));
|
||||
}
|
||||
return Mono.empty();
|
||||
}).flux();
|
||||
}
|
||||
|
||||
private Flux<McpSchema.JSONRPCMessage> eventStream(McpTransportStream<Disposable> stream, ClientResponse response) {
|
||||
McpTransportStream<Disposable> sessionStream = stream != null ? stream
|
||||
: new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect);
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
|
||||
var idWithMessages = response.bodyToFlux(PARAMETERIZED_TYPE_REF).map(this::parse);
|
||||
return Flux.from(sessionStream.consumeSseStream(idWithMessages));
|
||||
}
|
||||
|
||||
private static boolean isNotFound(ClientResponse response) {
|
||||
return response.statusCode().isSameCodeAs(HttpStatus.NOT_FOUND);
|
||||
}
|
||||
|
||||
private static boolean isNotAllowed(ClientResponse response) {
|
||||
return response.statusCode().isSameCodeAs(HttpStatus.METHOD_NOT_ALLOWED);
|
||||
}
|
||||
|
||||
private static boolean isEventStream(ClientResponse response) {
|
||||
return response.statusCode().is2xxSuccessful() && response.headers().contentType().isPresent()
|
||||
&& response.headers().contentType().get().isCompatibleWith(MediaType.TEXT_EVENT_STREAM);
|
||||
}
|
||||
|
||||
private static String sessionIdOrPlaceholder(McpTransportSession<?> transportSession) {
|
||||
return transportSession.sessionId().orElse("[missing_session_id]");
|
||||
}
|
||||
|
||||
private Flux<McpSchema.JSONRPCMessage> responseFlux(ClientResponse response) {
|
||||
return response.bodyToMono(String.class).<Iterable<McpSchema.JSONRPCMessage>>handle((responseMessage, s) -> {
|
||||
try {
|
||||
McpSchema.JSONRPCMessage jsonRpcResponse = McpSchema.deserializeJsonRpcMessage(objectMapper,
|
||||
responseMessage);
|
||||
s.next(List.of(jsonRpcResponse));
|
||||
}
|
||||
catch (IOException e) {
|
||||
s.error(e);
|
||||
}
|
||||
}).flatMapIterable(Function.identity());
|
||||
}
|
||||
|
||||
private Flux<McpSchema.JSONRPCMessage> newEventStream(ClientResponse response, String sessionRepresentation) {
|
||||
McpTransportStream<Disposable> sessionStream = new DefaultMcpTransportStream<>(this.resumableStreams,
|
||||
this::reconnect);
|
||||
logger.trace("Sent POST and opened a stream ({}) for session {}", sessionStream.streamId(),
|
||||
sessionRepresentation);
|
||||
return eventStream(sessionStream, response);
|
||||
}
|
||||
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
|
||||
return this.objectMapper.convertValue(data, typeRef);
|
||||
}
|
||||
|
||||
private Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> parse(ServerSentEvent<String> event) {
|
||||
if (MESSAGE_EVENT_TYPE.equals(event.event())) {
|
||||
try {
|
||||
// We don't support batching ATM and probably won't since the next version
|
||||
// considers removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, event.data());
|
||||
return Tuples.of(Optional.ofNullable(event.id()), List.of(message));
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
throw new McpError("Error parsing JSON-RPC message: " + event.data());
|
||||
}
|
||||
}
|
||||
else {
|
||||
throw new McpError("Received unrecognized SSE event type: " + event.event());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Builder for {@link WebClientStreamableHttpTransport}.
|
||||
*/
|
||||
public static class Builder {
|
||||
|
||||
private ObjectMapper objectMapper;
|
||||
|
||||
private WebClient.Builder webClientBuilder;
|
||||
|
||||
private String endpoint = DEFAULT_ENDPOINT;
|
||||
|
||||
private boolean resumableStreams = true;
|
||||
|
||||
private boolean openConnectionOnStartup = false;
|
||||
|
||||
private Builder(WebClient.Builder webClientBuilder) {
|
||||
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
||||
this.webClientBuilder = webClientBuilder;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure the {@link ObjectMapper} to use.
|
||||
* @param objectMapper instance to use
|
||||
* @return the builder instance
|
||||
*/
|
||||
public Builder objectMapper(ObjectMapper objectMapper) {
|
||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||
this.objectMapper = objectMapper;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure the {@link WebClient.Builder} to construct the {@link WebClient}.
|
||||
* @param webClientBuilder instance to use
|
||||
* @return the builder instance
|
||||
*/
|
||||
public Builder webClientBuilder(WebClient.Builder webClientBuilder) {
|
||||
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
||||
this.webClientBuilder = webClientBuilder;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure the endpoint to make HTTP requests against.
|
||||
* @param endpoint endpoint to use
|
||||
* @return the builder instance
|
||||
*/
|
||||
public Builder endpoint(String endpoint) {
|
||||
Assert.hasText(endpoint, "endpoint must be a non-empty String");
|
||||
this.endpoint = endpoint;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure whether to use the stream resumability feature by keeping track of
|
||||
* SSE event ids.
|
||||
* @param resumableStreams if {@code true} event ids will be tracked and upon
|
||||
* disconnection, the last seen id will be used upon reconnection as a header to
|
||||
* resume consuming messages.
|
||||
* @return the builder instance
|
||||
*/
|
||||
public Builder resumableStreams(boolean resumableStreams) {
|
||||
this.resumableStreams = resumableStreams;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Configure whether the client should open an SSE connection upon startup. Not
|
||||
* all servers support this (although it is in theory possible with the current
|
||||
* specification), so use with caution. By default, this value is {@code false}.
|
||||
* @param openConnectionOnStartup if {@code true} the {@link #connect(Function)}
|
||||
* method call will try to open an SSE connection before sending any JSON-RPC
|
||||
* request
|
||||
* @return the builder instance
|
||||
*/
|
||||
public Builder openConnectionOnStartup(boolean openConnectionOnStartup) {
|
||||
this.openConnectionOnStartup = openConnectionOnStartup;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Construct a fresh instance of {@link WebClientStreamableHttpTransport} using
|
||||
* the current builder configuration.
|
||||
* @return a new instance of {@link WebClientStreamableHttpTransport}
|
||||
*/
|
||||
public WebClientStreamableHttpTransport build() {
|
||||
ObjectMapper objectMapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
||||
|
||||
return new WebClientStreamableHttpTransport(objectMapper, this.webClientBuilder, endpoint, resumableStreams,
|
||||
openConnectionOnStartup);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -190,6 +190,9 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
|||
*/
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
|
||||
// TODO: Avoid eager connection opening and enable resilience
|
||||
// -> upon disconnects, re-establish connection
|
||||
// -> allow optimizing for eager connection start using a constructor flag
|
||||
Flux<ServerSentEvent<String>> events = eventStream();
|
||||
this.inboundSubscription = events.concatMap(event -> Mono.just(event).<JSONRPCMessage>handle((e, s) -> {
|
||||
if (ENDPOINT_EVENT_TYPE.equals(event.event())) {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,17 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
@Timeout(15)
|
||||
public class WebClientStreamableHttpAsyncClientResiliencyTests extends AbstractMcpAsyncClientResiliencyTests {
|
||||
|
||||
@Override
|
||||
protected McpClientTransport createMcpTransport() {
|
||||
return WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(host)).build();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,42 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import org.testcontainers.images.builder.ImageFromDockerfile;
|
||||
|
||||
@Timeout(15)
|
||||
public class WebClientStreamableHttpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js streamableHttp")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected McpClientTransport createMcpTransport() {
|
||||
return WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(host)).build();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
@Timeout(15)
|
||||
public class WebClientStreamableHttpSyncClientTests extends AbstractMcpSyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js streamableHttp")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected McpClientTransport createMcpTransport() {
|
||||
return WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(host)).build();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -26,7 +26,8 @@ class WebFluxSseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
|||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js sse")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
|
|
|||
|
|
@ -26,7 +26,8 @@ class WebFluxSseMcpSyncClientTests extends AbstractMcpSyncClientTests {
|
|||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js sse")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
|
|
|||
|
|
@ -41,7 +41,8 @@ class WebFluxSseClientTransportTests {
|
|||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js sse")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
|
|
|||
|
|
@ -9,13 +9,13 @@
|
|||
</appender>
|
||||
|
||||
<!-- Main MCP package -->
|
||||
<logger name="org.springframework.ai.mcp" level="INFO"/>
|
||||
<logger name="io.modelcontextprotocol" level="INFO"/>
|
||||
|
||||
<!-- Client packages -->
|
||||
<logger name="org.springframework.ai.mcp.client" level="INFO"/>
|
||||
|
||||
<logger name="io.modelcontextprotocol.client" level="INFO"/>
|
||||
|
||||
<!-- Spec package -->
|
||||
<logger name="org.springframework.ai.mcp.spec" level="INFO"/>
|
||||
<logger name="io.modelcontextprotocol.spec" level="INFO"/>
|
||||
|
||||
|
||||
<!-- Root logger -->
|
||||
|
|
|
|||
|
|
@ -68,6 +68,11 @@
|
|||
<artifactId>junit-jupiter</artifactId>
|
||||
<version>${testcontainers.version}</version>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>toxiproxy</artifactId>
|
||||
<version>${toxiproxy.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,222 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import eu.rekawek.toxiproxy.Proxy;
|
||||
import eu.rekawek.toxiproxy.ToxiproxyClient;
|
||||
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.Test;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.Network;
|
||||
import org.testcontainers.containers.ToxiproxyContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.function.Function;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
|
||||
/**
|
||||
* Resiliency test suite for the {@link McpAsyncClient} that can be used with different
|
||||
* {@link McpTransport} implementations that support Streamable HTTP.
|
||||
*
|
||||
* The purpose of these tests is to allow validating the transport layer resiliency
|
||||
* instead of the functionality offered by the logical layer of MCP concepts such as
|
||||
* tools, resources, prompts, etc.
|
||||
*
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
public abstract class AbstractMcpAsyncClientResiliencyTests {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(AbstractMcpAsyncClientResiliencyTests.class);
|
||||
|
||||
static Network network = Network.newNetwork();
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
static GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js streamableHttp")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withNetwork(network)
|
||||
.withNetworkAliases("everything-server")
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
static ToxiproxyContainer toxiproxy = new ToxiproxyContainer("ghcr.io/shopify/toxiproxy:2.5.0").withNetwork(network)
|
||||
.withExposedPorts(8474, 3000);
|
||||
|
||||
static Proxy proxy;
|
||||
|
||||
static {
|
||||
container.start();
|
||||
|
||||
toxiproxy.start();
|
||||
|
||||
final ToxiproxyClient toxiproxyClient = new ToxiproxyClient(toxiproxy.getHost(), toxiproxy.getControlPort());
|
||||
try {
|
||||
proxy = toxiproxyClient.createProxy("everything-server", "0.0.0.0:3000", "everything-server:3001");
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException("Can't create proxy!", e);
|
||||
}
|
||||
|
||||
final String ipAddressViaToxiproxy = toxiproxy.getHost();
|
||||
final int portViaToxiproxy = toxiproxy.getMappedPort(3000);
|
||||
|
||||
host = "http://" + ipAddressViaToxiproxy + ":" + portViaToxiproxy;
|
||||
}
|
||||
|
||||
private static void disconnect() {
|
||||
long start = System.nanoTime();
|
||||
try {
|
||||
// disconnect
|
||||
// proxy.toxics().bandwidth("CUT_CONNECTION_DOWNSTREAM",
|
||||
// ToxicDirection.DOWNSTREAM, 0);
|
||||
// proxy.toxics().bandwidth("CUT_CONNECTION_UPSTREAM",
|
||||
// ToxicDirection.UPSTREAM, 0);
|
||||
proxy.toxics().resetPeer("RESET_DOWNSTREAM", ToxicDirection.DOWNSTREAM, 0);
|
||||
proxy.toxics().resetPeer("RESET_UPSTREAM", ToxicDirection.UPSTREAM, 0);
|
||||
logger.info("Disconnect took {} ms", Duration.ofNanos(System.nanoTime() - start).toMillis());
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException("Failed to disconnect", e);
|
||||
}
|
||||
}
|
||||
|
||||
private static void reconnect() {
|
||||
long start = System.nanoTime();
|
||||
try {
|
||||
proxy.toxics().get("RESET_UPSTREAM").remove();
|
||||
proxy.toxics().get("RESET_DOWNSTREAM").remove();
|
||||
// proxy.toxics().get("CUT_CONNECTION_DOWNSTREAM").remove();
|
||||
// proxy.toxics().get("CUT_CONNECTION_UPSTREAM").remove();
|
||||
logger.info("Reconnect took {} ms", Duration.ofNanos(System.nanoTime() - start).toMillis());
|
||||
}
|
||||
catch (IOException e) {
|
||||
throw new RuntimeException("Failed to reconnect", e);
|
||||
}
|
||||
}
|
||||
|
||||
private static void restartMcpServer() {
|
||||
container.stop();
|
||||
container.start();
|
||||
}
|
||||
|
||||
abstract McpClientTransport createMcpTransport();
|
||||
|
||||
protected Duration getRequestTimeout() {
|
||||
return Duration.ofSeconds(14);
|
||||
}
|
||||
|
||||
protected Duration getInitializationTimeout() {
|
||||
return Duration.ofSeconds(2);
|
||||
}
|
||||
|
||||
McpAsyncClient client(McpClientTransport transport) {
|
||||
return client(transport, Function.identity());
|
||||
}
|
||||
|
||||
McpAsyncClient client(McpClientTransport transport, Function<McpClient.AsyncSpec, McpClient.AsyncSpec> customizer) {
|
||||
AtomicReference<McpAsyncClient> client = new AtomicReference<>();
|
||||
|
||||
assertThatCode(() -> {
|
||||
McpClient.AsyncSpec builder = McpClient.async(transport)
|
||||
.requestTimeout(getRequestTimeout())
|
||||
.initializationTimeout(getInitializationTimeout())
|
||||
.capabilities(McpSchema.ClientCapabilities.builder().roots(true).build());
|
||||
builder = customizer.apply(builder);
|
||||
client.set(builder.build());
|
||||
}).doesNotThrowAnyException();
|
||||
|
||||
return client.get();
|
||||
}
|
||||
|
||||
void withClient(McpClientTransport transport, Consumer<McpAsyncClient> c) {
|
||||
withClient(transport, Function.identity(), c);
|
||||
}
|
||||
|
||||
void withClient(McpClientTransport transport, Function<McpClient.AsyncSpec, McpClient.AsyncSpec> customizer,
|
||||
Consumer<McpAsyncClient> c) {
|
||||
var client = client(transport, customizer);
|
||||
try {
|
||||
c.accept(client);
|
||||
}
|
||||
finally {
|
||||
StepVerifier.create(client.closeGracefully()).expectComplete().verify(Duration.ofSeconds(10));
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPing() {
|
||||
withClient(createMcpTransport(), mcpAsyncClient -> {
|
||||
StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
|
||||
|
||||
disconnect();
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.ping()).expectError().verify();
|
||||
|
||||
reconnect();
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.ping()).expectNextCount(1).verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSessionInvalidation() {
|
||||
withClient(createMcpTransport(), mcpAsyncClient -> {
|
||||
StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
|
||||
|
||||
restartMcpServer();
|
||||
|
||||
// The first try will face the session mismatch exception and the second one
|
||||
// will go through the re-initialization process.
|
||||
StepVerifier.create(mcpAsyncClient.ping().retry(1)).expectNextCount(1).verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallTool() {
|
||||
withClient(createMcpTransport(), mcpAsyncClient -> {
|
||||
AtomicReference<List<McpSchema.Tool>> tools = new AtomicReference<>();
|
||||
StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
|
||||
StepVerifier.create(mcpAsyncClient.listTools())
|
||||
.consumeNextWith(list -> tools.set(list.tools()))
|
||||
.verifyComplete();
|
||||
|
||||
disconnect();
|
||||
|
||||
String name = tools.get().get(0).name();
|
||||
// Assuming this is the echo tool
|
||||
McpSchema.CallToolRequest request = new McpSchema.CallToolRequest(name, Map.of("message", "hello"));
|
||||
StepVerifier.create(mcpAsyncClient.callTool(request)).expectError().verify();
|
||||
|
||||
reconnect();
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.callTool(request)).expectNextCount(1).verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
@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();
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -110,14 +110,16 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
onClose();
|
||||
}
|
||||
|
||||
<T> void verifyInitializationTimeout(Function<McpAsyncClient, Mono<T>> operation, String action) {
|
||||
<T> void verifyNotificationSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation,
|
||||
String action) {
|
||||
withClient(createMcpTransport(), mcpAsyncClient -> {
|
||||
StepVerifier.withVirtualTime(() -> operation.apply(mcpAsyncClient))
|
||||
.expectSubscription()
|
||||
.thenAwait(getInitializationTimeout())
|
||||
.consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before " + action))
|
||||
.verify();
|
||||
StepVerifier.create(operation.apply(mcpAsyncClient)).verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
<T> void verifyCallSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation, String action) {
|
||||
withClient(createMcpTransport(), mcpAsyncClient -> {
|
||||
StepVerifier.create(operation.apply(mcpAsyncClient)).expectNextCount(1).verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -133,7 +135,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListToolsWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.listTools(null), "listing tools");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listTools(null), "listing tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -153,7 +155,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testPingWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.ping(), "pinging the server");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.ping(), "pinging the server");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -168,7 +170,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
@Test
|
||||
void testCallToolWithoutInitialization() {
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", ECHO_TEST_MESSAGE));
|
||||
verifyInitializationTimeout(client -> client.callTool(callToolRequest), "calling tools");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.callTool(callToolRequest), "calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -202,7 +204,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListResourcesWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.listResources(null), "listing resources");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(null), "listing resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -233,7 +235,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListPromptsWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.listPrompts(null), "listing " + "prompts");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listPrompts(null), "listing " + "prompts");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -258,7 +260,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
@Test
|
||||
void testGetPromptWithoutInitialization() {
|
||||
GetPromptRequest request = new GetPromptRequest("simple_prompt", Map.of());
|
||||
verifyInitializationTimeout(client -> client.getPrompt(request), "getting " + "prompts");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.getPrompt(request), "getting " + "prompts");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -279,7 +281,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testRootsListChangedWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.rootsListChangedNotification(),
|
||||
verifyNotificationSucceedsWithImplicitInitialization(client -> client.rootsListChangedNotification(),
|
||||
"sending roots list changed notification");
|
||||
}
|
||||
|
||||
|
|
@ -354,7 +356,8 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListResourceTemplatesWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.listResourceTemplates(), "listing resource templates");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listResourceTemplates(),
|
||||
"listing resource templates");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -447,8 +450,8 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testLoggingLevelsWithoutInitialization() {
|
||||
verifyInitializationTimeout(client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG),
|
||||
"setting logging level");
|
||||
verifyNotificationSucceedsWithImplicitInitialization(
|
||||
client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG), "setting logging level");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
|
@ -12,7 +13,6 @@ import java.util.function.Consumer;
|
|||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
|
||||
|
|
@ -112,33 +112,18 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
static final Object DUMMY_RETURN_VALUE = new Object();
|
||||
|
||||
<T> void verifyNotificationTimesOut(Consumer<McpSyncClient> operation, String action) {
|
||||
verifyCallTimesOut(client -> {
|
||||
<T> void verifyNotificationSucceedsWithImplicitInitialization(Consumer<McpSyncClient> operation, String action) {
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> {
|
||||
operation.accept(client);
|
||||
return DUMMY_RETURN_VALUE;
|
||||
}, action);
|
||||
}
|
||||
|
||||
<T> void verifyCallTimesOut(Function<McpSyncClient, T> blockingOperation, String action) {
|
||||
<T> void verifyCallSucceedsWithImplicitInitialization(Function<McpSyncClient, T> blockingOperation, String action) {
|
||||
withClient(createMcpTransport(), mcpSyncClient -> {
|
||||
// This scheduler is not replaced by virtual time scheduler
|
||||
Scheduler customScheduler = Schedulers.newBoundedElastic(1, 1, "actualBoundedElastic");
|
||||
|
||||
StepVerifier.withVirtualTime(() -> Mono.fromSupplier(() -> blockingOperation.apply(mcpSyncClient))
|
||||
// Offload the blocking call to the real scheduler
|
||||
.subscribeOn(customScheduler))
|
||||
.expectSubscription()
|
||||
// This works without actually waiting but executes all the
|
||||
// tasks pending execution on the VirtualTimeScheduler.
|
||||
// It is possible to execute the blocking code from the operation
|
||||
// because it is blocked on a dedicated Scheduler and the main
|
||||
// flow is not blocked and uses the VirtualTimeScheduler.
|
||||
.thenAwait(getInitializationTimeout())
|
||||
.consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before " + action))
|
||||
.verify();
|
||||
|
||||
customScheduler.dispose();
|
||||
StepVerifier.create(Mono.fromSupplier(() -> blockingOperation.apply(mcpSyncClient)))
|
||||
.expectNextCount(1)
|
||||
.verifyComplete();
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -154,7 +139,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListToolsWithoutInitialization() {
|
||||
verifyCallTimesOut(client -> client.listTools(null), "listing tools");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listTools(null), "listing tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -175,8 +160,8 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testCallToolsWithoutInitialization() {
|
||||
verifyCallTimesOut(client -> client.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))),
|
||||
"calling tools");
|
||||
verifyCallSucceedsWithImplicitInitialization(
|
||||
client -> client.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))), "calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -200,7 +185,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testPingWithoutInitialization() {
|
||||
verifyCallTimesOut(client -> client.ping(), "pinging the server");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.ping(), "pinging the server");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -214,7 +199,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
@Test
|
||||
void testCallToolWithoutInitialization() {
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
|
||||
verifyCallTimesOut(client -> client.callTool(callToolRequest), "calling tools");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.callTool(callToolRequest), "calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -243,7 +228,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testRootsListChangedWithoutInitialization() {
|
||||
verifyNotificationTimesOut(client -> client.rootsListChangedNotification(),
|
||||
verifyNotificationSucceedsWithImplicitInitialization(client -> client.rootsListChangedNotification(),
|
||||
"sending roots list changed notification");
|
||||
}
|
||||
|
||||
|
|
@ -257,7 +242,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListResourcesWithoutInitialization() {
|
||||
verifyCallTimesOut(client -> client.listResources(null), "listing resources");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(null), "listing resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -333,8 +318,14 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testReadResourceWithoutInitialization() {
|
||||
Resource resource = new Resource("test://uri", "Test Resource", null, null, null);
|
||||
verifyCallTimesOut(client -> client.readResource(resource), "reading resources");
|
||||
AtomicReference<List<Resource>> resources = new AtomicReference<>();
|
||||
withClient(createMcpTransport(), mcpSyncClient -> {
|
||||
mcpSyncClient.initialize();
|
||||
resources.set(mcpSyncClient.listResources().resources());
|
||||
});
|
||||
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.readResource(resources.get().get(0)),
|
||||
"reading resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -355,7 +346,8 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testListResourceTemplatesWithoutInitialization() {
|
||||
verifyCallTimesOut(client -> client.listResourceTemplates(null), "listing resource templates");
|
||||
verifyCallSucceedsWithImplicitInitialization(client -> client.listResourceTemplates(null),
|
||||
"listing resource templates");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -413,8 +405,8 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testLoggingLevelsWithoutInitialization() {
|
||||
verifyNotificationTimesOut(client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG),
|
||||
"setting logging level");
|
||||
verifyNotificationSucceedsWithImplicitInitialization(
|
||||
client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG), "setting logging level");
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
|
|||
|
|
@ -9,9 +9,9 @@ import java.util.HashMap;
|
|||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.TimeoutException;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
|
|
@ -32,7 +32,7 @@ import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
|
|||
import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
|
||||
import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Root;
|
||||
import io.modelcontextprotocol.spec.McpTransport;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import io.modelcontextprotocol.util.Utils;
|
||||
import org.slf4j.Logger;
|
||||
|
|
@ -77,29 +77,37 @@ import reactor.core.publisher.Sinks;
|
|||
* @see McpClient
|
||||
* @see McpSchema
|
||||
* @see McpClientSession
|
||||
* @see McpClientTransport
|
||||
*/
|
||||
public class McpAsyncClient {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(McpAsyncClient.class);
|
||||
|
||||
private static TypeReference<Void> VOID_TYPE_REFERENCE = new TypeReference<>() {
|
||||
private static final TypeReference<Void> VOID_TYPE_REFERENCE = new TypeReference<>() {
|
||||
};
|
||||
|
||||
protected final Sinks.One<McpSchema.InitializeResult> initializedSink = Sinks.one();
|
||||
public static final TypeReference<Object> OBJECT_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
private AtomicBoolean initialized = new AtomicBoolean(false);
|
||||
public static final TypeReference<PaginatedRequest> PAGINATED_REQUEST_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
public static final TypeReference<McpSchema.InitializeResult> INITIALIZE_RESULT_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
public static final TypeReference<CreateMessageRequest> CREATE_MESSAGE_REQUEST_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
public static final TypeReference<LoggingMessageNotification> LOGGING_MESSAGE_NOTIFICATION_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
private final AtomicReference<Initialization> initializationRef = new AtomicReference<>();
|
||||
|
||||
/**
|
||||
* The max timeout to await for the client-server connection to be initialized.
|
||||
*/
|
||||
private final Duration initializationTimeout;
|
||||
|
||||
/**
|
||||
* The MCP session implementation that manages bidirectional JSON-RPC communication
|
||||
* between clients and servers.
|
||||
*/
|
||||
private final McpClientSession mcpSession;
|
||||
|
||||
/**
|
||||
* Client capabilities.
|
||||
*/
|
||||
|
|
@ -110,21 +118,6 @@ public class McpAsyncClient {
|
|||
*/
|
||||
private final McpSchema.Implementation clientInfo;
|
||||
|
||||
/**
|
||||
* Server capabilities.
|
||||
*/
|
||||
private McpSchema.ServerCapabilities serverCapabilities;
|
||||
|
||||
/**
|
||||
* Server instructions.
|
||||
*/
|
||||
private String serverInstructions;
|
||||
|
||||
/**
|
||||
* Server implementation information.
|
||||
*/
|
||||
private McpSchema.Implementation serverInfo;
|
||||
|
||||
/**
|
||||
* Roots define the boundaries of where servers can operate within the filesystem,
|
||||
* allowing them to understand which directories and files they have access to.
|
||||
|
|
@ -155,13 +148,19 @@ public class McpAsyncClient {
|
|||
/**
|
||||
* Client transport implementation.
|
||||
*/
|
||||
private final McpTransport transport;
|
||||
private final McpClientTransport transport;
|
||||
|
||||
/**
|
||||
* Supported protocol versions.
|
||||
*/
|
||||
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
|
||||
|
||||
/**
|
||||
* The MCP session supplier that manages bidirectional JSON-RPC communication between
|
||||
* clients and servers.
|
||||
*/
|
||||
private final Supplier<McpClientSession> sessionSupplier;
|
||||
|
||||
/**
|
||||
* Create a new McpAsyncClient with the given transport and session request-response
|
||||
* timeout.
|
||||
|
|
@ -254,16 +253,38 @@ public class McpAsyncClient {
|
|||
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_MESSAGE,
|
||||
asyncLoggingNotificationHandler(loggingConsumersFinal));
|
||||
|
||||
this.mcpSession = new McpClientSession(requestTimeout, transport, requestHandlers, notificationHandlers);
|
||||
this.transport.setExceptionHandler(this::handleException);
|
||||
this.sessionSupplier = () -> new McpClientSession(requestTimeout, transport, requestHandlers,
|
||||
notificationHandlers);
|
||||
|
||||
}
|
||||
|
||||
private void handleException(Throwable t) {
|
||||
logger.warn("Handling exception", t);
|
||||
if (t instanceof McpTransportSessionNotFoundException) {
|
||||
Initialization previous = this.initializationRef.getAndSet(null);
|
||||
if (previous != null) {
|
||||
previous.close();
|
||||
}
|
||||
// Providing an empty operation since we are only interested in triggering the
|
||||
// implicit initialization step.
|
||||
withSession("re-initializing", result -> Mono.empty()).subscribe();
|
||||
}
|
||||
}
|
||||
|
||||
private McpSchema.InitializeResult currentInitializationResult() {
|
||||
Initialization current = this.initializationRef.get();
|
||||
McpSchema.InitializeResult initializeResult = current != null ? current.result.get() : null;
|
||||
return initializeResult;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the server capabilities that define the supported features and functionality.
|
||||
* @return The server capabilities
|
||||
*/
|
||||
public McpSchema.ServerCapabilities getServerCapabilities() {
|
||||
return this.serverCapabilities;
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.capabilities() : null;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -272,7 +293,8 @@ public class McpAsyncClient {
|
|||
* @return The server instructions
|
||||
*/
|
||||
public String getServerInstructions() {
|
||||
return this.serverInstructions;
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.instructions() : null;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -280,7 +302,8 @@ public class McpAsyncClient {
|
|||
* @return The server implementation details
|
||||
*/
|
||||
public McpSchema.Implementation getServerInfo() {
|
||||
return this.serverInfo;
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.serverInfo() : null;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -288,7 +311,8 @@ public class McpAsyncClient {
|
|||
* @return true if the client-server connection is initialized
|
||||
*/
|
||||
public boolean isInitialized() {
|
||||
return this.initialized.get();
|
||||
Initialization current = this.initializationRef.get();
|
||||
return current != null && (current.result.get() != null);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -311,7 +335,11 @@ public class McpAsyncClient {
|
|||
* Closes the client connection immediately.
|
||||
*/
|
||||
public void close() {
|
||||
this.mcpSession.close();
|
||||
Initialization current = this.initializationRef.getAndSet(null);
|
||||
if (current != null) {
|
||||
current.close();
|
||||
}
|
||||
this.transport.close();
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -319,14 +347,21 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes when the connection is closed
|
||||
*/
|
||||
public Mono<Void> closeGracefully() {
|
||||
return this.mcpSession.closeGracefully();
|
||||
return Mono.defer(() -> {
|
||||
Initialization current = this.initializationRef.getAndSet(null);
|
||||
Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
|
||||
return sessionClose.then(transport.closeGracefully());
|
||||
});
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
// Initialization
|
||||
// --------------------------
|
||||
/**
|
||||
* The initialization phase MUST be the first interaction between client and server.
|
||||
* The initialization phase should be the first interaction between client and server.
|
||||
* The client will ensure it happens in case it has not been explicitly called and in
|
||||
* case of transport session invalidation.
|
||||
* <p>
|
||||
* During this phase, the client and server:
|
||||
* <ul>
|
||||
* <li>Establish protocol version compatibility</li>
|
||||
|
|
@ -346,9 +381,13 @@ public class McpAsyncClient {
|
|||
* @see <a href=
|
||||
* "https://github.com/modelcontextprotocol/specification/blob/main/docs/specification/basic/lifecycle.md#initialization">MCP
|
||||
* Initialization Spec</a>
|
||||
* </p>
|
||||
*/
|
||||
public Mono<McpSchema.InitializeResult> initialize() {
|
||||
return withSession("by explicit API call", init -> Mono.just(init.get()));
|
||||
}
|
||||
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(McpClientSession mcpClientSession) {
|
||||
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
|
||||
|
||||
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
|
||||
|
|
@ -356,16 +395,10 @@ public class McpAsyncClient {
|
|||
this.clientCapabilities,
|
||||
this.clientInfo); // @formatter:on
|
||||
|
||||
Mono<McpSchema.InitializeResult> result = this.mcpSession.sendRequest(McpSchema.METHOD_INITIALIZE,
|
||||
initializeRequest, new TypeReference<McpSchema.InitializeResult>() {
|
||||
});
|
||||
Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
|
||||
initializeRequest, INITIALIZE_RESULT_TYPE_REF);
|
||||
|
||||
return result.flatMap(initializeResult -> {
|
||||
|
||||
this.serverCapabilities = initializeResult.capabilities();
|
||||
this.serverInstructions = initializeResult.instructions();
|
||||
this.serverInfo = initializeResult.serverInfo();
|
||||
|
||||
logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
|
||||
initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
|
||||
initializeResult.instructions());
|
||||
|
|
@ -375,28 +408,93 @@ public class McpAsyncClient {
|
|||
"Unsupported protocol version from the server: " + initializeResult.protocolVersion()));
|
||||
}
|
||||
|
||||
return this.mcpSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null).doOnSuccess(v -> {
|
||||
this.initialized.set(true);
|
||||
this.initializedSink.tryEmitValue(initializeResult);
|
||||
}).thenReturn(initializeResult);
|
||||
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
|
||||
.thenReturn(initializeResult);
|
||||
});
|
||||
}
|
||||
|
||||
private static class Initialization {
|
||||
|
||||
private final Sinks.One<McpSchema.InitializeResult> initSink = Sinks.one();
|
||||
|
||||
private final AtomicReference<McpSchema.InitializeResult> result = new AtomicReference<>();
|
||||
|
||||
private final AtomicReference<McpClientSession> mcpClientSession = new AtomicReference<>();
|
||||
|
||||
static Initialization create() {
|
||||
return new Initialization();
|
||||
}
|
||||
|
||||
void setMcpClientSession(McpClientSession mcpClientSession) {
|
||||
this.mcpClientSession.set(mcpClientSession);
|
||||
}
|
||||
|
||||
McpClientSession mcpSession() {
|
||||
return this.mcpClientSession.get();
|
||||
}
|
||||
|
||||
McpSchema.InitializeResult get() {
|
||||
return this.result.get();
|
||||
}
|
||||
|
||||
Mono<McpSchema.InitializeResult> await() {
|
||||
return this.initSink.asMono();
|
||||
}
|
||||
|
||||
void complete(McpSchema.InitializeResult initializeResult) {
|
||||
// first ensure the result is cached
|
||||
this.result.set(initializeResult);
|
||||
// inform all the subscribers waiting for the initialization
|
||||
this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
void error(Throwable t) {
|
||||
this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
void close() {
|
||||
this.mcpSession().close();
|
||||
}
|
||||
|
||||
Mono<Void> closeGracefully() {
|
||||
return this.mcpSession().closeGracefully();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility method to handle the common pattern of checking initialization before
|
||||
* Utility method to handle the common pattern of ensuring initialization before
|
||||
* executing an operation.
|
||||
* @param <T> The type of the result Mono
|
||||
* @param actionName The action to perform if the client is initialized
|
||||
* @param operation The operation to execute if the client is initialized
|
||||
* @param actionName The action to perform when the client is initialized
|
||||
* @param operation The operation to execute when the client is initialized
|
||||
* @return A Mono that completes with the result of the operation
|
||||
*/
|
||||
private <T> Mono<T> withInitializationCheck(String actionName,
|
||||
Function<McpSchema.InitializeResult, Mono<T>> operation) {
|
||||
return this.initializedSink.asMono()
|
||||
.timeout(this.initializationTimeout)
|
||||
.onErrorResume(TimeoutException.class,
|
||||
ex -> Mono.error(new McpError("Client must be initialized before " + actionName)))
|
||||
.flatMap(operation);
|
||||
private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) {
|
||||
return Mono.defer(() -> {
|
||||
Initialization newInit = Initialization.create();
|
||||
Initialization previous = this.initializationRef.compareAndExchange(null, newInit);
|
||||
|
||||
boolean needsToInitialize = previous == null;
|
||||
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
|
||||
if (needsToInitialize) {
|
||||
newInit.setMcpClientSession(this.sessionSupplier.get());
|
||||
}
|
||||
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
|
||||
? doInitialize(newInit.mcpSession()).doOnNext(newInit::complete).onErrorResume(ex -> {
|
||||
newInit.error(ex);
|
||||
return Mono.error(ex);
|
||||
}) : previous.await();
|
||||
|
||||
return initializationJob.map(initializeResult -> this.initializationRef.get())
|
||||
.timeout(this.initializationTimeout)
|
||||
.onErrorResume(ex -> {
|
||||
logger.warn("Failed to initialize", ex);
|
||||
return Mono.error(new McpError("Client failed to initialize " + actionName));
|
||||
})
|
||||
.flatMap(operation);
|
||||
});
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
|
|
@ -408,9 +506,8 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes with the server's ping response
|
||||
*/
|
||||
public Mono<Object> ping() {
|
||||
return this.withInitializationCheck("pinging the server", initializedResult -> this.mcpSession
|
||||
.sendRequest(McpSchema.METHOD_PING, null, new TypeReference<Object>() {
|
||||
}));
|
||||
return this.withSession("pinging the server",
|
||||
init -> init.mcpSession().sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF));
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
|
|
@ -490,16 +587,14 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes when the notification is sent.
|
||||
*/
|
||||
public Mono<Void> rootsListChangedNotification() {
|
||||
return this.withInitializationCheck("sending roots list changed notification",
|
||||
initResult -> this.mcpSession.sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
|
||||
return this.withSession("sending roots list changed notification",
|
||||
init -> init.mcpSession().sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
|
||||
}
|
||||
|
||||
private RequestHandler<McpSchema.ListRootsResult> rootsListRequestHandler() {
|
||||
return params -> {
|
||||
@SuppressWarnings("unused")
|
||||
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
|
||||
new TypeReference<McpSchema.PaginatedRequest>() {
|
||||
});
|
||||
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, PAGINATED_REQUEST_TYPE_REF);
|
||||
|
||||
List<Root> roots = this.roots.values().stream().toList();
|
||||
|
||||
|
|
@ -512,9 +607,7 @@ public class McpAsyncClient {
|
|||
// --------------------------
|
||||
private RequestHandler<CreateMessageResult> samplingCreateMessageHandler() {
|
||||
return params -> {
|
||||
McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params,
|
||||
new TypeReference<McpSchema.CreateMessageRequest>() {
|
||||
});
|
||||
McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params, CREATE_MESSAGE_REQUEST_TYPE_REF);
|
||||
|
||||
return this.samplingHandler.apply(request);
|
||||
};
|
||||
|
|
@ -553,11 +646,12 @@ public class McpAsyncClient {
|
|||
* @see #listTools()
|
||||
*/
|
||||
public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
|
||||
return this.withInitializationCheck("calling tools", initializedResult -> {
|
||||
if (this.serverCapabilities.tools() == null) {
|
||||
return this.withSession("calling tools", init -> {
|
||||
if (init.get().capabilities().tools() == null) {
|
||||
return Mono.error(new McpError("Server does not provide tools capability"));
|
||||
}
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -575,12 +669,13 @@ public class McpAsyncClient {
|
|||
* @return A Mono that emits the list of tools result
|
||||
*/
|
||||
public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
|
||||
return this.withInitializationCheck("listing tools", initializedResult -> {
|
||||
if (this.serverCapabilities.tools() == null) {
|
||||
return this.withSession("listing tools", init -> {
|
||||
if (init.get().capabilities().tools() == null) {
|
||||
return Mono.error(new McpError("Server does not provide tools capability"));
|
||||
}
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_TOOLS_RESULT_TYPE_REF);
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_TOOLS_RESULT_TYPE_REF);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -632,12 +727,13 @@ public class McpAsyncClient {
|
|||
* @see #readResource(McpSchema.Resource)
|
||||
*/
|
||||
public Mono<McpSchema.ListResourcesResult> listResources(String cursor) {
|
||||
return this.withInitializationCheck("listing resources", initializedResult -> {
|
||||
if (this.serverCapabilities.resources() == null) {
|
||||
return this.withSession("listing resources", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_RESOURCES_RESULT_TYPE_REF);
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_RESOURCES_RESULT_TYPE_REF);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -663,12 +759,12 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ReadResourceResult
|
||||
*/
|
||||
public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.ReadResourceRequest readResourceRequest) {
|
||||
return this.withInitializationCheck("reading resources", initializedResult -> {
|
||||
if (this.serverCapabilities.resources() == null) {
|
||||
return this.withSession("reading resources", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_READ, readResourceRequest,
|
||||
READ_RESOURCE_RESULT_TYPE_REF);
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_READ, readResourceRequest, READ_RESOURCE_RESULT_TYPE_REF);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -692,12 +788,13 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ListResourceTemplatesResult
|
||||
*/
|
||||
public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates(String cursor) {
|
||||
return this.withInitializationCheck("listing resource templates", initializedResult -> {
|
||||
if (this.serverCapabilities.resources() == null) {
|
||||
return this.withSession("listing resource templates", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST,
|
||||
new McpSchema.PaginatedRequest(cursor), LIST_RESOURCE_TEMPLATES_RESULT_TYPE_REF);
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_RESOURCE_TEMPLATES_RESULT_TYPE_REF);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -711,7 +808,7 @@ public class McpAsyncClient {
|
|||
* @see #unsubscribeResource(McpSchema.UnsubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
|
||||
return this.withInitializationCheck("subscribing to resources", initializedResult -> this.mcpSession
|
||||
return this.withSession("subscribing to resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -725,7 +822,7 @@ public class McpAsyncClient {
|
|||
* @see #subscribeResource(McpSchema.SubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
|
||||
return this.withInitializationCheck("unsubscribing from resources", initializedResult -> this.mcpSession
|
||||
return this.withSession("unsubscribing from resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -767,7 +864,7 @@ public class McpAsyncClient {
|
|||
* @see #getPrompt(GetPromptRequest)
|
||||
*/
|
||||
public Mono<ListPromptsResult> listPrompts(String cursor) {
|
||||
return this.withInitializationCheck("listing prompts", initializedResult -> this.mcpSession
|
||||
return this.withSession("listing prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -781,7 +878,7 @@ public class McpAsyncClient {
|
|||
* @see #listPrompts()
|
||||
*/
|
||||
public Mono<GetPromptResult> getPrompt(GetPromptRequest getPromptRequest) {
|
||||
return this.withInitializationCheck("getting prompts", initializedResult -> this.mcpSession
|
||||
return this.withSession("getting prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -812,8 +909,7 @@ public class McpAsyncClient {
|
|||
|
||||
return params -> {
|
||||
McpSchema.LoggingMessageNotification loggingMessageNotification = transport.unmarshalFrom(params,
|
||||
new TypeReference<McpSchema.LoggingMessageNotification>() {
|
||||
});
|
||||
LOGGING_MESSAGE_NOTIFICATION_TYPE_REF);
|
||||
|
||||
return Flux.fromIterable(loggingConsumers)
|
||||
.flatMap(consumer -> consumer.apply(loggingMessageNotification))
|
||||
|
|
@ -833,10 +929,9 @@ public class McpAsyncClient {
|
|||
return Mono.error(new McpError("Logging level must not be null"));
|
||||
}
|
||||
|
||||
return this.withInitializationCheck("setting logging level", initializedResult -> {
|
||||
return this.withSession("setting logging level", init -> {
|
||||
var params = new McpSchema.SetLevelRequest(loggingLevel);
|
||||
return this.mcpSession.sendRequest(McpSchema.METHOD_LOGGING_SET_LEVEL, params, new TypeReference<Object>() {
|
||||
}).then();
|
||||
return init.mcpSession().sendRequest(McpSchema.METHOD_LOGGING_SET_LEVEL, params, OBJECT_TYPE_REF).then();
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -866,7 +961,7 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.CompleteResult
|
||||
*/
|
||||
public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
|
||||
return this.withInitializationCheck("complete completions", initializedResult -> this.mcpSession
|
||||
return this.withSession("complete completions", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -112,6 +112,7 @@ public class StdioClientTransport implements McpClientTransport {
|
|||
@Override
|
||||
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
|
||||
return Mono.<Void>fromRunnable(() -> {
|
||||
logger.info("MCP server starting.");
|
||||
handleIncomingMessages(handler);
|
||||
handleIncomingErrors();
|
||||
|
||||
|
|
@ -142,6 +143,7 @@ public class StdioClientTransport implements McpClientTransport {
|
|||
startInboundProcessing();
|
||||
startOutboundProcessing();
|
||||
startErrorProcessing();
|
||||
logger.info("MCP server started");
|
||||
}).subscribeOn(Schedulers.boundedElastic());
|
||||
}
|
||||
|
||||
|
|
@ -366,6 +368,9 @@ public class StdioClientTransport implements McpClientTransport {
|
|||
if (process.exitValue() != 0) {
|
||||
logger.warn("Process terminated with code " + process.exitValue());
|
||||
}
|
||||
else {
|
||||
logger.info("MCP server process stopped");
|
||||
}
|
||||
}).then(Mono.fromRunnable(() -> {
|
||||
try {
|
||||
// The Threads are blocked on readLine so disposeGracefully would not
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
package io.modelcontextprotocol.spec;
|
||||
|
||||
import org.reactivestreams.Publisher;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.Disposables;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
/**
|
||||
* Default implementation of {@link McpTransportSession} which manages the open
|
||||
* connections using tye {@link Disposable} type and allows to perform clean up using the
|
||||
* {@link Disposable#dispose()} method.
|
||||
*
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
public class DefaultMcpTransportSession implements McpTransportSession<Disposable> {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(DefaultMcpTransportSession.class);
|
||||
|
||||
private final Disposable.Composite openConnections = Disposables.composite();
|
||||
|
||||
private final AtomicBoolean initialized = new AtomicBoolean(false);
|
||||
|
||||
private final AtomicReference<String> sessionId = new AtomicReference<>();
|
||||
|
||||
private final Supplier<Publisher<Void>> onClose;
|
||||
|
||||
public DefaultMcpTransportSession(Supplier<Publisher<Void>> onClose) {
|
||||
this.onClose = onClose;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<String> sessionId() {
|
||||
return Optional.ofNullable(this.sessionId.get());
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean markInitialized(String sessionId) {
|
||||
boolean flipped = this.initialized.compareAndSet(false, true);
|
||||
if (flipped) {
|
||||
this.sessionId.set(sessionId);
|
||||
logger.debug("Established session with id {}", sessionId);
|
||||
}
|
||||
else {
|
||||
if (sessionId != null && !sessionId.equals(this.sessionId.get())) {
|
||||
logger.warn("Different session id provided in response. Expecting {} but server returned {}",
|
||||
this.sessionId.get(), sessionId);
|
||||
}
|
||||
}
|
||||
return flipped;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void addConnection(Disposable connection) {
|
||||
this.openConnections.add(connection);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void removeConnection(Disposable connection) {
|
||||
this.openConnections.remove(connection);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
this.closeGracefully().subscribe();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.from(this.onClose.get()).then(Mono.fromRunnable(this.openConnections::dispose));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,74 @@
|
|||
package io.modelcontextprotocol.spec;
|
||||
|
||||
import org.reactivestreams.Publisher;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.util.function.Tuple2;
|
||||
|
||||
import java.util.Optional;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
/**
|
||||
* An implementation of {@link McpTransportStream} using Project Reactor types.
|
||||
*
|
||||
* @param <CONNECTION> the resource serving the stream
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
public class DefaultMcpTransportStream<CONNECTION> implements McpTransportStream<CONNECTION> {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(DefaultMcpTransportStream.class);
|
||||
|
||||
private static final AtomicLong counter = new AtomicLong();
|
||||
|
||||
private final AtomicReference<String> lastId = new AtomicReference<>();
|
||||
|
||||
// Used only for internal accounting
|
||||
private final long streamId;
|
||||
|
||||
private final boolean resumable;
|
||||
|
||||
private final Function<McpTransportStream<CONNECTION>, Publisher<CONNECTION>> reconnect;
|
||||
|
||||
/**
|
||||
* Constructs a new instance representing a particular stream that can resume using
|
||||
* the provided reconnect mechanism.
|
||||
* @param resumable whether the stream is resumable and should try to reconnect
|
||||
* @param reconnect the mechanism to use in case an error is observed on the current
|
||||
* event stream to asynchronously kick off a resumed stream consumption, potentially
|
||||
* using the stored {@link #lastId()}.
|
||||
*/
|
||||
public DefaultMcpTransportStream(boolean resumable,
|
||||
Function<McpTransportStream<CONNECTION>, Publisher<CONNECTION>> reconnect) {
|
||||
this.reconnect = reconnect;
|
||||
this.streamId = counter.getAndIncrement();
|
||||
this.resumable = resumable;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Optional<String> lastId() {
|
||||
return Optional.ofNullable(this.lastId.get());
|
||||
}
|
||||
|
||||
@Override
|
||||
public long streamId() {
|
||||
return this.streamId;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Publisher<McpSchema.JSONRPCMessage> consumeSseStream(
|
||||
Publisher<Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>>> eventStream) {
|
||||
return Flux.deferContextual(ctx -> Flux.from(eventStream).doOnError(e -> {
|
||||
if (resumable && !(e instanceof McpTransportSessionNotFoundException)) {
|
||||
Mono.from(reconnect.apply(this)).contextWrite(ctx).subscribe();
|
||||
}
|
||||
}).doOnNext(idAndMessage -> idAndMessage.getT1().ifPresent(id -> {
|
||||
String previousId = this.lastId.getAndSet(id);
|
||||
logger.debug("Updating last id {} -> {} for stream {}", previousId, id, this.streamId);
|
||||
})).flatMapIterable(Tuple2::getT2));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -4,20 +4,19 @@
|
|||
|
||||
package io.modelcontextprotocol.spec;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.MonoSink;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.MonoSink;
|
||||
|
||||
/**
|
||||
* Default implementation of the MCP (Model Context Protocol) session that manages
|
||||
* bidirectional JSON-RPC communication between clients and servers. This implementation
|
||||
|
|
@ -37,7 +36,6 @@ import reactor.core.publisher.MonoSink;
|
|||
*/
|
||||
public class McpClientSession implements McpSession {
|
||||
|
||||
/** Logger for this class */
|
||||
private static final Logger logger = LoggerFactory.getLogger(McpClientSession.class);
|
||||
|
||||
/** Duration to wait for request responses before timing out */
|
||||
|
|
@ -61,8 +59,6 @@ public class McpClientSession implements McpSession {
|
|||
/** Atomic counter for generating unique request IDs */
|
||||
private final AtomicLong requestCounter = new AtomicLong(0);
|
||||
|
||||
private final Disposable connection;
|
||||
|
||||
/**
|
||||
* Functional interface for handling incoming JSON-RPC requests. Implementations
|
||||
* should process the request parameters and return a response.
|
||||
|
|
@ -117,12 +113,15 @@ public class McpClientSession implements McpSession {
|
|||
this.requestHandlers.putAll(requestHandlers);
|
||||
this.notificationHandlers.putAll(notificationHandlers);
|
||||
|
||||
// TODO: consider mono.transformDeferredContextual where the Context contains
|
||||
// the
|
||||
// Observation associated with the individual message - it can be used to
|
||||
// create child Observation and emit it together with the message to the
|
||||
// consumer
|
||||
this.connection = this.transport.connect(mono -> mono.doOnNext(this::handle)).subscribe();
|
||||
this.transport.connect(mono -> mono.doOnNext(this::handle)).subscribe();
|
||||
}
|
||||
|
||||
private void dismissPendingResponses() {
|
||||
this.pendingResponses.forEach((id, sink) -> {
|
||||
logger.warn("Abruptly terminating exchange for request {}", id);
|
||||
sink.error(new RuntimeException("MCP session with server terminated"));
|
||||
});
|
||||
this.pendingResponses.clear();
|
||||
}
|
||||
|
||||
private void handle(McpSchema.JSONRPCMessage message) {
|
||||
|
|
@ -231,17 +230,15 @@ public class McpClientSession implements McpSession {
|
|||
String requestId = this.generateRequestId();
|
||||
|
||||
return Mono.deferContextual(ctx -> Mono.<McpSchema.JSONRPCResponse>create(sink -> {
|
||||
logger.debug("Sending message for method {}", method);
|
||||
this.pendingResponses.put(requestId, sink);
|
||||
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
|
||||
requestId, requestParams);
|
||||
this.transport.sendMessage(jsonrpcRequest)
|
||||
.contextWrite(ctx)
|
||||
// TODO: It's most efficient to create a dedicated Subscriber here
|
||||
.subscribe(v -> {
|
||||
}, error -> {
|
||||
this.pendingResponses.remove(requestId);
|
||||
sink.error(error);
|
||||
});
|
||||
this.transport.sendMessage(jsonrpcRequest).contextWrite(ctx).subscribe(v -> {
|
||||
}, error -> {
|
||||
this.pendingResponses.remove(requestId);
|
||||
sink.error(error);
|
||||
});
|
||||
})).timeout(this.requestTimeout).handle((jsonRpcResponse, sink) -> {
|
||||
if (jsonRpcResponse.error() != null) {
|
||||
logger.error("Error handling request: {}", jsonRpcResponse.error());
|
||||
|
|
@ -277,10 +274,7 @@ public class McpClientSession implements McpSession {
|
|||
*/
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
this.connection.dispose();
|
||||
return transport.closeGracefully();
|
||||
});
|
||||
return Mono.fromRunnable(this::dismissPendingResponses);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -288,8 +282,7 @@ public class McpClientSession implements McpSession {
|
|||
*/
|
||||
@Override
|
||||
public void close() {
|
||||
this.connection.dispose();
|
||||
transport.close();
|
||||
dismissPendingResponses();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,18 +3,38 @@
|
|||
*/
|
||||
package io.modelcontextprotocol.spec;
|
||||
|
||||
import java.util.function.Consumer;
|
||||
import java.util.function.Function;
|
||||
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
/**
|
||||
* Marker interface for the client-side MCP transport.
|
||||
* Interface for the client side of the {@link McpTransport}. It allows setting handlers
|
||||
* for messages that are incoming from the MCP server and hooking in to exceptions raised
|
||||
* on the transport layer.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
public interface McpClientTransport extends McpTransport {
|
||||
|
||||
/**
|
||||
* Used to register the incoming messages' handler and potentially (eagerly) connect
|
||||
* to the server.
|
||||
* @param handler a transformer for incoming messages
|
||||
* @return a {@link Mono} that terminates upon successful client setup. It can mean
|
||||
* establishing a connection (which can be later disposed) but it doesn't have to,
|
||||
* depending on the transport type. The successful termination of the returned
|
||||
* {@link Mono} simply means the client can now be used. An error can be retried
|
||||
* according to the application requirements.
|
||||
*/
|
||||
Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler);
|
||||
|
||||
/**
|
||||
* Sets the exception handler for exceptions raised on the transport layer.
|
||||
* @param handler Allows reacting to transport level exceptions by the higher layers
|
||||
*/
|
||||
default void setExceptionHandler(Consumer<Throwable> handler) {
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue