From 299eee944a5e37841379fe6165a5fe3d4426dca3 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sun, 12 Jan 2025 08:46:39 +0100 Subject: [PATCH] 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 --- .../transport/WebFluxSseClientTransport.java | 55 +++++- .../transport/WebFluxSseServerTransport.java | 141 ++++++++++++++- .../transport/WebMvcSseServerTransport.java | 132 ++++++++++++-- .../mcp/client/transport/FlowSseClient.java | 170 +++++++++++------- .../HttpClientSseClientTransport.java | 100 ++++++++++- .../ai/mcp/server/McpAsyncServer.java | 10 +- .../ai/mcp/server/McpServer.java | 6 +- 7 files changed, 507 insertions(+), 107 deletions(-) diff --git a/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/client/transport/WebFluxSseClientTransport.java b/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/client/transport/WebFluxSseClientTransport.java index ed24f6f..0505372 100644 --- a/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/client/transport/WebFluxSseClientTransport.java +++ b/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/client/transport/WebFluxSseClientTransport.java @@ -61,10 +61,9 @@ import org.springframework.web.reactive.function.client.WebClient; *
  • The client establishes an SSE connection to the server's /sse endpoint
    • *
  • The server sends an 'endpoint' event containing the URI for sending messages
    • * - * - * 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 + * The connection process follows these steps: + *
      + *
    1. Establishes an SSE connection to the server's /sse endpoint
      2. + *
    2. Waits for the server to send an 'endpoint' event with the message posting + * URI
      3. + *
    3. Sets up message handling for incoming JSON-RPC messages
      4. + *
    + * + *

    + * 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 connect(Function, Mono> handler) { Flux> 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. + * + *

    + * 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 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. + * + *

    + * 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 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 unmarshalFrom(Object data, TypeReference typeRef) { return this.objectMapper.convertValue(data, typeRef); diff --git a/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebFluxSseServerTransport.java b/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebFluxSseServerTransport.java index 0159258..88812ad 100644 --- a/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebFluxSseServerTransport.java +++ b/mcp-transport/mcp-webflux-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebFluxSseServerTransport.java @@ -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. + * + *

    + * Key features: + *

    + * + *

    + * The transport sets up two main endpoints: + *

    + * + *

    + * 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> 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 connect(Function, Mono> 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. + * + *

    + * The method: + *

    + * @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 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 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 unmarshalFrom(Object data, TypeReference 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. + * + *

    + * The shutdown process: + *

      + *
    • Marks the transport as closing to prevent new connections
      • + *
    • Closes each active session
      • + *
    • Removes closed sessions from the sessions map
      • + *
    • Times out after 5 seconds if shutdown takes too long
      • + *
    + * @return A Mono that completes when all sessions have been closed + */ @Override public Mono 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. + * + *

    + * The router function defines two endpoints: + *

      + *
    • GET {SSE_ENDPOINT} - For establishing SSE connections
      • + *
    • POST {messageEndpoint} - For receiving client messages
      • + *
    + * @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. + * + *

    + * The handler performs the following steps: + *

      + *
    • Generates a unique session ID
      • + *
    • Creates a new ClientSession instance
      • + *
    • Sends the message endpoint URI as an initial event
      • + *
    • Sets up message forwarding for the session
      • + *
    • Handles connection cleanup on completion or errors
      • + *
    + * @param request The incoming server request + * @return A response with the SSE event stream */ private Mono 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. + * + *

    + * The handler: + *

      + *
    • Deserializes the incoming JSON-RPC message
      • + *
    • Passes it through the message handler chain
      • + *
    • Returns appropriate HTTP responses based on processing results
      • + *
    • Handles various error conditions with appropriate error responses
      • + *
    + * @param request The incoming server request containing the JSON-RPC message + * @return A response indicating the message processing result */ private Mono 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. + * + *

    + * Each session: + *

      + *
    • Has a unique identifier
      • + *
    • Maintains its own message sink for event broadcasting
      • + *
    • Supports clean shutdown through the close method
      • + *
    */ private static class ClientSession { diff --git a/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java b/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java index cedd4b6..cc01f36 100644 --- a/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java +++ b/mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java @@ -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. + * + *

    + * Key features: + *

      + *
    • Implements bidirectional communication using HTTP POST for client-to-server + * messages and SSE for server-to-client messages
      • + *
    • Manages client sessions with unique IDs for reliable message delivery
      • + *
    • Supports graceful shutdown with proper session cleanup
      • + *
    • Provides JSON-RPC message handling through configured endpoints
      • + *
    • Includes built-in error handling and logging
      • + *
    + * + *

    + * The transport operates on two main endpoints: + *

      + *
    • {@code /sse} - The SSE endpoint where clients establish their event stream + * connection
      • + *
    • A configurable message endpoint where clients send their JSON-RPC messages via HTTP + * POST
      • + *
    + * + *

    + * 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> 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 connect(Function, Mono> handler) { - this.connectHandler = handler; + public Mono connect( + Function, Mono> 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 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: + *

      + *
    • Generates a unique session ID
      • + *
    • Creates a new ClientSession with an SSE builder
      • + *
    • Sends an initial endpoint event to inform the client where to send + * messages
      • + *
    • Maintains the session in the sessions map
      • + *
    + * @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: + *
      + *
    • Deserializes the request body into a JSON-RPC message
      • + *
    • Processes the message through the configured connect handler
      • + *
    • Returns appropriate HTTP responses based on the processing result
      • + *
    + * @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: + *
      + *
    • A unique session identifier
      • + *
    • An SSE builder for sending server events to the client
      • + *
    • Logging of session lifecycle events
      • + *
    */ 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 The target type + */ @Override public T unmarshalFrom(Object data, TypeReference typeRef) { return this.objectMapper.convertValue(data, typeRef); } + /** + * Initiates a graceful shutdown of the transport. This method: + *
      + *
    • Sets the closing flag to prevent new connections
      • + *
    • Closes all active SSE connections
      • + *
    • Removes all session records
      • + *
    + * @return A Mono that completes when all cleanup operations are finished + */ @Override public Mono 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: + *
      + *
    • GET /sse - For establishing SSE connections
      • + *
    • POST [messageEndpoint] - For receiving JSON-RPC messages from clients
      • + *
    + * @return The configured RouterFunction for handling HTTP requests */ public RouterFunction getRouterFunction() { return this.routerFunction; diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/FlowSseClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/FlowSseClient.java index c337e5d..3fa932a 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/FlowSseClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/FlowSseClient.java @@ -19,7 +19,6 @@ import java.net.URI; import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; -import java.time.Duration; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Flow; import java.util.concurrent.atomic.AtomicReference; @@ -27,23 +26,101 @@ import java.util.function.Function; import java.util.regex.Pattern; /** + * A Server-Sent Events (SSE) client implementation using Java's Flow API for reactive + * stream processing. This client establishes a connection to an SSE endpoint and + * processes the incoming event stream, parsing SSE-formatted messages into structured + * events. + * + *

    + * The client supports standard SSE event fields including: + *

      + *
    • event - The event type (defaults to "message" if not specified)
      • + *
    • id - The event ID
      • + *
    • data - The event payload data
      • + *
    + * + *

    + * Events are delivered to a provided {@link SseEventHandler} which can process events and + * handle any errors that occur during the connection. + * * @author Christian Tzolov + * @see SseEventHandler + * @see SseEvent */ public class FlowSseClient { private final HttpClient httpClient; + /** + * Pattern to extract the data content from SSE data field lines. Matches lines + * starting with "data:" and captures the remaining content. + */ private static final Pattern EVENT_DATA_PATTERN = Pattern.compile("^data:(.+)$", Pattern.MULTILINE); + /** + * Pattern to extract the event ID from SSE id field lines. Matches lines starting + * with "id:" and captures the ID value. + */ private static final Pattern EVENT_ID_PATTERN = Pattern.compile("^id:(.+)$", Pattern.MULTILINE); + /** + * Pattern to extract the event type from SSE event field lines. Matches lines + * starting with "event:" and captures the event type. + */ private static final Pattern EVENT_TYPE_PATTERN = Pattern.compile("^event:(.+)$", Pattern.MULTILINE); + /** + * Record class representing a Server-Sent Event with its standard fields. + * + * @param id the event ID (may be null) + * @param type the event type (defaults to "message" if not specified in the stream) + * @param data the event payload data + */ + public static record SseEvent(String id, String type, String data) { + } + + /** + * Interface for handling SSE events and errors. Implementations can process received + * events and handle any errors that occur during the SSE connection. + */ + public interface SseEventHandler { + + /** + * Called when an SSE event is received. + * @param event the received SSE event containing id, type, and data + */ + void onEvent(SseEvent event); + + /** + * Called when an error occurs during the SSE connection. + * @param error the error that occurred + */ + void onError(Throwable error); + + } + + /** + * Creates a new FlowSseClient with the specified HTTP client. + * @param httpClient the {@link HttpClient} instance to use for SSE connections + */ public FlowSseClient(HttpClient httpClient) { this.httpClient = httpClient; } - public void subscribe(String url, SSEEventHandler eventHandler) { + /** + * Subscribes to an SSE endpoint and processes the event stream. + * + *

    + * This method establishes a connection to the specified URL and begins processing the + * SSE stream. Events are parsed and delivered to the provided event handler. The + * connection remains active until either an error occurs or the server closes the + * connection. + * @param url the SSE endpoint URL to connect to + * @param eventHandler the handler that will receive SSE events and error + * notifications + * @throws RuntimeException if the connection fails with a non-200 status code + */ + public void subscribe(String url, SseEventHandler eventHandler) { HttpRequest request = HttpRequest.newBuilder() .uri(URI.create(url)) .header("Accept", "text/event-stream") @@ -70,7 +147,7 @@ public class FlowSseClient { // Empty line means end of event if (eventBuilder.length() > 0) { String eventData = eventBuilder.toString(); - SSEEvent event = new SSEEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); + SseEvent event = new SseEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); eventHandler.onEvent(event); eventBuilder.setLength(0); } @@ -108,7 +185,7 @@ public class FlowSseClient { // Handle any remaining event data if (eventBuilder.length() > 0) { String eventData = eventBuilder.toString(); - SSEEvent event = new SSEEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); + SseEvent event = new SseEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); eventHandler.onEvent(event); } } @@ -130,73 +207,30 @@ public class FlowSseClient { }); } - // Event handler interface - public interface SSEEventHandler { + // public static void main(String[] args) { - void onEvent(SSEEvent event); + // FlowSseClient sseClient = new FlowSseClient( + // HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build()); - void onError(Throwable error); + // var eventHandler = new SSEEventHandler() { + // @Override + // public void onEvent(SSEEvent event) { + // System.out.println("Received event: " + event); + // } - } + // @Override + // public void onError(Throwable error) { + // System.out.println("Error: " + error); + // } + // }; + // sseClient.subscribe("http://localhost:8080/sse", eventHandler); - // SSE Event class - public static class SSEEvent { - - private final String id; - - private final String type; - - private final String data; - - public SSEEvent(String id, String type, String data) { - this.id = id; - this.type = type; - this.data = data; - } - - public String getId() { - return id; - } - - public String getType() { - return type; - } - - public String getData() { - return data; - } - - @Override - public String toString() { - return "SSEEvent{" + "id='" + id + '\'' + ", type='" + type + '\'' + ", data='" + data + '\'' + '}'; - } - - } - - public static void main(String[] args) { - - FlowSseClient sseClient = new FlowSseClient( - HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build()); - - var eventHandler = new SSEEventHandler() { - @Override - public void onEvent(SSEEvent event) { - System.out.println("Received event: " + event); - } - - @Override - public void onError(Throwable error) { - System.out.println("Error: " + error); - } - }; - sseClient.subscribe("http://localhost:8080/sse", eventHandler); - - try { - Thread.sleep(100000); - } - catch (InterruptedException e) { - e.printStackTrace(); - } - } + // try { + // Thread.sleep(100000); + // } + // catch (InterruptedException e) { + // e.printStackTrace(); + // } + // } } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java index 3f457a4..9284806 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java @@ -20,7 +20,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import org.springframework.ai.mcp.client.transport.FlowSseClient.SSEEvent; +import org.springframework.ai.mcp.client.transport.FlowSseClient.SseEvent; import org.springframework.ai.mcp.spec.ClientMcpTransport; import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; @@ -46,40 +46,86 @@ import java.util.function.Function; * {@link org.springframework.ai.mcp.spec.McpTransport} that follows the MCP HTTP with SSE * transport specification, using Java's HttpClient. * + *

    + * This transport implementation establishes a bidirectional communication channel between + * client and server using SSE for server-to-client messages and HTTP POST requests for + * client-to-server messages. The transport: + *

      + *
    • Establishes an SSE connection to receive server messages
      • + *
    • Handles endpoint discovery through SSE events
      • + *
    • Manages message serialization/deserialization using Jackson
      • + *
    • Provides graceful connection termination
      • + *
    + * + *

    + * The transport supports two types of SSE events: + *

      + *
    • 'endpoint' - Contains the URL for sending client messages
      • + *
    • 'message' - Contains JSON-RPC message payload
      • + *
    + * * @author Christian Tzolov + * @see org.springframework.ai.mcp.spec.McpTransport + * @see org.springframework.ai.mcp.spec.ClientMcpTransport */ public class HttpClientSseClientTransport implements ClientMcpTransport { private final static Logger logger = LoggerFactory.getLogger(HttpClientSseClientTransport.class); + /** SSE event type for JSON-RPC messages */ private final static String MESSAGE_EVENT_TYPE = "message"; + /** SSE event type for endpoint discovery */ private final static String ENDPOINT_EVENT_TYPE = "endpoint"; + /** Default SSE endpoint path */ private final static String SSE_ENDPOINT = "/sse"; + /** Base URI for the MCP server */ private final String baseUri; + /** SSE client for handling server-sent events. Uses the /sse endpoint */ private final FlowSseClient sseClient; + /** + * HTTP client for sending messages to the server. Uses HTTP POST over the message + * endpoint + */ private final HttpClient httpClient; + /** JSON object mapper for message serialization/deserialization */ protected ObjectMapper objectMapper; + /** Flag indicating if the transport is in closing state */ private volatile boolean isClosing = false; + /** Latch for coordinating endpoint discovery */ private CountDownLatch closeLatch = new CountDownLatch(1); + /** Holds the discovered message endpoint URL */ private final AtomicReference messageEndpoint = new AtomicReference<>(); + /** Holds the SSE connection future */ private final AtomicReference> connectionFuture = new AtomicReference<>(); + /** + * Creates a new transport instance with default HTTP client and object mapper. + * @param baseUri the base URI of the MCP server + */ public HttpClientSseClientTransport(String baseUri) { this(HttpClient.newBuilder(), baseUri, new ObjectMapper()); } + /** + * Creates a new transport instance with custom HTTP client builder and object mapper. + * @param clientBuilder the HTTP client builder to use + * @param baseUri the base URI of the MCP server + * @param objectMapper the object mapper for JSON serialization/deserialization + * @throws IllegalArgumentException if objectMapper or clientBuilder is null + */ public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, String baseUri, ObjectMapper objectMapper) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.hasText(baseUri, "baseUri must not be empty"); Assert.notNull(clientBuilder, "clientBuilder must not be null"); this.baseUri = baseUri; this.objectMapper = objectMapper; @@ -87,31 +133,44 @@ public class HttpClientSseClientTransport implements ClientMcpTransport { this.sseClient = new FlowSseClient(this.httpClient); } + /** + * Establishes the SSE connection with the server and sets up message handling. + * + *

    + * This method: + *

      + *
    • Initiates the SSE connection
      • + *
    • Handles endpoint discovery events
      • + *
    • Processes incoming JSON-RPC messages
      • + *
    + * @param handler the function to process received JSON-RPC messages + * @return a Mono that completes when the connection is established + */ @Override public Mono connect(Function, Mono> handler) { CompletableFuture future = new CompletableFuture<>(); connectionFuture.set(future); - sseClient.subscribe(this.baseUri + SSE_ENDPOINT, new FlowSseClient.SSEEventHandler() { + sseClient.subscribe(this.baseUri + SSE_ENDPOINT, new FlowSseClient.SseEventHandler() { @Override - public void onEvent(SSEEvent event) { + public void onEvent(SseEvent event) { if (isClosing) { return; } try { - if (ENDPOINT_EVENT_TYPE.equals(event.getType())) { - String endpoint = event.getData(); + if (ENDPOINT_EVENT_TYPE.equals(event.type())) { + String endpoint = event.data(); messageEndpoint.set(endpoint); closeLatch.countDown(); future.complete(null); } - else if (MESSAGE_EVENT_TYPE.equals(event.getType())) { - JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.getData()); + else if (MESSAGE_EVENT_TYPE.equals(event.type())) { + JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.data()); handler.apply(Mono.just(message)).subscribe(); } else { - logger.error("Received unrecognized SSE event type: " + event.getType()); + logger.error("Received unrecognized SSE event type: " + event.type()); } } catch (IOException e) { @@ -132,6 +191,16 @@ public class HttpClientSseClientTransport implements ClientMcpTransport { return Mono.fromFuture(future); } + /** + * Sends a JSON-RPC message to the server. + * + *

    + * This method waits for the message endpoint to be discovered before sending the + * message. The message is serialized to JSON and sent as an HTTP POST request. + * @param message the JSON-RPC message to send + * @return a Mono that completes when the message is sent + * @throws McpError if the message endpoint is not available or the wait times out + */ @Override public Mono sendMessage(JSONRPCMessage message) { if (isClosing) { @@ -175,6 +244,14 @@ public class HttpClientSseClientTransport implements ClientMcpTransport { } } + /** + * Gracefully closes the transport connection. + * + *

    + * Sets the closing flag and cancels any pending connection future. This prevents new + * messages from being sent and allows ongoing operations to complete. + * @return a Mono that completes when the closing process is initiated + */ @Override public Mono closeGracefully() { return Mono.fromRunnable(() -> { @@ -186,6 +263,13 @@ public class HttpClientSseClientTransport implements ClientMcpTransport { }); } + /** + * Unmarshals data to the specified type using the configured object mapper. + * @param data the data to unmarshal + * @param typeRef the type reference for the target type + * @param the target type + * @return the unmarshalled object + */ @Override public T unmarshalFrom(Object data, TypeReference typeRef) { return this.objectMapper.convertValue(data, typeRef); diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java index 2a87b3a..fa8c512 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java @@ -487,14 +487,14 @@ public class McpAsyncServer { return Mono.error(new McpError("Server must be configured with prompt capabilities")); } - if (this.prompts.containsKey(promptRegistration.propmpt().name())) { + if (this.prompts.containsKey(promptRegistration.prompt().name())) { return Mono - .error(new McpError("Prompt with name '" + promptRegistration.propmpt().name() + "' already exists")); + .error(new McpError("Prompt with name '" + promptRegistration.prompt().name() + "' already exists")); } - this.prompts.put(promptRegistration.propmpt().name(), promptRegistration); + this.prompts.put(promptRegistration.prompt().name(), promptRegistration); - logger.info("Added prompt handler: {}", promptRegistration.propmpt().name()); + logger.info("Added prompt handler: {}", promptRegistration.prompt().name()); // Servers that declared the listChanged capability SHOULD send a notification, // when the list of available prompts changes @@ -546,7 +546,7 @@ public class McpAsyncServer { // new TypeReference() { // }); - var promptList = this.prompts.values().stream().map(PromptRegistration::propmpt).toList(); + var promptList = this.prompts.values().stream().map(PromptRegistration::prompt).toList(); return Mono.just(new McpSchema.ListPromptsResult(promptList, null)); }; diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java index 6ee01a3..362491b 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java @@ -190,7 +190,7 @@ public interface McpServer { * @param promptHandler The function that processes prompt requests and returns * formatted templates */ - public static record PromptRegistration(McpSchema.Prompt propmpt, + public static record PromptRegistration(McpSchema.Prompt prompt, Function promptHandler) { } @@ -498,7 +498,7 @@ public interface McpServer { */ public Builder prompts(List prompts) { for (PromptRegistration prompt : prompts) { - this.prompts.put(prompt.propmpt().name(), prompt); + this.prompts.put(prompt.prompt().name(), prompt); } return this; } @@ -521,7 +521,7 @@ public interface McpServer { */ public Builder prompts(PromptRegistration... prompts) { for (PromptRegistration prompt : prompts) { - this.prompts.put(prompt.propmpt().name(), prompt); + this.prompts.put(prompt.prompt().name(), prompt); } return this; }