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.

Author: tzolov

mcp/src/main/java/org/springframework/ai/mcp/client/McpAsyncClient.java

@@ -160,7 +160,10 @@ public McpAsyncClient(McpTransport transport, Duration requestTimeout, Implement
 		}
 
 		// 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());
 		}

mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java

@@ -761,6 +761,46 @@ public enum StopReason {
 			@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
 
 	// ---------------------------

mcp/src/test/java/org/springframework/ai/mcp/client/AbstractMcpAsyncClientTests.java

@@ -19,6 +19,7 @@
 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 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 @@ void testInitializeWithSamplingCapability() {
 
 		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 @@ void testInitializeWithAllCapabilities() {
 			.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));

mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java

@@ -27,11 +27,13 @@
 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 @@ void testSamplingCreateMessageRequestHandlingWithNullHandler() {
 		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");
 	}
 
 }