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
4 changed files with 511 additions and 111 deletions

View file

@ -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.
*
* <p>
* The client supports standard SSE event fields including:
* <ul>
* <li>event - The event type (defaults to "message" if not specified)</li>
* <li>id - The event ID</li>
* <li>data - The event payload data</li>
* </ul>
*
* <p>
* 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.
*
* <p>
* 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();
// }
// }
}

View file

@ -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.
*
* <p>
* 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:
* <ul>
* <li>Establishes an SSE connection to receive server messages</li>
* <li>Handles endpoint discovery through SSE events</li>
* <li>Manages message serialization/deserialization using Jackson</li>
* <li>Provides graceful connection termination</li>
* </ul>
*
* <p>
* The transport supports two types of SSE events:
* <ul>
* <li>'endpoint' - Contains the URL for sending client messages</li>
* <li>'message' - Contains JSON-RPC message payload</li>
* </ul>
*
* @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<String> messageEndpoint = new AtomicReference<>();
/** Holds the SSE connection future */
private final AtomicReference<CompletableFuture<Void>> 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.
*
* <p>
* This method:
* <ul>
* <li>Initiates the SSE connection</li>
* <li>Handles endpoint discovery events</li>
* <li>Processes incoming JSON-RPC messages</li>
* </ul>
* @param handler the function to process received JSON-RPC messages
* @return a Mono that completes when the connection is established
*/
@Override
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
CompletableFuture<Void> 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.
*
* <p>
* 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<Void> sendMessage(JSONRPCMessage message) {
if (isClosing) {
@ -175,6 +244,14 @@ public class HttpClientSseClientTransport implements ClientMcpTransport {
}
}
/**
* Gracefully closes the transport connection.
*
* <p>
* 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<Void> 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 <T> the target type
* @return the unmarshalled object
*/
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return this.objectMapper.convertValue(data, typeRef);

View file

@ -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<McpSchema.PaginatedRequest>() {
// });
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));
};

View file

@ -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<McpSchema.GetPromptRequest, McpSchema.GetPromptResult> promptHandler) {
}
@ -498,7 +498,7 @@ public interface McpServer {
*/
public Builder prompts(List<PromptRegistration> 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;
}