Module and Package restructoring
- mcp-transport -> mcp-spring - mcp-webflux-sse-transport -> mcp-spring-webflux - pom coordinates changed accordingly - mpc-webmvc-sse-transport -> mcp-spring-webmvc - pom coordinates changed accordingly - docs contains the mdx version of the documentation
This commit is contained in:
Christian Tzolov
parent
7e73152ada
commit
585f30cf33
10 changed files with 51 additions and 133 deletions
|
|
@ -0,0 +1,324 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
*/
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.function.BiConsumer;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.spec.ClientMcpTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.core.publisher.SynchronousSink;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.util.retry.Retry;
|
||||
import reactor.util.retry.Retry.RetrySignal;
|
||||
|
||||
import org.springframework.core.ParameterizedTypeReference;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Server-Sent Events (SSE) implementation of the
|
||||
* {@link io.modelcontextprotocol.spec.McpTransport} that follows the MCP HTTP with SSE
|
||||
* transport specification.
|
||||
*
|
||||
* <p>
|
||||
* This transport establishes a bidirectional communication channel where:
|
||||
* <ul>
|
||||
* <li>Inbound messages are received through an SSE connection from the server</li>
|
||||
* <li>Outbound messages are sent via HTTP POST requests to a server-provided
|
||||
* endpoint</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>
|
||||
* The message flow follows these steps:
|
||||
* <ol>
|
||||
* <li>The client establishes an SSE connection to the server's /sse endpoint</li>
|
||||
* <li>The server sends an 'endpoint' event containing the URI for sending messages</li>
|
||||
* </ol>
|
||||
*
|
||||
* This implementation uses {@link WebClient} for HTTP communications and supports JSON
|
||||
* serialization/deserialization of messages.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @see <a href=
|
||||
* "https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse">MCP
|
||||
* HTTP with SSE Transport Specification</a>
|
||||
*/
|
||||
public class WebFluxSseClientTransport implements ClientMcpTransport {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(WebFluxSseClientTransport.class);
|
||||
|
||||
/**
|
||||
* 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";
|
||||
|
||||
/**
|
||||
* Event type for receiving the message endpoint URI from the server. The server MUST
|
||||
* send this event when a client connects, providing the URI where the client should
|
||||
* send its messages via HTTP POST.
|
||||
*/
|
||||
private static final String ENDPOINT_EVENT_TYPE = "endpoint";
|
||||
|
||||
/**
|
||||
* Default SSE endpoint path as specified by the MCP transport specification. This
|
||||
* endpoint is used to establish the SSE connection with the server.
|
||||
*/
|
||||
private static final String SSE_ENDPOINT = "/sse";
|
||||
|
||||
/**
|
||||
* Type reference for parsing SSE events containing string data.
|
||||
*/
|
||||
private static final ParameterizedTypeReference<ServerSentEvent<String>> SSE_TYPE = new ParameterizedTypeReference<>() {
|
||||
};
|
||||
|
||||
/**
|
||||
* WebClient instance for handling both SSE connections and HTTP POST requests. Used
|
||||
* for establishing the SSE connection and sending outbound messages.
|
||||
*/
|
||||
private final WebClient webClient;
|
||||
|
||||
/**
|
||||
* ObjectMapper for serializing outbound messages and deserializing inbound messages.
|
||||
* Handles conversion between JSON-RPC messages and their string representation.
|
||||
*/
|
||||
protected ObjectMapper objectMapper;
|
||||
|
||||
/**
|
||||
* Subscription for the SSE connection handling inbound messages. Used for cleanup
|
||||
* during transport shutdown.
|
||||
*/
|
||||
private Disposable inboundSubscription;
|
||||
|
||||
/**
|
||||
* Flag indicating if the transport is in the process of shutting down. Used to
|
||||
* prevent new operations during shutdown and handle cleanup gracefully.
|
||||
*/
|
||||
private volatile boolean isClosing = false;
|
||||
|
||||
/**
|
||||
* Sink for managing the message endpoint URI provided by the server. Stores the most
|
||||
* recent endpoint URI and makes it available for outbound message processing.
|
||||
*/
|
||||
protected final Sinks.One<String> messageEndpointSink = Sinks.one();
|
||||
|
||||
/**
|
||||
* Constructs a new SseClientTransport with the specified WebClient builder. Uses a
|
||||
* default ObjectMapper instance for JSON processing.
|
||||
* @param webClientBuilder the WebClient.Builder to use for creating the WebClient
|
||||
* instance
|
||||
* @throws IllegalArgumentException if webClientBuilder is null
|
||||
*/
|
||||
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder) {
|
||||
this(webClientBuilder, new ObjectMapper());
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructs a new SseClientTransport with the specified WebClient builder and
|
||||
* ObjectMapper. Initializes both inbound and outbound message processing pipelines.
|
||||
* @param webClientBuilder the WebClient.Builder to use for creating the WebClient
|
||||
* instance
|
||||
* @param objectMapper the ObjectMapper to use for JSON processing
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
*/
|
||||
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
|
||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
||||
|
||||
this.objectMapper = objectMapper;
|
||||
this.webClient = webClientBuilder.build();
|
||||
}
|
||||
|
||||
/**
|
||||
* Establishes a connection to the MCP server using Server-Sent Events (SSE). This
|
||||
* method initiates the SSE connection and sets up the message processing pipeline.
|
||||
*
|
||||
* <p>
|
||||
* The connection process follows these steps:
|
||||
* <ol>
|
||||
* <li>Establishes an SSE connection to the server's /sse endpoint</li>
|
||||
* <li>Waits for the server to send an 'endpoint' event with the message posting
|
||||
* URI</li>
|
||||
* <li>Sets up message handling for incoming JSON-RPC messages</li>
|
||||
* </ol>
|
||||
*
|
||||
* <p>
|
||||
* The connection is considered established only after receiving the endpoint event
|
||||
* from the server.
|
||||
* @param handler a function that processes incoming JSON-RPC messages and returns
|
||||
* responses
|
||||
* @return a Mono that completes when the connection is fully established
|
||||
* @throws McpError if there's an error processing SSE events or if an unrecognized
|
||||
* event type is received
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
|
||||
Flux<ServerSentEvent<String>> events = eventStream();
|
||||
this.inboundSubscription = events.concatMap(event -> Mono.just(event).<JSONRPCMessage>handle((e, s) -> {
|
||||
if (ENDPOINT_EVENT_TYPE.equals(event.event())) {
|
||||
String messageEndpointUri = event.data();
|
||||
if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
|
||||
s.complete();
|
||||
}
|
||||
else {
|
||||
// TODO: clarify with the spec if multiple events can be
|
||||
// received
|
||||
s.error(new McpError("Failed to handle SSE endpoint event"));
|
||||
}
|
||||
}
|
||||
else if (MESSAGE_EVENT_TYPE.equals(event.event())) {
|
||||
try {
|
||||
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, event.data());
|
||||
s.next(message);
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
s.error(ioException);
|
||||
}
|
||||
}
|
||||
else {
|
||||
s.error(new McpError("Received unrecognized SSE event type: " + event.event()));
|
||||
}
|
||||
}).transform(handler)).subscribe();
|
||||
|
||||
// The connection is established once the server sends the endpoint event
|
||||
return messageEndpointSink.asMono().then();
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a JSON-RPC message to the server using the endpoint provided during
|
||||
* connection.
|
||||
*
|
||||
* <p>
|
||||
* Messages are sent via HTTP POST requests to the server-provided endpoint URI. The
|
||||
* message is serialized to JSON before transmission. If the transport is in the
|
||||
* process of closing, the message send operation is skipped gracefully.
|
||||
* @param message the JSON-RPC message to send
|
||||
* @return a Mono that completes when the message has been sent successfully
|
||||
* @throws RuntimeException if message serialization fails
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> sendMessage(JSONRPCMessage message) {
|
||||
// The messageEndpoint is the endpoint URI to send the messages
|
||||
// It is provided by the server as part of the endpoint event
|
||||
return messageEndpointSink.asMono().flatMap(messageEndpointUri -> {
|
||||
if (isClosing) {
|
||||
return Mono.empty();
|
||||
}
|
||||
try {
|
||||
String jsonText = this.objectMapper.writeValueAsString(message);
|
||||
return webClient.post()
|
||||
.uri(messageEndpointUri)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue(jsonText)
|
||||
.retrieve()
|
||||
.toBodilessEntity()
|
||||
.doOnSuccess(response -> {
|
||||
logger.debug("Message sent successfully");
|
||||
})
|
||||
.doOnError(error -> {
|
||||
if (!isClosing) {
|
||||
logger.error("Error sending message: {}", error.getMessage());
|
||||
}
|
||||
});
|
||||
}
|
||||
catch (IOException e) {
|
||||
if (!isClosing) {
|
||||
return Mono.error(new RuntimeException("Failed to serialize message", e));
|
||||
}
|
||||
return Mono.empty();
|
||||
}
|
||||
}).then(); // TODO: Consider non-200-ok response
|
||||
}
|
||||
|
||||
/**
|
||||
* Initializes and starts the inbound SSE event processing. Establishes the SSE
|
||||
* connection and sets up event handling for both message and endpoint events.
|
||||
* Includes automatic retry logic for handling transient connection failures.
|
||||
*/
|
||||
// visible for tests
|
||||
protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
|
||||
return this.webClient
|
||||
.get()
|
||||
.uri(SSE_ENDPOINT)
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.retrieve()
|
||||
.bodyToFlux(SSE_TYPE)
|
||||
.retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler)));
|
||||
} // @formatter:on
|
||||
|
||||
/**
|
||||
* Retry handler for the inbound SSE stream. Implements the retry logic for handling
|
||||
* connection failures and other errors.
|
||||
*/
|
||||
private BiConsumer<RetrySignal, SynchronousSink<Object>> inboundRetryHandler = (retrySpec, sink) -> {
|
||||
if (isClosing) {
|
||||
logger.debug("SSE connection closed during shutdown");
|
||||
sink.error(retrySpec.failure());
|
||||
return;
|
||||
}
|
||||
if (retrySpec.failure() instanceof IOException) {
|
||||
logger.debug("Retrying SSE connection after IO error");
|
||||
sink.next(retrySpec);
|
||||
return;
|
||||
}
|
||||
logger.error("Fatal SSE error, not retrying: {}", retrySpec.failure().getMessage());
|
||||
sink.error(retrySpec.failure());
|
||||
};
|
||||
|
||||
/**
|
||||
* Implements graceful shutdown of the transport. Cleans up all resources including
|
||||
* subscriptions and schedulers. Ensures orderly shutdown of both inbound and outbound
|
||||
* message processing.
|
||||
* @return a Mono that completes when shutdown is finished
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() { // @formatter:off
|
||||
return Mono.fromRunnable(() -> {
|
||||
isClosing = true;
|
||||
|
||||
// Dispose of subscriptions
|
||||
|
||||
if (inboundSubscription != null) {
|
||||
inboundSubscription.dispose();
|
||||
}
|
||||
|
||||
})
|
||||
.then()
|
||||
.subscribeOn(Schedulers.boundedElastic());
|
||||
} // @formatter:on
|
||||
|
||||
/**
|
||||
* Unmarshals data from a generic Object into the specified type using the configured
|
||||
* ObjectMapper.
|
||||
*
|
||||
* <p>
|
||||
* This method is particularly useful when working with JSON-RPC parameters or result
|
||||
* objects that need to be converted to specific Java types. It leverages Jackson's
|
||||
* type conversion capabilities to handle complex object structures.
|
||||
* @param <T> the target type to convert the data into
|
||||
* @param data the source object to convert
|
||||
* @param typeRef the TypeReference describing the target type
|
||||
* @return the unmarshaled object of type T
|
||||
* @throws IllegalArgumentException if the conversion cannot be performed
|
||||
*/
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
|
||||
return this.objectMapper.convertValue(data, typeRef);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,410 @@
|
|||
package io.modelcontextprotocol.server.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.ServerMcpTransport;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.reactive.function.server.RouterFunction;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
import org.springframework.web.reactive.function.server.ServerRequest;
|
||||
import org.springframework.web.reactive.function.server.ServerResponse;
|
||||
|
||||
/**
|
||||
* Server-side implementation of the MCP (Model Context Protocol) HTTP transport using
|
||||
* Server-Sent Events (SSE). This implementation provides a bidirectional communication
|
||||
* channel between MCP clients and servers using HTTP POST for client-to-server messages
|
||||
* and SSE for server-to-client messages.
|
||||
*
|
||||
* <p>
|
||||
* Key features:
|
||||
* <ul>
|
||||
* <li>Implements the {@link ServerMcpTransport} interface for MCP server transport
|
||||
* functionality</li>
|
||||
* <li>Uses WebFlux for non-blocking request handling and SSE support</li>
|
||||
* <li>Maintains client sessions for reliable message delivery</li>
|
||||
* <li>Supports graceful shutdown with session cleanup</li>
|
||||
* <li>Thread-safe message broadcasting to multiple clients</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>
|
||||
* The transport sets up two main endpoints:
|
||||
* <ul>
|
||||
* <li>SSE endpoint (/sse) - For establishing SSE connections with clients</li>
|
||||
* <li>Message endpoint (configurable) - For receiving JSON-RPC messages from clients</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>
|
||||
* This implementation is thread-safe and can handle multiple concurrent client
|
||||
* connections. It uses {@link ConcurrentHashMap} for session management and Reactor's
|
||||
* {@link Sinks} for thread-safe message broadcasting.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Alexandros Pappas
|
||||
* @see ServerMcpTransport
|
||||
* @see ServerSentEvent
|
||||
*/
|
||||
public class WebFluxSseServerTransport implements ServerMcpTransport {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(WebFluxSseServerTransport.class);
|
||||
|
||||
/**
|
||||
* Event type for JSON-RPC messages sent through the SSE connection.
|
||||
*/
|
||||
public static final String MESSAGE_EVENT_TYPE = "message";
|
||||
|
||||
/**
|
||||
* Event type for sending the message endpoint URI to clients.
|
||||
*/
|
||||
public static final String ENDPOINT_EVENT_TYPE = "endpoint";
|
||||
|
||||
/**
|
||||
* Default SSE endpoint path as specified by the MCP transport specification.
|
||||
*/
|
||||
public static final String DEFAULT_SSE_ENDPOINT = "/sse";
|
||||
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
private final String messageEndpoint;
|
||||
|
||||
private final String sseEndpoint;
|
||||
|
||||
private final RouterFunction<?> routerFunction;
|
||||
|
||||
/**
|
||||
* Map of active client sessions, keyed by session ID.
|
||||
*/
|
||||
private final ConcurrentHashMap<String, ClientSession> sessions = new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* Flag indicating if the transport is shutting down.
|
||||
*/
|
||||
private volatile boolean isClosing = false;
|
||||
|
||||
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
|
||||
|
||||
/**
|
||||
* Constructs a new WebFlux SSE server transport instance.
|
||||
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
|
||||
* of MCP messages. Must not be null.
|
||||
* @param messageEndpoint The endpoint URI where clients should send their JSON-RPC
|
||||
* messages. This endpoint will be communicated to clients during SSE connection
|
||||
* setup. Must not be null.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
*/
|
||||
public WebFluxSseServerTransport(ObjectMapper objectMapper, String messageEndpoint, String sseEndpoint) {
|
||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||
Assert.notNull(messageEndpoint, "Message endpoint must not be null");
|
||||
Assert.notNull(sseEndpoint, "SSE endpoint must not be null");
|
||||
|
||||
this.objectMapper = objectMapper;
|
||||
this.messageEndpoint = messageEndpoint;
|
||||
this.sseEndpoint = sseEndpoint;
|
||||
this.routerFunction = RouterFunctions.route()
|
||||
.GET(this.sseEndpoint, this::handleSseConnection)
|
||||
.POST(this.messageEndpoint, this::handleMessage)
|
||||
.build();
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructs a new WebFlux SSE server transport instance with the default SSE
|
||||
* endpoint.
|
||||
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
|
||||
* of MCP messages. Must not be null.
|
||||
* @param messageEndpoint The endpoint URI where clients should send their JSON-RPC
|
||||
* messages. This endpoint will be communicated to clients during SSE connection
|
||||
* setup. Must not be null.
|
||||
* @throws IllegalArgumentException if either parameter is null
|
||||
*/
|
||||
public WebFluxSseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
|
||||
this(objectMapper, messageEndpoint, DEFAULT_SSE_ENDPOINT);
|
||||
}
|
||||
|
||||
/**
|
||||
* Configures the message handler for this transport. In the WebFlux SSE
|
||||
* implementation, this method stores the handler for processing incoming messages but
|
||||
* doesn't establish any connections since the server accepts connections rather than
|
||||
* initiating them.
|
||||
* @param handler A function that processes incoming JSON-RPC messages and returns
|
||||
* responses. This handler will be called for each message received through the
|
||||
* message endpoint.
|
||||
* @return An empty Mono since the server doesn't initiate connections
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
this.connectHandler = handler;
|
||||
// Server-side transport doesn't initiate connections
|
||||
return Mono.empty().then();
|
||||
}
|
||||
|
||||
/**
|
||||
* Broadcasts a JSON-RPC message to all connected clients through their SSE
|
||||
* connections. The message is serialized to JSON and sent as a server-sent event to
|
||||
* each active session.
|
||||
*
|
||||
* <p>
|
||||
* The method:
|
||||
* <ul>
|
||||
* <li>Serializes the message to JSON</li>
|
||||
* <li>Creates a server-sent event with the message data</li>
|
||||
* <li>Attempts to send the event to all active sessions</li>
|
||||
* <li>Tracks and reports any delivery failures</li>
|
||||
* </ul>
|
||||
* @param message The JSON-RPC message to broadcast
|
||||
* @return A Mono that completes when the message has been sent to all sessions, or
|
||||
* errors if any session fails to receive the message
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (sessions.isEmpty()) {
|
||||
logger.debug("No active sessions to broadcast message to");
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
return Mono.<Void>create(sink -> {
|
||||
try {// @formatter:off
|
||||
String jsonText = objectMapper.writeValueAsString(message);
|
||||
ServerSentEvent<Object> event = ServerSentEvent.builder()
|
||||
.event(MESSAGE_EVENT_TYPE)
|
||||
.data(jsonText)
|
||||
.build();
|
||||
|
||||
logger.debug("Attempting to broadcast message to {} active sessions", sessions.size());
|
||||
|
||||
List<String> failedSessions = sessions.values().stream()
|
||||
.filter(session -> session.messageSink.tryEmitNext(event).isFailure())
|
||||
.map(session -> session.id)
|
||||
.toList();
|
||||
|
||||
if (failedSessions.isEmpty()) {
|
||||
logger.debug("Successfully broadcast message to all sessions");
|
||||
sink.success();
|
||||
}
|
||||
else {
|
||||
String error = "Failed to broadcast message to sessions: " + String.join(", ", failedSessions);
|
||||
logger.error(error);
|
||||
sink.error(new RuntimeException(error));
|
||||
} // @formatter:on
|
||||
}
|
||||
catch (IOException e) {
|
||||
logger.error("Failed to serialize message: {}", e.getMessage());
|
||||
sink.error(e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts data from one type to another using the configured ObjectMapper. This
|
||||
* method is primarily used for converting between different representations of
|
||||
* JSON-RPC message data.
|
||||
* @param <T> The target type to convert to
|
||||
* @param data The source data to convert
|
||||
* @param typeRef Type reference describing the target type
|
||||
* @return The converted data
|
||||
* @throws IllegalArgumentException if the conversion fails
|
||||
*/
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
|
||||
return this.objectMapper.convertValue(data, typeRef);
|
||||
}
|
||||
|
||||
/**
|
||||
* Initiates a graceful shutdown of the transport. This method ensures all active
|
||||
* sessions are properly closed and cleaned up.
|
||||
*
|
||||
* <p>
|
||||
* The shutdown process:
|
||||
* <ul>
|
||||
* <li>Marks the transport as closing to prevent new connections</li>
|
||||
* <li>Closes each active session</li>
|
||||
* <li>Removes closed sessions from the sessions map</li>
|
||||
* <li>Times out after 5 seconds if shutdown takes too long</li>
|
||||
* </ul>
|
||||
* @return A Mono that completes when all sessions have been closed
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.fromRunnable(() -> {
|
||||
isClosing = true;
|
||||
logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size());
|
||||
}).then(Mono.when(sessions.values().stream().map(session -> {
|
||||
String sessionId = session.id;
|
||||
return Mono.fromRunnable(() -> session.close())
|
||||
.then(Mono.delay(Duration.ofMillis(100)))
|
||||
.then(Mono.fromRunnable(() -> sessions.remove(sessionId)));
|
||||
}).toList()))
|
||||
.timeout(Duration.ofSeconds(5))
|
||||
.doOnSuccess(v -> logger.info("Graceful shutdown completed"))
|
||||
.doOnError(e -> logger.error("Error during graceful shutdown: {}", e.getMessage()));
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the WebFlux router function that defines the transport's HTTP endpoints.
|
||||
* This router function should be integrated into the application's web configuration.
|
||||
*
|
||||
* <p>
|
||||
* The router function defines two endpoints:
|
||||
* <ul>
|
||||
* <li>GET {sseEndpoint} - For establishing SSE connections</li>
|
||||
* <li>POST {messageEndpoint} - For receiving client messages</li>
|
||||
* </ul>
|
||||
* @return The configured {@link RouterFunction} for handling HTTP requests
|
||||
*/
|
||||
public RouterFunction<?> getRouterFunction() {
|
||||
return this.routerFunction;
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles new SSE connection requests from clients. Creates a new session for each
|
||||
* connection and sets up the SSE event stream.
|
||||
*
|
||||
* <p>
|
||||
* The handler performs the following steps:
|
||||
* <ul>
|
||||
* <li>Generates a unique session ID</li>
|
||||
* <li>Creates a new ClientSession instance</li>
|
||||
* <li>Sends the message endpoint URI as an initial event</li>
|
||||
* <li>Sets up message forwarding for the session</li>
|
||||
* <li>Handles connection cleanup on completion or errors</li>
|
||||
* </ul>
|
||||
* @param request The incoming server request
|
||||
* @return A response with the SSE event stream
|
||||
*/
|
||||
private Mono<ServerResponse> handleSseConnection(ServerRequest request) {
|
||||
if (isClosing) {
|
||||
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
|
||||
}
|
||||
String sessionId = UUID.randomUUID().toString();
|
||||
logger.debug("Creating new SSE connection for session: {}", sessionId);
|
||||
ClientSession session = new ClientSession(sessionId);
|
||||
this.sessions.put(sessionId, session);
|
||||
|
||||
return ServerResponse.ok()
|
||||
.contentType(MediaType.TEXT_EVENT_STREAM)
|
||||
.body(Flux.<ServerSentEvent<?>>create(sink -> {
|
||||
// Send initial endpoint event
|
||||
logger.debug("Sending initial endpoint event to session: {}", sessionId);
|
||||
sink.next(ServerSentEvent.builder().event(ENDPOINT_EVENT_TYPE).data(messageEndpoint).build());
|
||||
|
||||
// Subscribe to session messages
|
||||
session.messageSink.asFlux()
|
||||
.doOnSubscribe(s -> logger.debug("Session {} subscribed to message sink", sessionId))
|
||||
.doOnComplete(() -> {
|
||||
logger.debug("Session {} completed", sessionId);
|
||||
sessions.remove(sessionId);
|
||||
})
|
||||
.doOnError(error -> {
|
||||
logger.error("Error in session {}: {}", sessionId, error.getMessage());
|
||||
sessions.remove(sessionId);
|
||||
})
|
||||
.doOnCancel(() -> {
|
||||
logger.debug("Session {} cancelled", sessionId);
|
||||
sessions.remove(sessionId);
|
||||
})
|
||||
.subscribe(event -> {
|
||||
logger.debug("Forwarding event to session {}: {}", sessionId, event);
|
||||
sink.next(event);
|
||||
}, sink::error, sink::complete);
|
||||
|
||||
sink.onCancel(() -> {
|
||||
logger.debug("Session {} cancelled", sessionId);
|
||||
sessions.remove(sessionId);
|
||||
});
|
||||
}), ServerSentEvent.class);
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles incoming JSON-RPC messages from clients. Deserializes the message and
|
||||
* processes it through the configured message handler.
|
||||
*
|
||||
* <p>
|
||||
* The handler:
|
||||
* <ul>
|
||||
* <li>Deserializes the incoming JSON-RPC message</li>
|
||||
* <li>Passes it through the message handler chain</li>
|
||||
* <li>Returns appropriate HTTP responses based on processing results</li>
|
||||
* <li>Handles various error conditions with appropriate error responses</li>
|
||||
* </ul>
|
||||
* @param request The incoming server request containing the JSON-RPC message
|
||||
* @return A response indicating the message processing result
|
||||
*/
|
||||
private Mono<ServerResponse> handleMessage(ServerRequest request) {
|
||||
if (isClosing) {
|
||||
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
|
||||
}
|
||||
|
||||
return request.bodyToMono(String.class).flatMap(body -> {
|
||||
try {
|
||||
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body);
|
||||
return Mono.just(message)
|
||||
.transform(this.connectHandler)
|
||||
.flatMap(response -> ServerResponse.ok().build())
|
||||
.onErrorResume(error -> {
|
||||
logger.error("Error processing message: {}", error.getMessage());
|
||||
return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR)
|
||||
.bodyValue(new McpError(error.getMessage()));
|
||||
});
|
||||
}
|
||||
catch (IllegalArgumentException | IOException e) {
|
||||
logger.error("Failed to deserialize message: {}", e.getMessage());
|
||||
return ServerResponse.badRequest().bodyValue(new McpError("Invalid message format"));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents an active client SSE connection session. Manages the message sink for
|
||||
* sending events to the client and handles session lifecycle.
|
||||
*
|
||||
* <p>
|
||||
* Each session:
|
||||
* <ul>
|
||||
* <li>Has a unique identifier</li>
|
||||
* <li>Maintains its own message sink for event broadcasting</li>
|
||||
* <li>Supports clean shutdown through the close method</li>
|
||||
* </ul>
|
||||
*/
|
||||
private static class ClientSession {
|
||||
|
||||
private final String id;
|
||||
|
||||
private final Sinks.Many<ServerSentEvent<?>> messageSink;
|
||||
|
||||
ClientSession(String id) {
|
||||
this.id = id;
|
||||
logger.debug("Creating new session: {}", id);
|
||||
this.messageSink = Sinks.many().replay().latest();
|
||||
logger.debug("Session {} initialized with replay sink", id);
|
||||
}
|
||||
|
||||
void close() {
|
||||
logger.debug("Closing session: {}", id);
|
||||
Sinks.EmitResult result = messageSink.tryEmitComplete();
|
||||
if (result.isFailure()) {
|
||||
logger.warn("Failed to complete message sink for session {}: {}", id, result);
|
||||
}
|
||||
else {
|
||||
logger.debug("Successfully completed message sink for session {}", id);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,462 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
*/
|
||||
package io.modelcontextprotocol;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.client.McpClient;
|
||||
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
|
||||
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
|
||||
import io.modelcontextprotocol.server.McpServer;
|
||||
import io.modelcontextprotocol.server.McpServerFeatures;
|
||||
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Role;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Root;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Tool;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.params.ParameterizedTest;
|
||||
import org.junit.jupiter.params.provider.ValueSource;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.client.RestClient;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
public class WebFluxSseIntegrationTests {
|
||||
|
||||
private static final int PORT = 8182;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
private WebFluxSseServerTransport mcpServerTransport;
|
||||
|
||||
ConcurrentHashMap<String, McpClient.SyncSpec> clientBulders = new ConcurrentHashMap<>();
|
||||
|
||||
@BeforeEach
|
||||
public void before() {
|
||||
|
||||
this.mcpServerTransport = new WebFluxSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
|
||||
clientBulders.put("httpclient", McpClient.sync(new HttpClientSseClientTransport("http://localhost:" + PORT)));
|
||||
clientBulders.put("webflux",
|
||||
McpClient.sync(new WebFluxSseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT))));
|
||||
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void after() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
// Sampling Tests
|
||||
// ---------------------------------------
|
||||
@Test
|
||||
void testCreateMessageWithoutInitialization() {
|
||||
var mcpAsyncServer = McpServer.async(mcpServerTransport).serverInfo("test-server", "1.0.0").build();
|
||||
|
||||
var messages = List
|
||||
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
|
||||
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
|
||||
|
||||
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
|
||||
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized. Call the initialize method first!");
|
||||
});
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testCreateMessageWithoutSamplingCapabilities(String clientType) {
|
||||
|
||||
var mcpAsyncServer = McpServer.async(mcpServerTransport).serverInfo("test-server", "1.0.0").build();
|
||||
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build();
|
||||
|
||||
InitializeResult initResult = client.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
var messages = List
|
||||
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
|
||||
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
|
||||
|
||||
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
|
||||
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be configured with sampling capabilities");
|
||||
});
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testCreateMessageSuccess(String clientType) throws InterruptedException {
|
||||
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
var mcpAsyncServer = McpServer.async(mcpServerTransport).serverInfo("test-server", "1.0.0").build();
|
||||
|
||||
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
|
||||
assertThat(request.messages()).hasSize(1);
|
||||
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
|
||||
|
||||
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
|
||||
CreateMessageResult.StopReason.STOP_SEQUENCE);
|
||||
};
|
||||
|
||||
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
|
||||
.capabilities(ClientCapabilities.builder().sampling().build())
|
||||
.sampling(samplingHandler)
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = client.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
var messages = List
|
||||
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
|
||||
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
|
||||
|
||||
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
|
||||
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.createMessage(request)).consumeNextWith(result -> {
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.role()).isEqualTo(Role.USER);
|
||||
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
|
||||
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
|
||||
assertThat(result.model()).isEqualTo("MockModelName");
|
||||
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
|
||||
}).verifyComplete();
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
// Roots Tests
|
||||
// ---------------------------------------
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testRootsSuccess(String clientType) {
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
|
||||
|
||||
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
|
||||
.build();
|
||||
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.roots(roots)
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(rootsRef.get()).isNull();
|
||||
|
||||
assertThat(mcpServer.listRoots().roots()).containsAll(roots);
|
||||
|
||||
mcpClient.rootsListChangedNotification();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(roots);
|
||||
});
|
||||
|
||||
// Remove a root
|
||||
mcpClient.removeRoot(roots.get(0).uri());
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(roots.get(1)));
|
||||
});
|
||||
|
||||
// Add a new root
|
||||
var root3 = new Root("uri3://", "root3");
|
||||
mcpClient.addRoot(root3);
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3));
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testRootsWithoutCapability(String clientType) {
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
var mcpServer = McpServer.sync(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> {
|
||||
}).build();
|
||||
|
||||
// Create client without roots capability
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()) // No
|
||||
// roots
|
||||
// capability
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
// Attempt to list roots should fail
|
||||
assertThatThrownBy(() -> mcpServer.listRoots().roots()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Roots not supported");
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testRootsWithEmptyRootsList(String clientType) {
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
|
||||
.build();
|
||||
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.roots(List.of()) // Empty roots list
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
mcpClient.rootsListChangedNotification();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).isEmpty();
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testRootsWithMultipleConsumers(String clientType) {
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
List<Root> roots = List.of(new Root("uri1://", "root1"));
|
||||
|
||||
AtomicReference<List<Root>> rootsRef1 = new AtomicReference<>();
|
||||
AtomicReference<List<Root>> rootsRef2 = new AtomicReference<>();
|
||||
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.rootsChangeConsumer(rootsUpdate -> rootsRef1.set(rootsUpdate))
|
||||
.rootsChangeConsumer(rootsUpdate -> rootsRef2.set(rootsUpdate))
|
||||
.build();
|
||||
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.roots(roots)
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
mcpClient.rootsListChangedNotification();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef1.get()).containsAll(roots);
|
||||
assertThat(rootsRef2.get()).containsAll(roots);
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testRootsServerCloseWithActiveSubscription(String clientType) {
|
||||
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
List<Root> roots = List.of(new Root("uri1://", "root1"));
|
||||
|
||||
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
|
||||
.build();
|
||||
|
||||
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.roots(roots)
|
||||
.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
mcpClient.rootsListChangedNotification();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(roots);
|
||||
});
|
||||
|
||||
// Close server while subscription is active
|
||||
mcpServer.close();
|
||||
|
||||
// Verify client can handle server closure gracefully
|
||||
mcpClient.close();
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
// Tools Tests
|
||||
// ---------------------------------------
|
||||
|
||||
String emptyJsonSchema = """
|
||||
{
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
""";
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testToolCallSuccess(String clientType) {
|
||||
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
||||
McpServerFeatures.SyncToolRegistration tool1 = new McpServerFeatures.SyncToolRegistration(
|
||||
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), request -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/modelcontextprotocol/specification/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
return callResponse;
|
||||
});
|
||||
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tools(tool1)
|
||||
.build();
|
||||
|
||||
var mcpClient = clientBuilder.build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response).isEqualTo(callResponse);
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@ParameterizedTest(name = "{0} : {displayName} ")
|
||||
@ValueSource(strings = { "httpclient", "webflux" })
|
||||
void testToolListChangeHandlingSuccess(String clientType) {
|
||||
|
||||
var clientBuilder = clientBulders.get(clientType);
|
||||
|
||||
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
||||
McpServerFeatures.SyncToolRegistration tool1 = new McpServerFeatures.SyncToolRegistration(
|
||||
new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), request -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/modelcontextprotocol/specification/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
return callResponse;
|
||||
});
|
||||
|
||||
var mcpServer = McpServer.sync(mcpServerTransport)
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tools(tool1)
|
||||
.build();
|
||||
|
||||
AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
|
||||
var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/modelcontextprotocol/specification/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
rootsRef.set(toolsUpdate);
|
||||
}).build();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(rootsRef.get()).isNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
mcpServer.notifyToolsListChanged();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(tool1.tool()));
|
||||
});
|
||||
|
||||
// Remove a tool
|
||||
mcpServer.removeTool("tool1");
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).isEmpty();
|
||||
});
|
||||
|
||||
// Add a new tool
|
||||
McpServerFeatures.SyncToolRegistration tool2 = new McpServerFeatures.SyncToolRegistration(
|
||||
new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema), request -> callResponse);
|
||||
|
||||
mcpServer.addTool(tool2);
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
|
||||
import io.modelcontextprotocol.spec.ClientMcpTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpAsyncClient} with {@link WebFluxSseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class WebFluxSseMcpAsyncClientTests 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:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected ClientMcpTransport createMcpTransport() {
|
||||
return new WebFluxSseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@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,49 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport;
|
||||
import io.modelcontextprotocol.spec.ClientMcpTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpSyncClient} with {@link WebFluxSseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class WebFluxSseMcpSyncClientTests 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:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected ClientMcpTransport createMcpTransport() {
|
||||
return new WebFluxSseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,321 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Tests for the {@link WebFluxSseClientTransport} class.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
class WebFluxSseClientTransportTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
private TestSseClientTransport transport;
|
||||
|
||||
private WebClient.Builder webClientBuilder;
|
||||
|
||||
private ObjectMapper objectMapper;
|
||||
|
||||
// Test class to access protected methods
|
||||
static class TestSseClientTransport extends WebFluxSseClientTransport {
|
||||
|
||||
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
|
||||
|
||||
private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
|
||||
|
||||
public TestSseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
|
||||
super(webClientBuilder, objectMapper);
|
||||
}
|
||||
|
||||
// @Override
|
||||
// public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>,
|
||||
// Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
// simulateEndpointEvent("https://localhost:3001");
|
||||
// return super.connect(handler);
|
||||
// }
|
||||
|
||||
@Override
|
||||
protected Flux<ServerSentEvent<String>> eventStream() {
|
||||
return super.eventStream().mergeWith(events.asFlux());
|
||||
}
|
||||
|
||||
public String getLastEndpoint() {
|
||||
return messageEndpointSink.asMono().block();
|
||||
}
|
||||
|
||||
public int getInboundMessageCount() {
|
||||
return inboundMessageCount.get();
|
||||
}
|
||||
|
||||
public void simulateEndpointEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("endpoint").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
public void simulateMessageEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("message").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
void startContainer() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
startContainer();
|
||||
webClientBuilder = WebClient.builder().baseUrl(host);
|
||||
objectMapper = new ObjectMapper();
|
||||
transport = new TestSseClientTransport(webClientBuilder, objectMapper);
|
||||
transport.connect(Function.identity()).block();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void afterEach() {
|
||||
if (transport != null) {
|
||||
assertThatCode(() -> transport.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
cleanup();
|
||||
}
|
||||
|
||||
void cleanup() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testEndpointEventHandling() {
|
||||
assertThat(transport.getLastEndpoint()).startsWith("/message?");
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorValidation() {
|
||||
assertThatThrownBy(() -> new WebFluxSseClientTransport(null)).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("WebClient.Builder must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new WebFluxSseClientTransport(webClientBuilder, null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("ObjectMapper must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageProcessing() {
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Simulate receiving the message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "test-method",
|
||||
"id": "test-id",
|
||||
"params": {"key": "value"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Subscribe to messages and verify
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testResponseMessageProcessing() {
|
||||
// Simulate receiving a response message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"result": {"status": "success"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testErrorMessageProcessing() {
|
||||
// Simulate receiving an error message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid Request"
|
||||
}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotificationMessageProcessing() {
|
||||
// Simulate receiving a notification message (no id)
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "update",
|
||||
"params": {"status": "processing"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify the notification was processed
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
// Test graceful shutdown
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message is not processed after shutdown
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Message count should remain 0 after shutdown
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(0);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRetryBehavior() {
|
||||
// Create a WebClient that simulates connection failures
|
||||
WebClient.Builder failingWebClientBuilder = WebClient.builder().baseUrl("http://non-existent-host");
|
||||
|
||||
WebFluxSseClientTransport failingTransport = new WebFluxSseClientTransport(failingWebClientBuilder);
|
||||
|
||||
// Verify that the transport attempts to reconnect
|
||||
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
|
||||
|
||||
// Clean up
|
||||
failingTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMultipleMessageProcessing() {
|
||||
// Simulate receiving multiple messages in sequence
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method1",
|
||||
"id": "id1",
|
||||
"params": {"key": "value1"}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method2",
|
||||
"id": "id2",
|
||||
"params": {"key": "value2"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send corresponding messages
|
||||
JSONRPCRequest message1 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method1", "id1",
|
||||
Map.of("key", "value1"));
|
||||
|
||||
JSONRPCRequest message2 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method2", "id2",
|
||||
Map.of("key", "value2"));
|
||||
|
||||
// Verify both messages are processed
|
||||
StepVerifier.create(transport.sendMessage(message1).then(transport.sendMessage(message2))).verifyComplete();
|
||||
|
||||
// Verify message count
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageOrderPreservation() {
|
||||
// Simulate receiving messages in a specific order
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "first",
|
||||
"id": "1",
|
||||
"params": {"sequence": 1}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "second",
|
||||
"id": "2",
|
||||
"params": {"sequence": 2}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "third",
|
||||
"id": "3",
|
||||
"params": {"sequence": 3}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify message count and order
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(3);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.server;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransport;
|
||||
import io.modelcontextprotocol.spec.ServerMcpTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpAsyncServer} using {@link WebFluxSseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class WebFluxSseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
|
||||
|
||||
private static final int PORT = 8181;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
@Override
|
||||
protected ServerMcpTransport createMcpTransport() {
|
||||
var transport = new WebFluxSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.server;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransport;
|
||||
import io.modelcontextprotocol.spec.ServerMcpTransport;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpSyncServer} using {@link WebFluxSseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class WebFluxSseMcpSyncServerTests extends AbstractMcpSyncServerTests {
|
||||
|
||||
private static final int PORT = 8182;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
private WebFluxSseServerTransport transport;
|
||||
|
||||
@Override
|
||||
protected ServerMcpTransport createMcpTransport() {
|
||||
transport = new WebFluxSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,69 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
*/
|
||||
package io.modelcontextprotocol.server.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.util.concurrent.BlockingQueue;
|
||||
import java.util.concurrent.LinkedBlockingQueue;
|
||||
|
||||
public class BlockingInputStream extends InputStream {
|
||||
|
||||
private final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>();
|
||||
|
||||
private volatile boolean completed = false;
|
||||
|
||||
private volatile boolean closed = false;
|
||||
|
||||
@Override
|
||||
public int read() throws IOException {
|
||||
if (closed) {
|
||||
throw new IOException("Stream is closed");
|
||||
}
|
||||
|
||||
try {
|
||||
Integer value = queue.poll();
|
||||
if (value == null) {
|
||||
if (completed) {
|
||||
return -1;
|
||||
}
|
||||
value = queue.take(); // Blocks until data is available
|
||||
if (value == null && completed) {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
return value;
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new IOException("Read interrupted", e);
|
||||
}
|
||||
}
|
||||
|
||||
public void write(int b) {
|
||||
if (!closed && !completed) {
|
||||
queue.offer(b);
|
||||
}
|
||||
}
|
||||
|
||||
public void write(byte[] data) {
|
||||
if (!closed && !completed) {
|
||||
for (byte b : data) {
|
||||
queue.offer((int) b & 0xFF);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public void complete() {
|
||||
this.completed = true;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
this.closed = true;
|
||||
this.completed = true;
|
||||
this.queue.clear();
|
||||
}
|
||||
|
||||
}
|
||||
25
mcp-spring/mcp-spring-webflux/src/test/resources/logback.xml
Normal file
25
mcp-spring/mcp-spring-webflux/src/test/resources/logback.xml
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!DOCTYPE configuration>
|
||||
|
||||
<configuration>
|
||||
<appender name="CONSOLE" class="ch.qos.logback.core.ConsoleAppender">
|
||||
<encoder>
|
||||
<pattern>%d{yyyy-MM-dd HH:mm:ss} [%thread] %-5level %logger{36} - %msg%n</pattern>
|
||||
</encoder>
|
||||
</appender>
|
||||
|
||||
<!-- Main MCP package -->
|
||||
<logger name="org.springframework.ai.mcp" level="INFO"/>
|
||||
|
||||
<!-- Client packages -->
|
||||
<logger name="org.springframework.ai.mcp.client" level="INFO"/>
|
||||
|
||||
<!-- Spec package -->
|
||||
<logger name="org.springframework.ai.mcp.spec" level="INFO"/>
|
||||
|
||||
|
||||
<!-- Root logger -->
|
||||
<root level="INFO">
|
||||
<appender-ref ref="CONSOLE"/>
|
||||
</root>
|
||||
</configuration>
|
||||
Loading…
Reference in a new issue