modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: enforce sampling handler presence and add CreateMessageResult builder

Browse source 
Improves error handling by enforcing sampling handler presence when sampling capabilities are enabled. Adds a builder pattern for CreateMessageResult to improve API ergonomics.
This commit is contained in:
Christian Tzolov 2025-01-02 18:19:33 +01:00
commit 51b75dca19
4 changed files with 64 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

5
mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java
View file

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

40
mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java
View file

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

16
mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java
View file

@ -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<CreateMessageRequest, CreateMessageResult> 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));

31
mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java
View file

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