diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java index a8cad48..9ecd8a7 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java +++ b/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 { diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java index 460bc01..02e713a 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java @@ -8,6 +8,7 @@ import java.time.Duration; 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; @@ -40,6 +41,7 @@ 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.junit.jupiter.api.Assertions.assertInstanceOf; /** * Test suite for the {@link McpAsyncClient} that can be used with different @@ -193,22 +195,6 @@ 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 -> { @@ -525,4 +511,52 @@ public abstract class AbstractMcpAsyncClientTests { }); } + @Test + void testSampling() { + McpClientTransport transport = createMcpTransport(); + + final String message = "Hello, world!"; + final String response = "Goodbye, world!"; + final int maxTokens = 100; + + AtomicReference receivedPrompt = new AtomicReference<>(); + AtomicReference 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(); + }); + } + } diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java index 3785fd6..5c7256e 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java @@ -8,6 +8,7 @@ import java.time.Duration; import java.util.List; import java.util.Map; 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; @@ -31,13 +32,12 @@ import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; -import reactor.core.scheduler.Scheduler; -import reactor.core.scheduler.Schedulers; 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.junit.jupiter.api.Assertions.assertInstanceOf; /** * Unit tests for MCP Client Session functionality. @@ -438,4 +438,48 @@ public abstract class AbstractMcpSyncClientTests { .hasMessageContaining("Logging level must not be null")); } + @Test + void testSampling() { + McpClientTransport transport = createMcpTransport(); + + final String message = "Hello, world!"; + final String response = "Goodbye, world!"; + final int maxTokens = 100; + + AtomicReference receivedPrompt = new AtomicReference<>(); + AtomicReference 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 new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response), + "modelId", McpSchema.CreateMessageResult.StopReason.END_TURN); + }), client -> { + client.initialize(); + + McpSchema.CallToolResult result = client.callTool( + new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens))); + + // 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); + }); + } + } diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java index 1b7c77c..e146656 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java @@ -8,6 +8,7 @@ import java.time.Duration; 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; @@ -40,6 +41,7 @@ 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.junit.jupiter.api.Assertions.assertInstanceOf; /** * Test suite for the {@link McpAsyncClient} that can be used with different @@ -194,22 +196,6 @@ 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 -> { @@ -486,7 +472,6 @@ public abstract class AbstractMcpAsyncClientTests { // --------------------------------------- // Logging Tests // --------------------------------------- - @Test void testLoggingLevelsWithoutInitialization() { verifyNotificationSucceedsWithImplicitInitialization( @@ -526,4 +511,52 @@ public abstract class AbstractMcpAsyncClientTests { }); } + @Test + void testSampling() { + McpClientTransport transport = createMcpTransport(); + + final String message = "Hello, world!"; + final String response = "Goodbye, world!"; + final int maxTokens = 100; + + AtomicReference receivedPrompt = new AtomicReference<>(); + AtomicReference 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(); + }); + } + } diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java index 7798957..c626628 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java @@ -8,6 +8,7 @@ import java.time.Duration; import java.util.List; import java.util.Map; 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; @@ -38,6 +39,7 @@ 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.junit.jupiter.api.Assertions.assertInstanceOf; /** * Unit tests for MCP Client Session functionality. @@ -439,4 +441,48 @@ public abstract class AbstractMcpSyncClientTests { .hasMessageContaining("Logging level must not be null")); } + @Test + void testSampling() { + McpClientTransport transport = createMcpTransport(); + + final String message = "Hello, world!"; + final String response = "Goodbye, world!"; + final int maxTokens = 100; + + AtomicReference receivedPrompt = new AtomicReference<>(); + AtomicReference 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 new McpSchema.CreateMessageResult(McpSchema.Role.USER, new McpSchema.TextContent(response), + "modelId", McpSchema.CreateMessageResult.StopReason.END_TURN); + }), client -> { + client.initialize(); + + McpSchema.CallToolResult result = client.callTool( + new McpSchema.CallToolRequest("sampleLLM", Map.of("prompt", message, "maxTokens", maxTokens))); + + // 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); + }); + } + }