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 2025-01-29 13:40:48 +01:00
commit 585f30cf33
19 changed files with 51 additions and 133 deletions

View file

@ -0,0 +1,55 @@
# WebFlux SSE Transport
```xml
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-spring-webflux</artifactId>
</dependency>
```
```java
String MESSAGE_ENDPOINT = "/mcp/message";
@Configuration
static class MyConfig {
// SSE transport
@Bean
public WebFluxSseServerTransport sseServerTransport() {
return new WebFluxSseServerTransport(new ObjectMapper(), "/mcp/message");
}
// Router function for SSE transport used by Spring WebFlux to start an HTTP
// server.
@Bean
public RouterFunction<?> mcpRouterFunction(WebFluxSseServerTransport transport) {
return transport.getRouterFunction();
}
@Bean
public McpAsyncServer mcpServer(ServerMcpTransport transport, OpenLibrary openLibrary) {
// Configure server capabilities with resource support
var capabilities = McpSchema.ServerCapabilities.builder()
.resources(false, true) // No subscribe support, but list changes notifications
.tools(true) // Tool support with list changes notifications
.prompts(true) // Prompt support with list changes notifications
.logging() // Logging support
.build();
// Create the server with both tool and resource capabilities
var server = McpServer.using(transport)
.serverInfo("MCP Demo Server", "1.0.0")
.capabilities(capabilities)
.resources(systemInfoResourceRegistration())
.prompts(greetingPromptRegistration())
.tools(openLibraryToolRegistrations(openLibrary))
.async();
return server;
}
// ...
}
```

View file

@ -0,0 +1,121 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-parent</artifactId>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>mcp-spring-webflux</artifactId>
<packaging>jar</packaging>
<name>WebFlux implementation of the Java MCP SSE transport</name>
<description></description>
<url>https://github.com/modelcontextprotocol/java-sdk</url>
<scm>
<url>https://github.com/modelcontextprotocol/java-sdk</url>
<connection>git://github.com/modelcontextprotocol/java-sdk.git</connection>
<developerConnection>git@github.com/modelcontextprotocol/java-sdk.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp</artifactId>
<version>0.7.0-SNAPSHOT</version>
</dependency>
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-test</artifactId>
<version>0.7.0-SNAPSHOT</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>${springframework.version}</version>
</dependency>
<dependency>
<groupId>io.projectreactor.netty</groupId>
<artifactId>reactor-netty-http</artifactId>
<scope>test</scope>
</dependency>
<!-- The Spring Context is required due to the reactor-netty connector being dependant on
the Spring Lifecycle, as discussed here:
https://github.com/spring-projects/spring-framework/issues/31180 -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context</artifactId>
<version>${springframework.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-test</artifactId>
<version>${springframework.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
<version>${assert4j.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-api</artifactId>
<version>${junit.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>${mockito.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<version>${awaitility.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>ch.qos.logback</groupId>
<artifactId>logback-classic</artifactId>
<version>${logback.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-params</artifactId>
<version>${junit-jupiter.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View file

@ -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);
}
}

View file

@ -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);
}
}
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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);
}
}

View file

@ -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();
}
}
}

View file

@ -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();
}
}
}

View file

@ -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();
}
}

View 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>

View file

@ -0,0 +1,29 @@
# WebMVC SSE Server Transport
```xml
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-spring-webmvc</artifactId>
</dependency>
```
```java
String MESSAGE_ENDPOINT = "/mcp/message";
@Configuration
@EnableWebMvc
static class MyConfig {
@Bean
public WebMvcSseServerTransport webMvcSseServerTransport() {
return new WebMvcSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
}
@Bean
public RouterFunction<ServerResponse> routerFunction(WebMvcSseServerTransport transport) {
return transport.getRouterFunction();
}
}
```

View file

@ -0,0 +1,128 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-parent</artifactId>
<version>0.7.0-SNAPSHOT</version>
<relativePath>../../pom.xml</relativePath>
</parent>
<artifactId>mcp-spring-webmvc</artifactId>
<packaging>jar</packaging>
<name>Spring Web MVC implementation of the Java MCP SSE transport</name>
<description></description>
<url>https://github.com/modelcontextprotocol/java-sdk</url>
<scm>
<url>https://github.com/modelcontextprotocol/java-sdk</url>
<connection>git://github.com/modelcontextprotocol/java-sdk.git</connection>
<developerConnection>git@github.com/modelcontextprotocol/java-sdk.git</developerConnection>
</scm>
<dependencies>
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp</artifactId>
<version>0.7.0-SNAPSHOT</version>
</dependency>
<dependency>
<groupId>io.modelcontextprotocol.sdk</groupId>
<artifactId>mcp-test</artifactId>
<version>0.7.0-SNAPSHOT</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webmvc</artifactId>
<version>${springframework.version}</version>
</dependency>
<!-- The Spring Context is required due to the reactor-netty connector being dependant on
the Spring Lifecycle, as discussed here:
https://github.com/spring-projects/spring-framework/issues/31180 -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-context</artifactId>
<version>${springframework.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-test</artifactId>
<version>${springframework.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
<version>${assert4j.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter-api</artifactId>
<version>${junit.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<version>${mockito.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.awaitility</groupId>
<artifactId>awaitility</artifactId>
<version>${awaitility.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>ch.qos.logback</groupId>
<artifactId>logback-classic</artifactId>
<version>${logback.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.projectreactor.netty</groupId>
<artifactId>reactor-netty-http</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>jakarta.servlet</groupId>
<artifactId>jakarta.servlet-api</artifactId>
<version>${jakarta.servlet.version}</version>
<scope>provided</scope>
</dependency>
<dependency>
<groupId>org.apache.tomcat.embed</groupId>
<artifactId>tomcat-embed-core</artifactId>
<version>${tomcat.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View file

@ -0,0 +1,373 @@
/*
* Copyright 2024-2024 the original author or authors.
*/
package io.modelcontextprotocol.server.transport;
import java.io.IOException;
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.Mono;
import org.springframework.http.HttpStatus;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.RouterFunctions;
import org.springframework.web.servlet.function.ServerRequest;
import org.springframework.web.servlet.function.ServerResponse;
import org.springframework.web.servlet.function.ServerResponse.SseBuilder;
/**
* Server-side implementation of the Model Context Protocol (MCP) transport layer using
* HTTP with Server-Sent Events (SSE) through Spring WebMVC. This implementation provides
* a bridge between synchronous WebMVC operations and reactive programming patterns to
* maintain compatibility with the reactive transport interface.
*
* <p>
* Key features:
* <ul>
* <li>Implements bidirectional communication using HTTP POST for client-to-server
* messages and SSE for server-to-client messages</li>
* <li>Manages client sessions with unique IDs for reliable message delivery</li>
* <li>Supports graceful shutdown with proper session cleanup</li>
* <li>Provides JSON-RPC message handling through configured endpoints</li>
* <li>Includes built-in error handling and logging</li>
* </ul>
*
* <p>
* The transport operates on two main endpoints:
* <ul>
* <li>{@code /sse} - The SSE endpoint where clients establish their event stream
* connection</li>
* <li>A configurable message endpoint where clients send their JSON-RPC messages via HTTP
* POST</li>
* </ul>
*
* <p>
* This implementation uses {@link ConcurrentHashMap} to safely manage multiple client
* sessions in a thread-safe manner. Each client session is assigned a unique ID and
* maintains its own SSE connection.
*
* @author Christian Tzolov
* @author Alexandros Pappas
* @see ServerMcpTransport
* @see RouterFunction
*/
public class WebMvcSseServerTransport implements ServerMcpTransport {
private static final Logger logger = LoggerFactory.getLogger(WebMvcSseServerTransport.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<ServerResponse> 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;
/**
* The function to process incoming JSON-RPC messages and produce responses.
*/
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
/**
* Constructs a new WebMvcSseServerTransport instance.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* of messages.
* @param messageEndpoint The endpoint URI where clients should send their JSON-RPC
* messages via HTTP POST. This endpoint will be communicated to clients through the
* SSE connection's initial endpoint event.
* @throws IllegalArgumentException if either objectMapper or messageEndpoint is null
*/
public WebMvcSseServerTransport(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 WebMvcSseServerTransport instance with the default SSE endpoint.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* of messages.
* @param messageEndpoint The endpoint URI where clients should send their JSON-RPC
* messages via HTTP POST. This endpoint will be communicated to clients through the
* SSE connection's initial endpoint event.
* @throws IllegalArgumentException if either objectMapper or messageEndpoint is null
*/
public WebMvcSseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
this(objectMapper, messageEndpoint, DEFAULT_SSE_ENDPOINT);
}
/**
* Sets up the message handler for this transport. In the WebMVC SSE implementation,
* this method only stores the handler for later use, as connections are initiated by
* clients rather than the server.
* @param connectionHandler The function to process incoming JSON-RPC messages and
* produce responses
* @return An empty Mono since the server doesn't initiate connections
*/
@Override
public Mono<Void> connect(
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectionHandler) {
this.connectHandler = connectionHandler;
// Server-side transport doesn't initiate connections
return Mono.empty();
}
/**
* Broadcasts a message to all connected clients through their SSE connections. The
* message is serialized to JSON and sent as an SSE event with type "message". If any
* errors occur during sending to a particular client, they are logged but don't
* prevent sending to other clients.
* @param message The JSON-RPC message to broadcast to all connected clients
* @return A Mono that completes when the broadcast attempt is finished
*/
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
return Mono.fromRunnable(() -> {
if (sessions.isEmpty()) {
logger.debug("No active sessions to broadcast message to");
return;
}
try {
String jsonText = objectMapper.writeValueAsString(message);
logger.debug("Attempting to broadcast message to {} active sessions", sessions.size());
sessions.values().forEach(session -> {
try {
session.sseBuilder.id(session.id).event(MESSAGE_EVENT_TYPE).data(jsonText);
}
catch (Exception e) {
logger.error("Failed to send message to session {}: {}", session.id, e.getMessage());
session.sseBuilder.error(e);
}
});
}
catch (IOException e) {
logger.error("Failed to serialize message: {}", e.getMessage());
}
});
}
/**
* Handles new SSE connection requests from clients by creating a new session and
* establishing an SSE connection. This method:
* <ul>
* <li>Generates a unique session ID</li>
* <li>Creates a new ClientSession with an SSE builder</li>
* <li>Sends an initial endpoint event to inform the client where to send
* messages</li>
* <li>Maintains the session in the sessions map</li>
* </ul>
* @param request The incoming server request
* @return A ServerResponse configured for SSE communication, or an error response if
* the server is shutting down or the connection fails
*/
private ServerResponse handleSseConnection(ServerRequest request) {
if (this.isClosing) {
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
}
String sessionId = UUID.randomUUID().toString();
logger.debug("Creating new SSE connection for session: {}", sessionId);
// Send initial endpoint event
try {
return ServerResponse.sse(sseBuilder -> {
ClientSession session = new ClientSession(sessionId, sseBuilder);
this.sessions.put(sessionId, session);
try {
session.sseBuilder.id(session.id).event(ENDPOINT_EVENT_TYPE).data(messageEndpoint);
}
catch (Exception e) {
logger.error("Failed to poll event from session queue: {}", e.getMessage());
sseBuilder.error(e);
}
});
}
catch (Exception e) {
logger.error("Failed to send initial endpoint event to session {}: {}", sessionId, e.getMessage());
sessions.remove(sessionId);
return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).build();
}
}
/**
* Handles incoming JSON-RPC messages from clients. This method:
* <ul>
* <li>Deserializes the request body into a JSON-RPC message</li>
* <li>Processes the message through the configured connect handler</li>
* <li>Returns appropriate HTTP responses based on the processing result</li>
* </ul>
* @param request The incoming server request containing the JSON-RPC message
* @return A ServerResponse indicating success (200 OK) or appropriate error status
* with error details in case of failures
*/
private ServerResponse handleMessage(ServerRequest request) {
if (this.isClosing) {
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).body("Server is shutting down");
}
try {
String body = request.body(String.class);
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body);
// Convert the message to a Mono, apply the handler, and block for the
// response
@SuppressWarnings("unused")
McpSchema.JSONRPCMessage response = Mono.just(message).transform(connectHandler).block();
return ServerResponse.ok().build();
}
catch (IllegalArgumentException | IOException e) {
logger.error("Failed to deserialize message: {}", e.getMessage());
return ServerResponse.badRequest().body(new McpError("Invalid message format"));
}
catch (Exception e) {
logger.error("Error handling message: {}", e.getMessage());
return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).body(new McpError(e.getMessage()));
}
}
/**
* Represents an active client session with its associated SSE connection. Each
* session maintains:
* <ul>
* <li>A unique session identifier</li>
* <li>An SSE builder for sending server events to the client</li>
* <li>Logging of session lifecycle events</li>
* </ul>
*/
private static class ClientSession {
private final String id;
private final SseBuilder sseBuilder;
/**
* Creates a new client session with the specified ID and SSE builder.
* @param id The unique identifier for this session
* @param sseBuilder The SSE builder for sending server events to the client
*/
ClientSession(String id, SseBuilder sseBuilder) {
this.id = id;
this.sseBuilder = sseBuilder;
logger.debug("Session {} initialized with SSE emitter", id);
}
/**
* Closes this session by completing the SSE connection. Any errors during
* completion are logged but do not prevent the session from being marked as
* closed.
*/
void close() {
logger.debug("Closing session: {}", id);
try {
sseBuilder.complete();
logger.debug("Successfully completed SSE emitter for session {}", id);
}
catch (Exception e) {
logger.warn("Failed to complete SSE emitter for session {}: {}", id, e.getMessage());
// sseBuilder.error(e);
}
}
}
/**
* Converts data from one type to another using the configured ObjectMapper. This is
* particularly useful for handling complex JSON-RPC parameter types.
* @param data The source data object to convert
* @param typeRef The target type reference
* @return The converted object of type T
* @param <T> The target type
*/
@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:
* <ul>
* <li>Sets the closing flag to prevent new connections</li>
* <li>Closes all active SSE connections</li>
* <li>Removes all session records</li>
* </ul>
* @return A Mono that completes when all cleanup operations are finished
*/
@Override
public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> {
this.isClosing = true;
logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size());
sessions.values().forEach(session -> {
String sessionId = session.id;
session.close();
sessions.remove(sessionId);
});
logger.info("Graceful shutdown completed");
});
}
/**
* Returns the RouterFunction that defines the HTTP endpoints for this transport. The
* router function handles two endpoints:
* <ul>
* <li>GET /sse - For establishing SSE connections</li>
* <li>POST [messageEndpoint] - For receiving JSON-RPC messages from clients</li>
* </ul>
* @return The configured RouterFunction for handling HTTP requests
*/
public RouterFunction<ServerResponse> getRouterFunction() {
return this.routerFunction;
}
}

View file

@ -0,0 +1,117 @@
/*
* Copyright 2024-2024 the original author or authors.
*/
package io.modelcontextprotocol.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransport;
import io.modelcontextprotocol.spec.ServerMcpTransport;
import org.apache.catalina.Context;
import org.apache.catalina.LifecycleException;
import org.apache.catalina.startup.Tomcat;
import org.junit.jupiter.api.Timeout;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
import org.springframework.web.servlet.DispatcherServlet;
import org.springframework.web.servlet.config.annotation.EnableWebMvc;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerResponse;
@Timeout(15)
class WebMvcSseAsyncServerTransportTests extends AbstractMcpAsyncServerTests {
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private static final int PORT = 8181;
private Tomcat tomcat;
private WebMvcSseServerTransport transport;
@Configuration
@EnableWebMvc
static class TestConfig {
@Bean
public WebMvcSseServerTransport webMvcSseServerTransport() {
return new WebMvcSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
}
@Bean
public RouterFunction<ServerResponse> routerFunction(WebMvcSseServerTransport transport) {
return transport.getRouterFunction();
}
}
private AnnotationConfigWebApplicationContext appContext;
@Override
protected ServerMcpTransport createMcpTransport() {
// Set up Tomcat first
tomcat = new Tomcat();
tomcat.setPort(PORT);
// Set Tomcat base directory to java.io.tmpdir to avoid permission issues
String baseDir = System.getProperty("java.io.tmpdir");
tomcat.setBaseDir(baseDir);
// Use the same directory for document base
Context context = tomcat.addContext("", baseDir);
// Create and configure Spring WebMvc context
appContext = new AnnotationConfigWebApplicationContext();
appContext.register(TestConfig.class);
appContext.setServletContext(context.getServletContext());
appContext.refresh();
// Get the transport from Spring context
transport = appContext.getBean(WebMvcSseServerTransport.class);
// Create DispatcherServlet with our Spring context
DispatcherServlet dispatcherServlet = new DispatcherServlet(appContext);
// dispatcherServlet.setThrowExceptionIfNoHandlerFound(true);
// Add servlet to Tomcat and get the wrapper
var wrapper = Tomcat.addServlet(context, "dispatcherServlet", dispatcherServlet);
wrapper.setLoadOnStartup(1);
context.addServletMappingDecoded("/*", "dispatcherServlet");
try {
tomcat.start();
tomcat.getConnector(); // Create and start the connector
}
catch (LifecycleException e) {
throw new RuntimeException("Failed to start Tomcat", e);
}
return transport;
}
@Override
protected void onStart() {
}
@Override
protected void onClose() {
if (transport != null) {
transport.closeGracefully().block();
}
if (appContext != null) {
appContext.close();
}
if (tomcat != null) {
try {
tomcat.stop();
tomcat.destroy();
}
catch (LifecycleException e) {
throw new RuntimeException("Failed to stop Tomcat", e);
}
}
}
}

View file

@ -0,0 +1,507 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*/
package io.modelcontextprotocol.server;
import java.time.Duration;
import java.util.List;
import java.util.Map;
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.server.transport.WebMvcSseServerTransport;
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.apache.catalina.Context;
import org.apache.catalina.LifecycleException;
import org.apache.catalina.LifecycleState;
import org.apache.catalina.startup.Tomcat;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import reactor.test.StepVerifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.client.RestClient;
import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
import org.springframework.web.servlet.DispatcherServlet;
import org.springframework.web.servlet.config.annotation.EnableWebMvc;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerResponse;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await;
public class WebMvcSseIntegrationTests {
private static final int PORT = 8183;
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private WebMvcSseServerTransport mcpServerTransport;
McpClient.SyncSpec clientBuilder;
@Configuration
@EnableWebMvc
static class TestConfig {
@Bean
public WebMvcSseServerTransport webMvcSseServerTransport() {
return new WebMvcSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
}
@Bean
public RouterFunction<ServerResponse> routerFunction(WebMvcSseServerTransport transport) {
return transport.getRouterFunction();
}
}
private Tomcat tomcat;
private AnnotationConfigWebApplicationContext appContext;
@BeforeEach
public void before() {
// Set up Tomcat first
tomcat = new Tomcat();
tomcat.setPort(PORT);
// Set Tomcat base directory to java.io.tmpdir to avoid permission issues
String baseDir = System.getProperty("java.io.tmpdir");
tomcat.setBaseDir(baseDir);
// Use the same directory for document base
Context context = tomcat.addContext("", baseDir);
// Create and configure Spring WebMvc context
appContext = new AnnotationConfigWebApplicationContext();
appContext.register(TestConfig.class);
appContext.setServletContext(context.getServletContext());
appContext.refresh();
// Get the transport from Spring context
mcpServerTransport = appContext.getBean(WebMvcSseServerTransport.class);
// Create DispatcherServlet with our Spring context
DispatcherServlet dispatcherServlet = new DispatcherServlet(appContext);
// dispatcherServlet.setThrowExceptionIfNoHandlerFound(true);
// Add servlet to Tomcat and get the wrapper
var wrapper = Tomcat.addServlet(context, "dispatcherServlet", dispatcherServlet);
wrapper.setLoadOnStartup(1);
wrapper.setAsyncSupported(true);
context.addServletMappingDecoded("/*", "dispatcherServlet");
try {
// Configure and start the connector with async support
var connector = tomcat.getConnector();
connector.setAsyncTimeout(3000); // 3 seconds timeout for async requests
tomcat.start();
assertThat(tomcat.getServer().getState() == LifecycleState.STARTED);
}
catch (Exception e) {
throw new RuntimeException("Failed to start Tomcat", e);
}
this.clientBuilder = McpClient.sync(new HttpClientSseClientTransport("http://localhost:" + PORT));
}
@AfterEach
public void after() {
if (mcpServerTransport != null) {
mcpServerTransport.closeGracefully().block();
}
if (appContext != null) {
appContext.close();
}
if (tomcat != null) {
try {
tomcat.stop();
tomcat.destroy();
}
catch (LifecycleException e) {
throw new RuntimeException("Failed to stop Tomcat", e);
}
}
}
// ---------------------------------------
// 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!");
});
}
@Test
void testCreateMessageWithoutSamplingCapabilities() {
var mcpAsyncServer = McpServer.async(mcpServerTransport).serverInfo("test-server", "1.0.0").build();
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");
});
}
@Test
void testCreateMessageSuccess() throws InterruptedException {
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
// ---------------------------------------
@Test
void testRootsSuccess() {
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();
}
@Test
void testRootsWithoutCapability() {
var mcpServer = McpServer.sync(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> {
}).build();
// Create client without roots capability
// No roots capability
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).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();
}
@Test
void testRootsWithEmptyRootsList() {
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();
}
@Test
void testRootsWithMultipleConsumers() {
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();
}
@Test
void testRootsServerCloseWithActiveSubscription() {
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": {}
}
""";
@Test
void testToolCallSuccess() {
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();
}
@Test
void testToolListChangeHandlingSuccess() {
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();
}
@Test
void testInitialize() {
var mcpServer = McpServer.sync(mcpServerTransport).build();
var mcpClient = clientBuilder.build();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.close();
mcpServer.close();
}
}

View file

@ -0,0 +1,117 @@
/*
* Copyright 2024-2024 the original author or authors.
*/
package io.modelcontextprotocol.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransport;
import io.modelcontextprotocol.spec.ServerMcpTransport;
import org.apache.catalina.Context;
import org.apache.catalina.LifecycleException;
import org.apache.catalina.startup.Tomcat;
import org.junit.jupiter.api.Timeout;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.context.support.AnnotationConfigWebApplicationContext;
import org.springframework.web.servlet.DispatcherServlet;
import org.springframework.web.servlet.config.annotation.EnableWebMvc;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerResponse;
@Timeout(15)
class WebMvcSseSyncServerTransportTests extends AbstractMcpSyncServerTests {
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private static final int PORT = 8181;
private Tomcat tomcat;
private WebMvcSseServerTransport transport;
@Configuration
@EnableWebMvc
static class TestConfig {
@Bean
public WebMvcSseServerTransport webMvcSseServerTransport() {
return new WebMvcSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
}
@Bean
public RouterFunction<ServerResponse> routerFunction(WebMvcSseServerTransport transport) {
return transport.getRouterFunction();
}
}
private AnnotationConfigWebApplicationContext appContext;
@Override
protected ServerMcpTransport createMcpTransport() {
// Set up Tomcat first
tomcat = new Tomcat();
tomcat.setPort(PORT);
// Set Tomcat base directory to java.io.tmpdir to avoid permission issues
String baseDir = System.getProperty("java.io.tmpdir");
tomcat.setBaseDir(baseDir);
// Use the same directory for document base
Context context = tomcat.addContext("", baseDir);
// Create and configure Spring WebMvc context
appContext = new AnnotationConfigWebApplicationContext();
appContext.register(TestConfig.class);
appContext.setServletContext(context.getServletContext());
appContext.refresh();
// Get the transport from Spring context
transport = appContext.getBean(WebMvcSseServerTransport.class);
// Create DispatcherServlet with our Spring context
DispatcherServlet dispatcherServlet = new DispatcherServlet(appContext);
// dispatcherServlet.setThrowExceptionIfNoHandlerFound(true);
// Add servlet to Tomcat and get the wrapper
var wrapper = Tomcat.addServlet(context, "dispatcherServlet", dispatcherServlet);
wrapper.setLoadOnStartup(1);
context.addServletMappingDecoded("/*", "dispatcherServlet");
try {
tomcat.start();
tomcat.getConnector(); // Create and start the connector
}
catch (LifecycleException e) {
throw new RuntimeException("Failed to start Tomcat", e);
}
return transport;
}
@Override
protected void onStart() {
}
@Override
protected void onClose() {
if (transport != null) {
transport.closeGracefully().block();
}
if (appContext != null) {
appContext.close();
}
if (tomcat != null) {
try {
tomcat.stop();
tomcat.destroy();
}
catch (LifecycleException e) {
throw new RuntimeException("Failed to stop Tomcat", e);
}
}
}
}

View file

@ -0,0 +1,27 @@
<?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="io.modelcontextprotocol" level="DEBUG"/>
<!-- Client packages -->
<logger name="io.modelcontextprotocol.client" level="DEBUG"/>
<!-- Server transport package -->
<logger name="io.modelcontextprotocol.server.transport" level="DEBUG"/>
<!-- Spec package -->
<logger name="io.modelcontextprotocol.spec" level="DEBUG"/>
<!-- Root logger -->
<root level="INFO">
<appender-ref ref="CONSOLE"/>
</root>
</configuration>