Fix Streamable HTTP WebClient GET SSE handling

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

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

@@ -125,13 +125,14 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
 	}
 
 	private DefaultMcpTransportSession createTransportSession() {
-		Supplier<Publisher<Void>> onClose = () -> {
+		Function<String, Publisher<Void>> onClose = sessionId -> sessionId == null ? Mono.empty()
-			DefaultMcpTransportSession transportSession = this.activeSession.get();
+				: webClient.delete().uri(this.endpoint).headers(httpHeaders -> {
-			return transportSession.sessionId().isEmpty() ? Mono.empty()
+					httpHeaders.add("mcp-session-id", sessionId);
-					: webClient.delete().uri(this.endpoint).headers(httpHeaders -> {
+				})
-						httpHeaders.add("mcp-session-id", transportSession.sessionId().get());
+					.retrieve()
-					}).retrieve().toBodilessEntity().doOnError(e -> logger.info("Got response {}", e)).then();
+					.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

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

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

mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java

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

mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java

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