modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Validate message endpoint in SSE WebFlux client transport

Browse source 
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
Daniel Garnier-Moiroux 2026-04-30 18:25:47 +02:00 • committed by Daniel Garnier-Moiroux
commit d8e899ece7
2 changed files with 112 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

69
mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java
View file

@ -5,7 +5,9 @@
package io.modelcontextprotocol.client.transport;
import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiConsumer;
import java.util.function.Function;
@ -128,7 +130,17 @@ public class WebFluxSseClientTransport implements McpClientTransport {
* The SSE endpoint URI provided by the server. Used for sending outbound messages via
* HTTP POST requests.
*/
private String sseEndpoint;
private final String sseEndpoint;
/**
* Used to capture the full SSE URI from the web client when connecting.
*/
private final AtomicReference<URI> sseUri = new AtomicReference<>();
/**
* Validator for the message endpoint.
*/
private final SseMessageEndpointValidator messageEndpointValidator;
/**
* Constructs a new SseClientTransport with the specified WebClient builder and
@ -152,13 +164,30 @@ public class WebFluxSseClientTransport implements McpClientTransport {
* @throws IllegalArgumentException if either parameter is null
*/
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper, String sseEndpoint) {
this(webClientBuilder, jsonMapper, sseEndpoint, new DefaultSseMessageEndpointValidator());
}
/**
* 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 jsonMapper the ObjectMapper to use for JSON processing
* @param sseEndpoint the SSE endpoint URI to use for establishing the connection
* @param messageEndpointValidator validator for the message endpoint
* @throws IllegalArgumentException if either parameter is null
*/
public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper, String sseEndpoint,
SseMessageEndpointValidator messageEndpointValidator) {
Assert.notNull(jsonMapper, "jsonMapper must not be null");
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null");
Assert.hasText(sseEndpoint, "SSE endpoint must not be null or empty");
this.jsonMapper = jsonMapper;
this.webClient = webClientBuilder.build();
this.sseEndpoint = sseEndpoint;
this.messageEndpointValidator = messageEndpointValidator;
}
@Override
@ -195,6 +224,14 @@ public class WebFluxSseClientTransport implements McpClientTransport {
this.inboundSubscription = events.concatMap(event -> Mono.just(event).<JSONRPCMessage>handle((e, s) -> {
if (ENDPOINT_EVENT_TYPE.equals(event.event())) {
String messageEndpointUri = event.data();
try {
this.messageEndpointValidator.validate(this.sseUri.get(), messageEndpointUri);
}
catch (InvalidSseMessageEndpointException ex) {
messageEndpointSink.tryEmitError(ex);
s.error(ex);
return;
}
if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
s.complete();
}
@ -276,16 +313,17 @@ public class WebFluxSseClientTransport implements McpClientTransport {
* Includes automatic retry logic for handling transient connection failures.
*/
// visible for tests
protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
return this.webClient
.get()
protected Flux<ServerSentEvent<String>> eventStream() {
return this.webClient.get()
.uri(this.sseEndpoint)
.accept(MediaType.TEXT_EVENT_STREAM)
.header(HttpHeaders.PROTOCOL_VERSION, MCP_PROTOCOL_VERSION)
.retrieve()
.bodyToFlux(SSE_TYPE)
.exchangeToFlux(exchange -> {
this.sseUri.set(exchange.request().getURI());
return exchange.bodyToFlux(SSE_TYPE);
})
.retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler)));
} // @formatter:on
}
/**
* Retry handler for the inbound SSE stream. Implements the retry logic for handling
@ -368,6 +406,8 @@ public class WebFluxSseClientTransport implements McpClientTransport {
private McpJsonMapper jsonMapper;
private SseMessageEndpointValidator messageEndpointValidator = new DefaultSseMessageEndpointValidator();
/**
* Creates a new builder with the specified WebClient.Builder.
* @param webClientBuilder the WebClient.Builder to use
@ -399,13 +439,26 @@ public class WebFluxSseClientTransport implements McpClientTransport {
return this;
}
/**
* Sets the validator that ensure the message endpoint returned over the SSE
* connection is valid.
* @param messageEndpointValidator the validator
* @return this builder
*/
public Builder messageEndpointValidator(SseMessageEndpointValidator messageEndpointValidator) {
Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null");
this.messageEndpointValidator = messageEndpointValidator;
return this;
}
/**
* Builds a new {@link WebFluxSseClientTransport} instance.
* @return a new transport instance
*/
public WebFluxSseClientTransport build() {
return new WebFluxSseClientTransport(webClientBuilder,
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, sseEndpoint);
jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, sseEndpoint,
messageEndpointValidator);
}
}

54
mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransportTests.java
View file

@ -4,6 +4,7 @@
package io.modelcontextprotocol.client.transport;
import java.net.URI;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
@ -21,6 +22,7 @@ import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.mockito.ArgumentCaptor;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import reactor.core.publisher.Flux;
@ -35,6 +37,9 @@ import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.matches;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
/**
* Tests for the {@link WebFluxSseClientTransport} class.
@ -57,6 +62,8 @@ class WebFluxSseClientTransportTests {
private WebClient.Builder webClientBuilder;
private SseMessageEndpointValidator sseMessageEndpointValidator = mock(SseMessageEndpointValidator.class);
// Test class to access protected methods
static class TestSseClientTransport extends WebFluxSseClientTransport {
@ -64,8 +71,9 @@ class WebFluxSseClientTransportTests {
private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
public TestSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper) {
super(webClientBuilder, jsonMapper);
public TestSseClientTransport(WebClient.Builder webClientBuilder, McpJsonMapper jsonMapper,
SseMessageEndpointValidator sseMessageEndpointValidator) {
super(webClientBuilder, jsonMapper, "/sse", sseMessageEndpointValidator);
}
@Override
@ -113,7 +121,7 @@ class WebFluxSseClientTransportTests {
@BeforeEach
void setUp() {
webClientBuilder = WebClient.builder().baseUrl(host);
transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER);
transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER, sseMessageEndpointValidator);
transport.connect(Function.identity()).block();
}
@ -368,4 +376,44 @@ class WebFluxSseClientTransportTests {
assertThat(transport.getInboundMessageCount()).isEqualTo(3);
}
@Test
void testMessageEndpointValidation() throws InvalidSseMessageEndpointException {
var uriCaptor = ArgumentCaptor.forClass(URI.class);
verify(sseMessageEndpointValidator).validate(uriCaptor.capture(), matches("/message\\?sessionId=[a-z0-9-]+"));
assertThat(uriCaptor.getValue().toString()).matches("http://localhost:\\d+/sse");
}
@Test
void testMessageEndpointValidationRejects() {
TestSseClientTransport transport = new TestSseClientTransport(webClientBuilder, JSON_MAPPER,
(sseUri, messageEndpoint) -> {
throw new InvalidSseMessageEndpointException("boom", messageEndpoint);
});
try {
// fails to connect
StepVerifier.create(transport.connect(Function.identity()))
.verifyErrorMatches(WebFluxSseClientTransportTests::isInvalidEndpointError);
// Since connection failed, there is no message endpoint, and no message can
// be sent
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
StepVerifier.create(transport.sendMessage(testMessage))
.verifyErrorMatches(WebFluxSseClientTransportTests::isInvalidEndpointError);
}
finally {
transport.closeGracefully();
}
}
private static boolean isInvalidEndpointError(Throwable e) {
if (e instanceof InvalidSseMessageEndpointException ismee) {
return ismee.getMessageEndpoint().matches("/message\\?sessionId=[a-z0-9-]+")
&& ismee.getMessage().equals("boom");
}
return false;
}
}