Enhances WebFlux SSE transport implementation with customizable endpoint paths: - Add configurable SSE endpoint support in both client and server transports - Update tests to verify custom SSE endpoint functionality - Implement builder pattern to support the new configuration options Co-authored-by: haidao <wwulfric@gmail.com> Co-authored-by: Harry <34418180+HarryFQG@users.noreply.github.com> Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
parent
2934635d99
commit
55ee15604a
8 changed files with 210 additions and 19 deletions
|
|
@ -79,7 +79,7 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
||||||
* Default SSE endpoint path as specified by the MCP transport specification. This
|
* Default SSE endpoint path as specified by the MCP transport specification. This
|
||||||
* endpoint is used to establish the SSE connection with the server.
|
* endpoint is used to establish the SSE connection with the server.
|
||||||
*/
|
*/
|
||||||
private static final String SSE_ENDPOINT = "/sse";
|
private static final String DEFAULT_SSE_ENDPOINT = "/sse";
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Type reference for parsing SSE events containing string data.
|
* Type reference for parsing SSE events containing string data.
|
||||||
|
|
@ -117,6 +117,12 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
||||||
*/
|
*/
|
||||||
protected final Sinks.One<String> messageEndpointSink = Sinks.one();
|
protected final Sinks.One<String> messageEndpointSink = Sinks.one();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* The SSE endpoint URI provided by the server. Used for sending outbound messages via
|
||||||
|
* HTTP POST requests.
|
||||||
|
*/
|
||||||
|
private String sseEndpoint;
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Constructs a new SseClientTransport with the specified WebClient builder. Uses a
|
* Constructs a new SseClientTransport with the specified WebClient builder. Uses a
|
||||||
* default ObjectMapper instance for JSON processing.
|
* default ObjectMapper instance for JSON processing.
|
||||||
|
|
@ -137,11 +143,27 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
||||||
* @throws IllegalArgumentException if either parameter is null
|
* @throws IllegalArgumentException if either parameter is null
|
||||||
*/
|
*/
|
||||||
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
|
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
|
||||||
|
this(webClientBuilder, objectMapper, DEFAULT_SSE_ENDPOINT);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 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
|
||||||
|
* @param sseEndpoint the SSE endpoint URI to use for establishing the connection
|
||||||
|
* @throws IllegalArgumentException if either parameter is null
|
||||||
|
*/
|
||||||
|
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper,
|
||||||
|
String sseEndpoint) {
|
||||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||||
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
||||||
|
Assert.hasText(sseEndpoint, "SSE endpoint must not be null or empty");
|
||||||
|
|
||||||
this.objectMapper = objectMapper;
|
this.objectMapper = objectMapper;
|
||||||
this.webClient = webClientBuilder.build();
|
this.webClient = webClientBuilder.build();
|
||||||
|
this.sseEndpoint = sseEndpoint;
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|
@ -254,7 +276,7 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
||||||
protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
|
protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
|
||||||
return this.webClient
|
return this.webClient
|
||||||
.get()
|
.get()
|
||||||
.uri(SSE_ENDPOINT)
|
.uri(this.sseEndpoint)
|
||||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||||
.retrieve()
|
.retrieve()
|
||||||
.bodyToFlux(SSE_TYPE)
|
.bodyToFlux(SSE_TYPE)
|
||||||
|
|
@ -321,4 +343,66 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
||||||
return this.objectMapper.convertValue(data, typeRef);
|
return this.objectMapper.convertValue(data, typeRef);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Creates a new builder for {@link WebFluxSseClientTransport}.
|
||||||
|
* @param webClientBuilder the WebClient.Builder to use for creating the WebClient
|
||||||
|
* instance
|
||||||
|
* @return a new builder instance
|
||||||
|
*/
|
||||||
|
public static Builder builder(WebClient.Builder webClientBuilder) {
|
||||||
|
return new Builder(webClientBuilder);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Builder for {@link WebFluxSseClientTransport}.
|
||||||
|
*/
|
||||||
|
public static class Builder {
|
||||||
|
|
||||||
|
private final WebClient.Builder webClientBuilder;
|
||||||
|
|
||||||
|
private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
|
||||||
|
|
||||||
|
private ObjectMapper objectMapper = new ObjectMapper();
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Creates a new builder with the specified WebClient.Builder.
|
||||||
|
* @param webClientBuilder the WebClient.Builder to use
|
||||||
|
*/
|
||||||
|
public Builder(WebClient.Builder webClientBuilder) {
|
||||||
|
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
|
||||||
|
this.webClientBuilder = webClientBuilder;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sets the SSE endpoint path.
|
||||||
|
* @param sseEndpoint the SSE endpoint path
|
||||||
|
* @return this builder
|
||||||
|
*/
|
||||||
|
public Builder sseEndpoint(String sseEndpoint) {
|
||||||
|
Assert.hasText(sseEndpoint, "sseEndpoint must not be empty");
|
||||||
|
this.sseEndpoint = sseEndpoint;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sets the object mapper for JSON serialization/deserialization.
|
||||||
|
* @param objectMapper the object mapper
|
||||||
|
* @return this builder
|
||||||
|
*/
|
||||||
|
public Builder objectMapper(ObjectMapper objectMapper) {
|
||||||
|
Assert.notNull(objectMapper, "objectMapper must not be null");
|
||||||
|
this.objectMapper = objectMapper;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Builds a new {@link WebFluxSseClientTransport} instance.
|
||||||
|
* @return a new transport instance
|
||||||
|
*/
|
||||||
|
public WebFluxSseClientTransport build() {
|
||||||
|
return new WebFluxSseClientTransport(webClientBuilder, objectMapper, sseEndpoint);
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -346,4 +346,74 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public static Builder builder() {
|
||||||
|
return new Builder();
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Builder for creating instances of {@link WebFluxSseServerTransportProvider}.
|
||||||
|
* <p>
|
||||||
|
* This builder provides a fluent API for configuring and creating instances of
|
||||||
|
* WebFluxSseServerTransportProvider with custom settings.
|
||||||
|
*/
|
||||||
|
public static class Builder {
|
||||||
|
|
||||||
|
private ObjectMapper objectMapper;
|
||||||
|
|
||||||
|
private String messageEndpoint;
|
||||||
|
|
||||||
|
private String sseEndpoint = DEFAULT_SSE_ENDPOINT;
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sets the ObjectMapper to use for JSON serialization/deserialization of MCP
|
||||||
|
* messages.
|
||||||
|
* @param objectMapper The ObjectMapper instance. Must not be null.
|
||||||
|
* @return this builder instance
|
||||||
|
* @throws IllegalArgumentException if objectMapper is null
|
||||||
|
*/
|
||||||
|
public Builder objectMapper(ObjectMapper objectMapper) {
|
||||||
|
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||||
|
this.objectMapper = objectMapper;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sets the endpoint URI where clients should send their JSON-RPC messages.
|
||||||
|
* @param messageEndpoint The message endpoint URI. Must not be null.
|
||||||
|
* @return this builder instance
|
||||||
|
* @throws IllegalArgumentException if messageEndpoint is null
|
||||||
|
*/
|
||||||
|
public Builder messageEndpoint(String messageEndpoint) {
|
||||||
|
Assert.notNull(messageEndpoint, "Message endpoint must not be null");
|
||||||
|
this.messageEndpoint = messageEndpoint;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Sets the SSE endpoint path.
|
||||||
|
* @param sseEndpoint The SSE endpoint path. Must not be null.
|
||||||
|
* @return this builder instance
|
||||||
|
* @throws IllegalArgumentException if sseEndpoint is null
|
||||||
|
*/
|
||||||
|
public Builder sseEndpoint(String sseEndpoint) {
|
||||||
|
Assert.notNull(sseEndpoint, "SSE endpoint must not be null");
|
||||||
|
this.sseEndpoint = sseEndpoint;
|
||||||
|
return this;
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Builds a new instance of {@link WebFluxSseServerTransportProvider} with the
|
||||||
|
* configured settings.
|
||||||
|
* @return A new WebFluxSseServerTransportProvider instance
|
||||||
|
* @throws IllegalStateException if required parameters are not set
|
||||||
|
*/
|
||||||
|
public WebFluxSseServerTransportProvider build() {
|
||||||
|
Assert.notNull(objectMapper, "ObjectMapper must be set");
|
||||||
|
Assert.notNull(messageEndpoint, "Message endpoint must be set");
|
||||||
|
|
||||||
|
return new WebFluxSseServerTransportProvider(objectMapper, messageEndpoint, sseEndpoint);
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -46,14 +46,17 @@ import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||||
|
|
||||||
import static org.assertj.core.api.Assertions.assertThat;
|
import static org.assertj.core.api.Assertions.assertThat;
|
||||||
import static org.awaitility.Awaitility.await;
|
import static org.awaitility.Awaitility.await;
|
||||||
import static org.junit.Assert.assertThat;
|
|
||||||
import static org.mockito.Mockito.mock;
|
import static org.mockito.Mockito.mock;
|
||||||
|
|
||||||
public class WebFluxSseIntegrationTests {
|
public class WebFluxSseIntegrationTests {
|
||||||
|
|
||||||
private static final int PORT = 8182;
|
private static final int PORT = 8182;
|
||||||
|
|
||||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
// private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||||
|
|
||||||
|
private static final String CUSTOM_SSE_ENDPOINT = "/somePath/sse";
|
||||||
|
|
||||||
|
private static final String CUSTOM_MESSAGE_ENDPOINT = "/otherPath/mcp/message";
|
||||||
|
|
||||||
private DisposableServer httpServer;
|
private DisposableServer httpServer;
|
||||||
|
|
||||||
|
|
@ -64,15 +67,25 @@ public class WebFluxSseIntegrationTests {
|
||||||
@BeforeEach
|
@BeforeEach
|
||||||
public void before() {
|
public void before() {
|
||||||
|
|
||||||
this.mcpServerTransportProvider = new WebFluxSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT);
|
this.mcpServerTransportProvider = new WebFluxSseServerTransportProvider.Builder()
|
||||||
|
.objectMapper(new ObjectMapper())
|
||||||
|
.messageEndpoint(CUSTOM_MESSAGE_ENDPOINT)
|
||||||
|
.sseEndpoint(CUSTOM_SSE_ENDPOINT)
|
||||||
|
.build();
|
||||||
|
|
||||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransportProvider.getRouterFunction());
|
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransportProvider.getRouterFunction());
|
||||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||||
this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||||
|
|
||||||
clientBulders.put("httpclient", McpClient.sync(new HttpClientSseClientTransport("http://localhost:" + PORT)));
|
clientBulders.put("httpclient",
|
||||||
|
McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + PORT)
|
||||||
|
.sseEndpoint(CUSTOM_SSE_ENDPOINT)
|
||||||
|
.build()));
|
||||||
clientBulders.put("webflux",
|
clientBulders.put("webflux",
|
||||||
McpClient.sync(new WebFluxSseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT))));
|
McpClient
|
||||||
|
.sync(WebFluxSseClientTransport.builder(WebClient.builder().baseUrl("http://localhost:" + PORT))
|
||||||
|
.sseEndpoint(CUSTOM_SSE_ENDPOINT)
|
||||||
|
.build()));
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -33,7 +33,7 @@ class WebFluxSseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
protected McpClientTransport createMcpTransport() {
|
protected McpClientTransport createMcpTransport() {
|
||||||
return new WebFluxSseClientTransport(WebClient.builder().baseUrl(host));
|
return WebFluxSseClientTransport.builder(WebClient.builder().baseUrl(host)).build();
|
||||||
}
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
|
|
|
||||||
|
|
@ -33,7 +33,7 @@ class WebFluxSseMcpSyncClientTests extends AbstractMcpSyncClientTests {
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
protected McpClientTransport createMcpTransport() {
|
protected McpClientTransport createMcpTransport() {
|
||||||
return new WebFluxSseClientTransport(WebClient.builder().baseUrl(host));
|
return WebFluxSseClientTransport.builder(WebClient.builder().baseUrl(host)).build();
|
||||||
}
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
|
|
|
||||||
|
|
@ -63,13 +63,6 @@ class WebFluxSseClientTransportTests {
|
||||||
super(webClientBuilder, 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
|
@Override
|
||||||
protected Flux<ServerSentEvent<String>> eventStream() {
|
protected Flux<ServerSentEvent<String>> eventStream() {
|
||||||
return super.eventStream().mergeWith(events.asFlux());
|
return super.eventStream().mergeWith(events.asFlux());
|
||||||
|
|
@ -137,6 +130,33 @@ class WebFluxSseClientTransportTests {
|
||||||
.hasMessageContaining("ObjectMapper must not be null");
|
.hasMessageContaining("ObjectMapper must not be null");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void testBuilderPattern() {
|
||||||
|
// Test default builder
|
||||||
|
WebFluxSseClientTransport transport1 = WebFluxSseClientTransport.builder(webClientBuilder).build();
|
||||||
|
assertThatCode(() -> transport1.closeGracefully().block()).doesNotThrowAnyException();
|
||||||
|
|
||||||
|
// Test builder with custom ObjectMapper
|
||||||
|
ObjectMapper customMapper = new ObjectMapper();
|
||||||
|
WebFluxSseClientTransport transport2 = WebFluxSseClientTransport.builder(webClientBuilder)
|
||||||
|
.objectMapper(customMapper)
|
||||||
|
.build();
|
||||||
|
assertThatCode(() -> transport2.closeGracefully().block()).doesNotThrowAnyException();
|
||||||
|
|
||||||
|
// Test builder with custom SSE endpoint
|
||||||
|
WebFluxSseClientTransport transport3 = WebFluxSseClientTransport.builder(webClientBuilder)
|
||||||
|
.sseEndpoint("/custom-sse")
|
||||||
|
.build();
|
||||||
|
assertThatCode(() -> transport3.closeGracefully().block()).doesNotThrowAnyException();
|
||||||
|
|
||||||
|
// Test builder with all custom parameters
|
||||||
|
WebFluxSseClientTransport transport4 = WebFluxSseClientTransport.builder(webClientBuilder)
|
||||||
|
.objectMapper(customMapper)
|
||||||
|
.sseEndpoint("/custom-sse")
|
||||||
|
.build();
|
||||||
|
assertThatCode(() -> transport4.closeGracefully().block()).doesNotThrowAnyException();
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
void testMessageProcessing() {
|
void testMessageProcessing() {
|
||||||
// Create a test message
|
// Create a test message
|
||||||
|
|
@ -240,7 +260,7 @@ class WebFluxSseClientTransportTests {
|
||||||
// Create a WebClient that simulates connection failures
|
// Create a WebClient that simulates connection failures
|
||||||
WebClient.Builder failingWebClientBuilder = WebClient.builder().baseUrl("http://non-existent-host");
|
WebClient.Builder failingWebClientBuilder = WebClient.builder().baseUrl("http://non-existent-host");
|
||||||
|
|
||||||
WebFluxSseClientTransport failingTransport = new WebFluxSseClientTransport(failingWebClientBuilder);
|
WebFluxSseClientTransport failingTransport = WebFluxSseClientTransport.builder(failingWebClientBuilder).build();
|
||||||
|
|
||||||
// Verify that the transport attempts to reconnect
|
// Verify that the transport attempts to reconnect
|
||||||
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
|
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
|
||||||
|
|
|
||||||
|
|
@ -31,7 +31,9 @@ class WebFluxSseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
protected McpServerTransportProvider createMcpTransportProvider() {
|
protected McpServerTransportProvider createMcpTransportProvider() {
|
||||||
var transportProvider = new WebFluxSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT);
|
var transportProvider = new WebFluxSseServerTransportProvider.Builder().objectMapper(new ObjectMapper())
|
||||||
|
.messageEndpoint(MESSAGE_ENDPOINT)
|
||||||
|
.build();
|
||||||
|
|
||||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transportProvider.getRouterFunction());
|
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transportProvider.getRouterFunction());
|
||||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||||
|
|
|
||||||
|
|
@ -33,7 +33,9 @@ class WebFluxSseMcpSyncServerTests extends AbstractMcpSyncServerTests {
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
protected McpServerTransportProvider createMcpTransportProvider() {
|
protected McpServerTransportProvider createMcpTransportProvider() {
|
||||||
transportProvider = new WebFluxSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT);
|
transportProvider = new WebFluxSseServerTransportProvider.Builder().objectMapper(new ObjectMapper())
|
||||||
|
.messageEndpoint(MESSAGE_ENDPOINT)
|
||||||
|
.build();
|
||||||
return transportProvider;
|
return transportProvider;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue