docs: improve transport layer documentation and code cleanup

- Add comprehensive Javadoc for SSE transport implementations
- Document key features and implementation details for WebFlux/WebMVC transports
- Refactor FlowSseClient to use records for better immutability
- Fix typo in prompt property name
- Remove unused code and improve readability
This commit is contained in:
Christian Tzolov 2025-01-12 08:46:39 +01:00
commit 299eee944a
3 changed files with 511 additions and 111 deletions

View file

@ -61,10 +61,9 @@ import org.springframework.web.reactive.function.client.WebClient;
* <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 handles automatic reconnection for transient failures and provides
* graceful shutdown capabilities. It uses {@link WebClient} for HTTP communications and
* supports JSON serialization/deserialization of messages.
*
* This implementation uses {@link WebClient} for HTTP communications and supports JSON
* serialization/deserialization of messages.
*
* @author Christian Tzolov
* @see <a href=
@ -157,6 +156,28 @@ public class WebFluxSseClientTransport implements ClientMcpTransport {
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();
@ -190,6 +211,18 @@ public class WebFluxSseClientTransport implements ClientMcpTransport {
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
@ -281,6 +314,20 @@ public class WebFluxSseClientTransport implements ClientMcpTransport {
.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

@ -28,10 +28,37 @@ import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.reactive.function.server.ServerResponse;
/**
* Server-side implementation of the MCP HTTP with SSE transport specification. Provides
* endpoints for SSE connections and message handling.
* 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
* @see ServerMcpTransport
* @see ServerSentEvent
*/
public class WebFluxSseServerTransport implements ServerMcpTransport {
@ -71,9 +98,13 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
/**
* Constructs a new SseServerTransport.
* Constructs a new WebFlux SSE server transport instance.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* @param messageEndpoint The endpoint URI where clients should send messages
* 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) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
@ -87,6 +118,16 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
.build();
}
/**
* 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;
@ -94,6 +135,23 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
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()) {
@ -133,11 +191,35 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
});
}
/**
* 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(() -> {
@ -155,14 +237,36 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
}
/**
* Get the router function for configuring the web server.
* 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 {SSE_ENDPOINT} - 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.
* 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) {
@ -208,7 +312,19 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
}
/**
* Handles incoming messages from clients.
* 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) {
@ -239,7 +355,16 @@ public class WebFluxSseServerTransport implements ServerMcpTransport {
}
/**
* Represents an active client session.
* 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 {

View file

@ -39,11 +39,39 @@ import org.springframework.web.servlet.function.ServerResponse;
import org.springframework.web.servlet.function.ServerResponse.SseBuilder;
/**
* Server-side implementation of the MCP HTTP with SSE transport specification using
* Spring WebMVC. This implementation wraps synchronous WebMVC operations in reactive
* types to maintain compatibility with the reactive transport interface.
* 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
* @see ServerMcpTransport
* @see RouterFunction
*/
public class WebMvcSseServerTransport implements ServerMcpTransport {
@ -80,12 +108,19 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
*/
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 FunctionalSseServerTransport.
* Constructs a new WebMvcSseServerTransport instance.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* @param messageEndpoint The endpoint URI where clients should send messages
* 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) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
@ -99,13 +134,30 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
.build();
}
/**
* 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>> handler) {
this.connectHandler = handler;
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(() -> {
@ -135,7 +187,18 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
}
/**
* Handles new SSE connection requests from clients.
* 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) {
@ -169,7 +232,15 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
}
/**
* Handles incoming messages from clients.
* 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) {
@ -197,7 +268,13 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
}
/**
* Represents an active client session.
* 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 {
@ -205,12 +282,22 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
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 {
@ -225,11 +312,28 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
}
/**
* 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(() -> {
@ -247,7 +351,13 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
}
/**
* Get the router function for configuring the web server.
* 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;