Fix Streamable HTTP WebClient GET SSE handling
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
parent
9a63538840
commit
8e9857f181
2 changed files with 75 additions and 14 deletions
|
|
@ -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));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 -> {
|
||||
|
|
|
|||
Loading…
Reference in a new issue