From 51b75dca196ae1a3615242f8005e677d259869a5 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Thu, 2 Jan 2025 18:19:33 +0100 Subject: [PATCH] feat: enforce sampling handler presence and add CreateMessageResult builder Improves error handling by enforcing sampling handler presence when sampling capabilities are enabled. Adds a builder pattern for CreateMessageResult to improve API ergonomics. --- .../ai/mcp/client/McpAsyncClient.java | 5 ++- .../ai/mcp/spec/McpSchema.java | 40 +++++++++++++++++++ .../client/AbstractMcpAsyncClientTests.java | 16 +++++++- .../McpAsyncClientResponseHandlerTests.java | 31 +++----------- 4 files changed, 64 insertions(+), 28 deletions(-) diff --git a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java index bf1f17c..c4473a4 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java @@ -160,7 +160,10 @@ public class McpAsyncClient { } // Sampling Handler - if (samplingHandler != null && this.clientCapabilities.sampling() != null) { + if (this.clientCapabilities.sampling() != null) { + if (samplingHandler == null) { + throw new McpError("Sampling handler must not be null when client capabilities include sampling"); + } this.samplingHandler = samplingHandler; requestHanlers.put(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, samplingCreateMessageHandler()); } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java index 663cc0c..bb3b301 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java @@ -761,6 +761,46 @@ public class McpSchema { @JsonProperty("stop_sequence") STOP_SEQUENCE, @JsonProperty("max_tokens") MAX_TOKENS } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + private Role role = Role.ASSISTANT; + private Content content; + private String model; + private StopReason stopReason = StopReason.END_TURN; + + public Builder role(Role role) { + this.role = role; + return this; + } + + public Builder content(Content content) { + this.content = content; + return this; + } + + public Builder model(String model) { + this.model = model; + return this; + } + + public Builder stopReason(StopReason stopReason) { + this.stopReason = stopReason; + return this; + } + + public Builder message(String message) { + this.content = new TextContent(message); + return this; + } + + public CreateMessageResult build() { + return new CreateMessageResult(role, content, model, stopReason); + } + } }// @formatter:on // --------------------------- diff --git a/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java b/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java index c8d357d..acde147 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java @@ -19,6 +19,7 @@ package org.springframework.ai.mcp.client; import java.time.Duration; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Function; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; @@ -29,6 +30,8 @@ import reactor.test.StepVerifier; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest; import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities; +import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest; +import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult; import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest; import org.springframework.ai.mcp.spec.McpSchema.Prompt; import org.springframework.ai.mcp.spec.McpSchema.Resource; @@ -295,7 +298,9 @@ public abstract class AbstractMcpAsyncClientTests { var capabilities = ClientCapabilities.builder().sampling().build(); - var client = McpClient.using(transport).requestTimeout(TIMEOUT).capabilities(capabilities).async(); + var client = McpClient.using(transport).requestTimeout(TIMEOUT).capabilities(capabilities).sampling(request -> { + return CreateMessageResult.builder().message("test").model("test-model").build(); + }).async(); assertThatCode(() -> { client.initialize().block(Duration.ofSeconds(10)); @@ -313,7 +318,14 @@ public abstract class AbstractMcpAsyncClientTests { .sampling() .build(); - var client = McpClient.using(transport).requestTimeout(TIMEOUT).capabilities(capabilities).async(); + Function samplingHandler = request -> { + return CreateMessageResult.builder().message("test").model("test-model").build(); + }; + var client = McpClient.using(transport) + .requestTimeout(TIMEOUT) + .capabilities(capabilities) + .sampling(samplingHandler) + .async(); assertThatCode(() -> { var result = client.initialize().block(Duration.ofSeconds(10)); diff --git a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java index 0b2c374..da18e5f 100644 --- a/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java +++ b/mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java @@ -27,11 +27,13 @@ import com.fasterxml.jackson.core.type.TypeReference; import org.junit.jupiter.api.Test; import org.springframework.ai.mcp.MockMcpTransport; +import org.springframework.ai.mcp.spec.McpError; import org.springframework.ai.mcp.spec.McpSchema; import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities; import org.springframework.ai.mcp.spec.McpSchema.Root; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.awaitility.Awaitility.await; class McpAsyncClientResponseHandlerTests { @@ -285,31 +287,10 @@ class McpAsyncClientResponseHandlerTests { MockMcpTransport transport = new MockMcpTransport(); // Create client with sampling capability but null handler - McpAsyncClient asyncMcpClient = McpClient.using(transport) - .capabilities(ClientCapabilities.builder().sampling().build()) - .async(); - - // Create a mock create message request - var messageRequest = new McpSchema.CreateMessageRequest( - List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))), - null, null, null, null, 0, null, null); - - // Simulate incoming request - McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, - McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, "test-id", messageRequest); - transport.simulateIncomingMessage(request); - - // Verify error response - McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); - assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class); - - McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; - assertThat(response.id()).isEqualTo("test-id"); - assertThat(response.result()).isNull(); - assertThat(response.error()).isNotNull(); - assertThat(response.error().message()).contains("Method not found: sampling/createMessage"); - - asyncMcpClient.closeGracefully(); + assertThatThrownBy( + () -> McpClient.using(transport).capabilities(ClientCapabilities.builder().sampling().build()).async()) + .isInstanceOf(McpError.class) + .hasMessage("Sampling handler must not be null when client capabilities include sampling"); } }