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

Fix Streamable HTTP WebClient GET SSE handling

Browse source 
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Dariusz Jędrzejczyk 2025-06-16 11:21:26 +02:00 • committed by Christian Tzolov
commit 8e9857f181
4 changed files with 75 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

18
mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java
View file

@ -125,13 +125,14 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
}
private DefaultMcpTransportSession createTransportSession() {
Supplier<Publisher<Void>> onClose = () -> {
DefaultMcpTransportSession transportSession = this.activeSession.get();
return transportSession.sessionId().isEmpty() ? Mono.empty()
: webClient.delete().uri(this.endpoint).headers(httpHeaders -> {
httpHeaders.add("mcp-session-id", transportSession.sessionId().get());
}).retrieve().toBodilessEntity().doOnError(e -> logger.info("Got response {}", e)).then();
};
Function<String, Publisher<Void>> onClose = sessionId -> sessionId == null ? Mono.empty()
: webClient.delete().uri(this.endpoint).headers(httpHeaders -> {
httpHeaders.add("mcp-session-id", sessionId);
})
.retrieve()
.toBodilessEntity()
.doOnError(e -> logger.warn("Got error when closing transport", e))
.then();
return new DefaultMcpTransportSession(onClose);
}
@ -192,6 +193,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
})
.exchangeToFlux(response -> {
if (isEventStream(response)) {
logger.debug("Established SSE stream via GET");
return eventStream(stream, response);
}
else if (isNotAllowed(response)) {
@ -208,6 +210,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
}).flux();
}
})
.flatMap(jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
.onErrorComplete(t -> {
this.handleException(t);
return true;
@ -274,6 +277,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
else {
MediaType mediaType = contentType.get();
if (mediaType.isCompatibleWith(MediaType.TEXT_EVENT_STREAM)) {
logger.debug("Established SSE stream via POST");
// communicate to caller that the message was delivered
sink.success();
// starting a stream

42
mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java
View file

@ -19,6 +19,8 @@ import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
import io.modelcontextprotocol.spec.McpSchema.Prompt;
import io.modelcontextprotocol.spec.McpSchema.Resource;
@ -77,7 +79,9 @@ public abstract class AbstractMcpAsyncClientTests {
McpClient.AsyncSpec builder = McpClient.async(transport)
.requestTimeout(getRequestTimeout())
.initializationTimeout(getInitializationTimeout())
.capabilities(ClientCapabilities.builder().roots(true).build());
.sampling(req -> Mono.just(new CreateMessageResult(McpSchema.Role.USER,
new McpSchema.TextContent("Oh, hi!"), "modelId", CreateMessageResult.StopReason.END_TURN)))
.capabilities(ClientCapabilities.builder().roots(true).sampling().build());
builder = customizer.apply(builder);
client.set(builder.build());
}).doesNotThrowAnyException();
@ -189,6 +193,22 @@ public abstract class AbstractMcpAsyncClientTests {
});
}
@Test
void testSampling() {
withClient(createMcpTransport(), mcpAsyncClient -> {
CallToolRequest callToolRequest = new CallToolRequest("sampleLLM",
Map.of("prompt", "Hello MCP Spring AI!"));
StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(callToolRequest)))
.consumeNextWith(callToolResult -> {
assertThat(callToolResult).isNotNull().satisfies(result -> {
assertThat(result.content()).isNotNull();
assertThat(result.isError()).isNull();
});
})
.verifyComplete();
});
}
@Test
void testCallToolWithInvalidTool() {
withClient(createMcpTransport(), mcpAsyncClient -> {
@ -424,6 +444,20 @@ public abstract class AbstractMcpAsyncClientTests {
});
}
@Test
void testInitializeWithElicitationCapability() {
ClientCapabilities capabilities = ClientCapabilities.builder().elicitation().build();
ElicitResult elicitResult = ElicitResult.builder()
.message(ElicitResult.Action.ACCEPT)
.content(Map.of("foo", "bar"))
.build();
withClient(createMcpTransport(),
builder -> builder.capabilities(capabilities).elicitation(request -> Mono.just(elicitResult)),
client -> {
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
});
}
@Test
void testInitializeWithAllCapabilities() {
var capabilities = ClientCapabilities.builder()
@ -435,7 +469,11 @@ public abstract class AbstractMcpAsyncClientTests {
Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler = request -> Mono
.just(CreateMessageResult.builder().message("test").model("test-model").build());
withClient(createMcpTransport(), builder -> builder.capabilities(capabilities).sampling(samplingHandler),
Function<ElicitRequest, Mono<ElicitResult>> elicitationHandler = request -> Mono
.just(ElicitResult.builder().message(ElicitResult.Action.ACCEPT).content(Map.of("foo", "bar")).build());
withClient(createMcpTransport(),
builder -> builder.capabilities(capabilities).sampling(samplingHandler).elicitation(elicitationHandler),
client ->
StepVerifier.create(client.initialize()).assertNext(result -> {

9
mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java
View file

@ -10,7 +10,7 @@ import reactor.core.publisher.Mono;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;
import java.util.function.Function;
/**
* Default implementation of {@link McpTransportSession} which manages the open
@ -29,9 +29,9 @@ public class DefaultMcpTransportSession implements McpTransportSession<Disposabl
private final AtomicReference<String> sessionId = new AtomicReference<>();
private final Supplier<Publisher<Void>> onClose;
private final Function<String, Publisher<Void>> onClose;
public DefaultMcpTransportSession(Supplier<Publisher<Void>> onClose) {
public DefaultMcpTransportSession(Function<String, Publisher<Void>> onClose) {
this.onClose = onClose;
}
@ -73,7 +73,8 @@ public class DefaultMcpTransportSession implements McpTransportSession<Disposabl
@Override
public Mono<Void> closeGracefully() {
return Mono.from(this.onClose.get()).then(Mono.fromRunnable(this.openConnections::dispose));
return Mono.from(this.onClose.apply(this.sessionId.get()))
.then(Mono.fromRunnable(this.openConnections::dispose));
}
}

20
mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java
View file

@ -80,7 +80,9 @@ public abstract class AbstractMcpAsyncClientTests {
McpClient.AsyncSpec builder = McpClient.async(transport)
.requestTimeout(getRequestTimeout())
.initializationTimeout(getInitializationTimeout())
.capabilities(ClientCapabilities.builder().roots(true).build());
.sampling(req -> Mono.just(new CreateMessageResult(McpSchema.Role.USER,
new McpSchema.TextContent("Oh, hi!"), "modelId", CreateMessageResult.StopReason.END_TURN)))
.capabilities(ClientCapabilities.builder().roots(true).sampling().build());
builder = customizer.apply(builder);
client.set(builder.build());
}).doesNotThrowAnyException();
@ -192,6 +194,22 @@ public abstract class AbstractMcpAsyncClientTests {
});
}
@Test
void testSampling() {
withClient(createMcpTransport(), mcpAsyncClient -> {
CallToolRequest callToolRequest = new CallToolRequest("sampleLLM",
Map.of("prompt", "Hello MCP Spring AI!"));
StepVerifier.create(mcpAsyncClient.initialize().then(mcpAsyncClient.callTool(callToolRequest)))
.consumeNextWith(callToolResult -> {
assertThat(callToolResult).isNotNull().satisfies(result -> {
assertThat(result.content()).isNotNull();
assertThat(result.isError()).isNull();
});
})
.verifyComplete();
});
}
@Test
void testCallToolWithInvalidTool() {
withClient(createMcpTransport(), mcpAsyncClient -> {