Add SEP-1577 sampling with tools support

MIGRATION-2.0.md

@@ -173,3 +173,13 @@ The 1.x canonical 4-arg constructors continue to compile.
 When a `JsonSchemaValidator` is available (including the default from `McpJsonDefaults.getSchemaValidator()` when you do not configure one explicitly) and `validateToolInputs` is left at its default of `true`, the server validates incoming tool arguments against `tool.inputSchema()` before invoking the tool. Failed validation produces a `CallToolResult` with `isError` set and a textual error in the content.
 
 **Action:** Ensure `inputSchema` maps are valid for your validator, tighten client arguments, or disable validation with `validateToolInputs(false)` on the server builder if you must preserve pre-2.0 behaviour.
+
+---
+
+## Sampling with tools (SEP-1577)
+
+### New `StopReason.TOOL_USE` enum constant
+
+`McpSchema.CreateMessageResult.StopReason` gains a new value, `TOOL_USE("toolUse")`. Exhaustive `switch` *expressions* over the enum (which the compiler enforces at compile time) that covered all pre-existing constants without a `default` branch will no longer compile.
+
+**Action:** Add a `case TOOL_USE` branch or a `default` branch to any exhaustive `switch` expression on `StopReason`.

mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

@@ -103,6 +103,9 @@ public class McpAsyncClient {
 	public static final TypeRef<CreateMessageRequest> CREATE_MESSAGE_REQUEST_TYPE_REF = new TypeRef<>() {
 	};
 
+	public static final TypeRef<McpSchema.CreateMessageWithToolsRequest> CREATE_MESSAGE_WITH_TOOLS_REQUEST_TYPE_REF = new TypeRef<>() {
+	};
+
 	public static final TypeRef<LoggingMessageNotification> LOGGING_MESSAGE_NOTIFICATION_TYPE_REF = new TypeRef<>() {
 	};
 
@@ -142,6 +145,8 @@ public class McpAsyncClient {
 	 */
 	private Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler;
 
+	private Function<McpSchema.CreateMessageWithToolsRequest, Mono<McpSchema.CreateMessageWithToolsResult>> samplingWithToolsHandler;
+
 	/**
 	 * MCP provides a standardized way for servers to request additional information from
 	 * users through the client during interactions. This flow allows clients to maintain
@@ -227,11 +232,21 @@ public class McpAsyncClient {
 
 		// Sampling Handler
 		if (this.clientCapabilities.sampling() != null) {
-			if (features.samplingHandler() == null) {
-				throw new IllegalArgumentException(
-						"Sampling handler must not be null when client capabilities include sampling");
+			boolean withTools = this.clientCapabilities.sampling().tools() != null;
+			if (withTools) {
+				if (features.samplingWithToolsHandler() == null) {
+					throw new IllegalArgumentException(
+							"Sampling-with-tools handler must not be null when client capabilities include sampling.tools");
+				}
+				this.samplingWithToolsHandler = features.samplingWithToolsHandler();
+			}
+			else {
+				if (features.samplingHandler() == null) {
+					throw new IllegalArgumentException(
+							"Sampling handler must not be null when client capabilities include sampling");
+				}
+				this.samplingHandler = features.samplingHandler();
 			}
-			this.samplingHandler = features.samplingHandler();
 			requestHandlers.put(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, samplingCreateMessageHandler());
 		}
 
@@ -575,11 +590,29 @@ private RequestHandler<McpSchema.ListRootsResult> rootsListRequestHandler() {
 	// --------------------------
 	// Sampling
 	// --------------------------
-	private RequestHandler<CreateMessageResult> samplingCreateMessageHandler() {
+	private RequestHandler<Object> samplingCreateMessageHandler() {
 		return params -> {
-			McpSchema.CreateMessageRequest request = transport.unmarshalFrom(params, CREATE_MESSAGE_REQUEST_TYPE_REF);
-
-			return this.samplingHandler.apply(request);
+			// Dispatch to V2 handler if the request carries tools/toolChoice, otherwise
+			// V1
+			McpSchema.CreateMessageWithToolsRequest v2Request = transport.unmarshalFrom(params,
+					CREATE_MESSAGE_WITH_TOOLS_REQUEST_TYPE_REF);
+			boolean isV2Request = (v2Request.tools() != null && !v2Request.tools().isEmpty())
+					|| v2Request.toolChoice() != null;
+			if (isV2Request) {
+				if (this.samplingWithToolsHandler == null) {
+					return Mono.error(new IllegalStateException(
+							"Received sampling request with tools, but no samplingWithTools handler is registered. "
+									+ "Use McpClient.async(transport).samplingWithTools(handler) to register one."));
+				}
+				return this.samplingWithToolsHandler.apply(v2Request).cast(Object.class);
+			}
+			// V1 path: unmarshal as the legacy type for handler type-safety
+			McpSchema.CreateMessageRequest v1Request = transport.unmarshalFrom(params, CREATE_MESSAGE_REQUEST_TYPE_REF);
+			if (this.samplingHandler == null) {
+				return Mono.error(
+						new IllegalStateException("Received sampling request, but no sampling handler is registered."));
+			}
+			return this.samplingHandler.apply(v1Request).cast(Object.class);
 		};
 	}
 

mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java

@@ -190,6 +190,8 @@ class SyncSpec {
 
 		private Function<CreateMessageRequest, CreateMessageResult> samplingHandler;
 
+		private Function<McpSchema.CreateMessageWithToolsRequest, McpSchema.CreateMessageWithToolsResult> samplingWithToolsHandler;
+
 		private Function<ElicitFormRequest, ElicitResult> formElicitationHandler;
 
 		private Function<ElicitUrlRequest, ElicitResult> urlElicitationHandler;
@@ -310,6 +312,23 @@ public SyncSpec sampling(Function<CreateMessageRequest, CreateMessageResult> sam
 			return this;
 		}
 
+		/**
+		 * Sets a sampling handler that supports tool use (SEP-1577) and registers the
+		 * {@code sampling.tools} client capability. Use this instead of
+		 * {@link #sampling(Function)} when the server may include {@code tools} and
+		 * {@code toolChoice} in its {@code sampling/createMessage} requests.
+		 * @param samplingWithToolsHandler A function that processes sampling-with-tools
+		 * requests and returns results. Must not be null.
+		 * @return This builder instance for method chaining
+		 * @throws IllegalArgumentException if samplingWithToolsHandler is null
+		 */
+		public SyncSpec samplingWithTools(
+				Function<McpSchema.CreateMessageWithToolsRequest, McpSchema.CreateMessageWithToolsResult> samplingWithToolsHandler) {
+			Assert.notNull(samplingWithToolsHandler, "Sampling-with-tools handler must not be null");
+			this.samplingWithToolsHandler = samplingWithToolsHandler;
+			return this;
+		}
+
 		/**
 		 * Sets a custom elicitation handler for processing elicitation message requests.
 		 * The elicitation handler can modify or validate messages before they are sent to
@@ -554,7 +573,8 @@ public McpSyncClient build() {
 					this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
 					this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
 					this.elicitationCompleteConsumers, this.samplingHandler, this.formElicitationHandler,
-					this.urlElicitationHandler, this.enableCallToolSchemaCaching, this.applyElicitationDefaults);
+					this.urlElicitationHandler, this.enableCallToolSchemaCaching, this.applyElicitationDefaults,
+					this.samplingWithToolsHandler);
 
 			McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
 
@@ -611,6 +631,8 @@ class AsyncSpec {
 
 		private Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler;
 
+		private Function<McpSchema.CreateMessageWithToolsRequest, Mono<McpSchema.CreateMessageWithToolsResult>> samplingWithToolsHandler;
+
 		private Function<ElicitFormRequest, Mono<ElicitResult>> formElicitationHandler;
 
 		private Function<ElicitUrlRequest, Mono<ElicitResult>> urlElicitationHandler;
@@ -729,6 +751,23 @@ public AsyncSpec sampling(Function<CreateMessageRequest, Mono<CreateMessageResul
 			return this;
 		}
 
+		/**
+		 * Sets a sampling handler that supports tool use (SEP-1577) and registers the
+		 * {@code sampling.tools} client capability. Use this instead of
+		 * {@link #sampling(Function)} when the server may include {@code tools} and
+		 * {@code toolChoice} in its {@code sampling/createMessage} requests.
+		 * @param samplingWithToolsHandler A function that processes sampling-with-tools
+		 * requests and returns results. Must not be null.
+		 * @return This builder instance for method chaining
+		 * @throws IllegalArgumentException if samplingWithToolsHandler is null
+		 */
+		public AsyncSpec samplingWithTools(
+				Function<McpSchema.CreateMessageWithToolsRequest, Mono<McpSchema.CreateMessageWithToolsResult>> samplingWithToolsHandler) {
+			Assert.notNull(samplingWithToolsHandler, "Sampling-with-tools handler must not be null");
+			this.samplingWithToolsHandler = samplingWithToolsHandler;
+			return this;
+		}
+
 		/**
 		 * Sets a custom elicitation handler for processing elicitation message requests.
 		 * The elicitation handler can modify or validate messages before they are sent to
@@ -964,8 +1003,8 @@ public McpAsyncClient build() {
 							this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
 							this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
 							this.elicitationCompleteConsumers, this.samplingHandler, this.formElicitationHandler,
-							this.urlElicitationHandler, this.enableCallToolSchemaCaching,
-							this.applyElicitationDefaults));
+							this.urlElicitationHandler, this.enableCallToolSchemaCaching, this.applyElicitationDefaults,
+							this.samplingWithToolsHandler));
 		}
 
 	}