Merge branch 'modelcontextprotocol:main' into fix-notification-handler

Author: pantanurag555

README.md

@@ -7,7 +7,13 @@ This SDK enables Java applications to interact with AI models and tools through
 ## 📚 Reference Documentation 
 
 #### MCP Java SDK documentation
-For comprehensive guides and SDK API documentation, visit the [MCP Java SDK Reference Documentation](https://modelcontextprotocol.io/sdk/java/mcp-overview).
+For comprehensive guides and SDK API documentation
+
+- [Features](https://modelcontextprotocol.io/sdk/java/mcp-overview#features) - Overview the features provided by the Java MCP SDK
+- [Acrchitecture](https://modelcontextprotocol.io/sdk/java/mcp-overview#architecture) - Java MCP SDK architecture overview.
+- [Java Dependencies / BOM](https://modelcontextprotocol.io/sdk/java/mcp-overview#dependencies) - Java dependencies and BOM.
+- [Java MCP Client](https://modelcontextprotocol.io/sdk/java/mcp-client) - Learn how to use the MCP client to interact with MCP servers.
+- [Java MCP Server](https://modelcontextprotocol.io/sdk/java/mcp-server) - Learn how to implement and configure a MCP servers.
 
 #### Spring AI MCP documentation
 [Spring AI MCP](https://docs.spring.io/spring-ai/reference/api/mcp/mcp-overview.html) extends the MCP Java SDK with Spring Boot integration, providing both [client](https://docs.spring.io/spring-ai/reference/api/mcp/mcp-client-boot-starter-docs.html) and [server](https://docs.spring.io/spring-ai/reference/api/mcp/mcp-server-boot-starter-docs.html) starters. Bootstrap your AI applications with MCP support using [Spring Initializer](https://start.spring.io).

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

@@ -125,13 +125,14 @@ public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchem
 	}
 
 	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 @@ private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
 				})
 				.exchangeToFlux(response -> {
 					if (isEventStream(response)) {
+						logger.debug("Established SSE stream via GET");
 						return eventStream(stream, response);
 					}
 					else if (isNotAllowed(response)) {
@@ -208,6 +210,7 @@ else if (isNotFound(response)) {
 						}).flux();
 					}
 				})
+				.flatMap(jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
 				.onErrorComplete(t -> {
 					this.handleException(t);
 					return true;
@@ -274,6 +277,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 						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-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java

@@ -1,13 +1,13 @@
 package io.modelcontextprotocol.client;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
 import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import org.junit.jupiter.api.Timeout;
-import org.springframework.web.reactive.function.client.WebClient;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
 
+import org.springframework.web.reactive.function.client.WebClient;
+
 @Timeout(15)
 public class WebClientStreamableHttpSyncClientTests extends AbstractMcpSyncClientTests {
 

mcp-test/pom.xml

@@ -54,6 +54,11 @@
 			<artifactId>junit-jupiter-api</artifactId>
 			<version>${junit.version}</version>
 		</dependency>
+		<dependency>
+			<groupId>org.junit.jupiter</groupId>
+			<artifactId>junit-jupiter-params</artifactId>
+			<version>${junit.version}</version>
+		</dependency>
 		<dependency>
 			<groupId>org.mockito</groupId>
 			<artifactId>mockito-core</artifactId>

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

@@ -8,6 +8,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -19,6 +20,8 @@
 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;
@@ -31,13 +34,17 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 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.assertj.core.api.Assertions.fail;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 
 /**
  * Test suite for the {@link McpAsyncClient} that can be used with different
@@ -77,7 +84,9 @@ McpAsyncClient client(McpClientTransport transport, Function<McpClient.AsyncSpec
 			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();
@@ -202,6 +211,64 @@ void testCallToolWithInvalidTool() {
 		});
 	}
 
+	@ParameterizedTest
+	@ValueSource(strings = { "success", "error", "debug" })
+	void testCallToolWithMessageAnnotations(String messageType) {
+		McpClientTransport transport = createMcpTransport();
+
+		withClient(transport, mcpAsyncClient -> {
+			StepVerifier.create(mcpAsyncClient.initialize()
+				.then(mcpAsyncClient.callTool(new McpSchema.CallToolRequest("annotatedMessage",
+						Map.of("messageType", messageType, "includeImage", true)))))
+				.consumeNextWith(result -> {
+					assertThat(result).isNotNull();
+					assertThat(result.isError()).isNotEqualTo(true);
+					assertThat(result.content()).isNotEmpty();
+					assertThat(result.content()).allSatisfy(content -> {
+						switch (content.type()) {
+							case "text":
+								McpSchema.TextContent textContent = assertInstanceOf(McpSchema.TextContent.class,
+										content);
+								assertThat(textContent.text()).isNotEmpty();
+								assertThat(textContent.annotations()).isNotNull();
+
+								switch (messageType) {
+									case "error":
+										assertThat(textContent.annotations().priority()).isEqualTo(1.0);
+										assertThat(textContent.annotations().audience())
+											.containsOnly(McpSchema.Role.USER, McpSchema.Role.ASSISTANT);
+										break;
+									case "success":
+										assertThat(textContent.annotations().priority()).isEqualTo(0.7);
+										assertThat(textContent.annotations().audience())
+											.containsExactly(McpSchema.Role.USER);
+										break;
+									case "debug":
+										assertThat(textContent.annotations().priority()).isEqualTo(0.3);
+										assertThat(textContent.annotations().audience())
+											.containsExactly(McpSchema.Role.ASSISTANT);
+										break;
+									default:
+										throw new IllegalStateException("Unexpected value: " + content.type());
+								}
+								break;
+							case "image":
+								McpSchema.ImageContent imageContent = assertInstanceOf(McpSchema.ImageContent.class,
+										content);
+								assertThat(imageContent.data()).isNotEmpty();
+								assertThat(imageContent.annotations()).isNotNull();
+								assertThat(imageContent.annotations().priority()).isEqualTo(0.5);
+								assertThat(imageContent.annotations().audience()).containsExactly(McpSchema.Role.USER);
+								break;
+							default:
+								fail("Unexpected content type: " + content.type());
+						}
+					});
+				})
+				.verifyComplete();
+		});
+	}
+
 	@Test
 	void testListResourcesWithoutInitialization() {
 		verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(null), "listing resources");
@@ -424,6 +491,20 @@ void testInitializeWithSamplingCapability() {
 				});
 	}
 
+	@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 +516,11 @@ void testInitializeWithAllCapabilities() {
 		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 -> {
@@ -487,4 +572,52 @@ void testLoggingWithNullNotification() {
 		});
 	}
 
+	@Test
+	void testSampling() {
+		McpClientTransport transport = createMcpTransport();
+
+		final String message = "Hello, world!";
+		final String response = "Goodbye, world!";
+		final int maxTokens = 100;
+
+		AtomicReference<String> receivedPrompt = new AtomicReference<>();
+		AtomicReference<String> receivedMessage = new AtomicReference<>();
+		AtomicInteger receivedMaxTokens = new AtomicInteger();
+
+		withClient(transport, spec -> spec.capabilities(McpSchema.ClientCapabilities.builder().sampling().build())
+			.sampling(request -> {
+				McpSchema.TextContent messageText = assertInstanceOf(McpSchema.TextContent.class,
+						request.messages().get(0).content());
+				receivedPrompt.set(request.systemPrompt());
+				receivedMessage.set(messageText.text());
+				receivedMaxTokens.set(request.maxTokens());
+
+				return Mono
+					.just(new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response),
+							"modelId", McpSchema.CreateMessageResult.StopReason.END_TURN));
+			}), client -> {
+				StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
+
+				StepVerifier.create(client.callTool(
+						new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens))))
+					.consumeNextWith(result -> {
+						// Verify tool response to ensure our sampling response was passed
+						// through
+						assertThat(result.content()).hasAtLeastOneElementOfType(McpSchema.TextContent.class);
+						assertThat(result.content()).allSatisfy(content -> {
+							if (!(content instanceof McpSchema.TextContent text))
+								return;
+
+							assertThat(text.text()).endsWith(response); // Prefixed
+						});
+
+						// Verify sampling request parameters received in our callback
+						assertThat(receivedPrompt.get()).isNotEmpty();
+						assertThat(receivedMessage.get()).endsWith(message); // Prefixed
+						assertThat(receivedMaxTokens.get()).isEqualTo(maxTokens);
+					})
+					.verifyComplete();
+			});
+	}
+
 }