Fix Streamable HTTP WebClient GET SSE handling

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
2 changed files with 75 additions and 14 deletions

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));
}
}

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 -> {