Unify sampling and elicitation code paths with shared validation

This refactoring ensures all sampling and elicitation code paths use
consistent validation and support the same features.

Sampling changes:
- Add shared validation module (mcp/server/validation.py) with
  validate_sampling_tools() and validate_tool_use_result_messages()
- Add tools and tool_choice parameters to all sampling methods:
  - _build_create_message_request()
  - ExperimentalServerSessionFeatures.create_message_as_task()
  - ServerTaskContext.create_message()
  - ServerTaskContext.create_message_as_task()
- Refactor ServerSession.create_message() to use shared validation

Elicitation changes:
- Rename _build_elicit_request to _build_elicit_form_request for clarity
- Add _build_elicit_url_request() for URL mode elicitation
- Add ServerTaskContext.elicit_url() so URL elicitation can be used
  from inside task-augmented tool calls (e.g., for OAuth flows)

This fixes a gap where task-augmented code paths were missing:
- tools/tool_choice parameters for sampling
- URL mode for elicitation

Author: maxisbey

src/mcp/server/experimental/session_features.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, Any, TypeVar
 
 import mcp.types as types
+from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
 from mcp.shared.experimental.tasks.capabilities import (
     require_task_augmented_elicitation,
     require_task_augmented_sampling,
@@ -156,6 +157,8 @@ async def create_message_as_task(
         stop_sequences: list[str] | None = None,
         metadata: dict[str, Any] | None = None,
         model_preferences: types.ModelPreferences | None = None,
+        tools: list[types.Tool] | None = None,
+        tool_choice: types.ToolChoice | None = None,
     ) -> types.CreateMessageResult:
         """
         Send a task-augmented sampling request and poll until complete.
@@ -173,15 +176,20 @@ async def create_message_as_task(
             stop_sequences: Stop sequences
             metadata: Additional metadata
             model_preferences: Model selection preferences
+            tools: Optional list of tools the LLM can use during sampling
+            tool_choice: Optional control over tool usage behavior
 
         Returns:
             The sampling result from the client
 
         Raises:
-            McpError: If client doesn't support task-augmented sampling
+            McpError: If client doesn't support task-augmented sampling or tools
+            ValueError: If tool_use or tool_result message structure is invalid
         """
         client_caps = self._session.client_params.capabilities if self._session.client_params else None
         require_task_augmented_sampling(client_caps)
+        validate_sampling_tools(client_caps, tools, tool_choice)
+        validate_tool_use_result_messages(messages)
 
         create_result = await self._session.send_request(
             types.ServerRequest(
@@ -195,6 +203,8 @@ async def create_message_as_task(
                         stopSequences=stop_sequences,
                         metadata=metadata,
                         modelPreferences=model_preferences,
+                        tools=tools,
+                        toolChoice=tool_choice,
                         task=types.TaskMetadata(ttl=ttl),
                     )
                 )

src/mcp/server/experimental/task_context.py

@@ -13,6 +13,7 @@
 
 from mcp.server.experimental.task_result_handler import TaskResultHandler
 from mcp.server.session import ServerSession
+from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
 from mcp.shared.exceptions import McpError
 from mcp.shared.experimental.tasks.capabilities import (
     require_task_augmented_elicitation,
@@ -44,6 +45,8 @@
     TaskMetadata,
     TaskStatusNotification,
     TaskStatusNotificationParams,
+    Tool,
+    ToolChoice,
 )
 
 
@@ -231,7 +234,7 @@ async def elicit(
         await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
 
         # Build the request using session's helper
-        request = self._session._build_elicit_request(  # pyright: ignore[reportPrivateUsage]
+        request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
             message=message,
             requestedSchema=requestedSchema,
             related_task_id=self.task_id,
@@ -263,6 +266,77 @@ async def elicit(
             await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
             raise
 
+    async def elicit_url(
+        self,
+        message: str,
+        url: str,
+        elicitation_id: str,
+    ) -> ElicitResult:
+        """
+        Send a URL mode elicitation request via the task message queue.
+
+        This directs the user to an external URL for out-of-band interactions
+        like OAuth flows, credential collection, or payment processing.
+
+        This method:
+        1. Checks client capability
+        2. Updates task status to "input_required"
+        3. Queues the elicitation request
+        4. Waits for the response (delivered via tasks/result round-trip)
+        5. Updates task status back to "working"
+        6. Returns the result
+
+        Args:
+            message: Human-readable explanation of why the interaction is needed
+            url: The URL the user should navigate to
+            elicitation_id: Unique identifier for tracking this elicitation
+
+        Returns:
+            The client's response indicating acceptance, decline, or cancellation
+
+        Raises:
+            McpError: If client doesn't support elicitation capability
+            RuntimeError: If handler is not configured
+        """
+        self._check_elicitation_capability()
+
+        if self._handler is None:
+            raise RuntimeError("handler is required for elicit_url(). Pass handler= to ServerTaskContext.")
+
+        # Update status to input_required
+        await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
+
+        # Build the request using session's helper
+        request = self._session._build_elicit_url_request(  # pyright: ignore[reportPrivateUsage]
+            message=message,
+            url=url,
+            elicitation_id=elicitation_id,
+            related_task_id=self.task_id,
+        )
+        request_id: RequestId = request.id
+
+        # Create resolver and register with handler for response routing
+        resolver: Resolver[dict[str, Any]] = Resolver()
+        self._handler._pending_requests[request_id] = resolver  # pyright: ignore[reportPrivateUsage]
+
+        # Queue the request
+        queued = QueuedMessage(
+            type="request",
+            message=request,
+            resolver=resolver,
+            original_request_id=request_id,
+        )
+        await self._queue.enqueue(self.task_id, queued)
+
+        try:
+            # Wait for response (routed back via TaskResultHandler)
+            response_data = await resolver.wait()
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            return ElicitResult.model_validate(response_data)
+        except anyio.get_cancelled_exc_class():  # pragma: no cover
+            await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
+            raise
+
     async def create_message(
         self,
         messages: list[SamplingMessage],
@@ -274,6 +348,8 @@ async def create_message(
         stop_sequences: list[str] | None = None,
         metadata: dict[str, Any] | None = None,
         model_preferences: ModelPreferences | None = None,
+        tools: list[Tool] | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> CreateMessageResult:
         """
         Send a sampling request via the task message queue.
@@ -295,14 +371,20 @@ async def create_message(
             stop_sequences: Stop sequences
             metadata: Additional metadata
             model_preferences: Model selection preferences
+            tools: Optional list of tools the LLM can use during sampling
+            tool_choice: Optional control over tool usage behavior
 
         Returns:
             The sampling result from the client
 
         Raises:
-            McpError: If client doesn't support sampling capability
+            McpError: If client doesn't support sampling capability or tools
+            ValueError: If tool_use or tool_result message structure is invalid
         """
         self._check_sampling_capability()
+        client_caps = self._session.client_params.capabilities if self._session.client_params else None
+        validate_sampling_tools(client_caps, tools, tool_choice)
+        validate_tool_use_result_messages(messages)
 
         if self._handler is None:
             raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.")
@@ -320,6 +402,8 @@ async def create_message(
             stop_sequences=stop_sequences,
             metadata=metadata,
             model_preferences=model_preferences,
+            tools=tools,
+            tool_choice=tool_choice,
             related_task_id=self.task_id,
         )
         request_id: RequestId = request.id
@@ -386,7 +470,7 @@ async def elicit_as_task(
         await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
 
         # Build request WITH task field for task-augmented elicitation
-        request = self._session._build_elicit_request(  # pyright: ignore[reportPrivateUsage]
+        request = self._session._build_elicit_form_request(  # pyright: ignore[reportPrivateUsage]
             message=message,
             requestedSchema=requestedSchema,
             related_task_id=self.task_id,
@@ -442,6 +526,8 @@ async def create_message_as_task(
         stop_sequences: list[str] | None = None,
         metadata: dict[str, Any] | None = None,
         model_preferences: ModelPreferences | None = None,
+        tools: list[Tool] | None = None,
+        tool_choice: ToolChoice | None = None,
     ) -> CreateMessageResult:
         """
         Send a task-augmented sampling request via the queue, then poll client.
@@ -461,16 +547,21 @@ async def create_message_as_task(
             stop_sequences: Stop sequences
             metadata: Additional metadata
             model_preferences: Model selection preferences
+            tools: Optional list of tools the LLM can use during sampling
+            tool_choice: Optional control over tool usage behavior
 
         Returns:
             The sampling result from the client
 
         Raises:
-            McpError: If client doesn't support task-augmented sampling
+            McpError: If client doesn't support task-augmented sampling or tools
+            ValueError: If tool_use or tool_result message structure is invalid
             RuntimeError: If handler is not configured
         """
         client_caps = self._session.client_params.capabilities if self._session.client_params else None
         require_task_augmented_sampling(client_caps)
+        validate_sampling_tools(client_caps, tools, tool_choice)
+        validate_tool_use_result_messages(messages)
 
         if self._handler is None:
             raise RuntimeError("handler is required for create_message_as_task()")
@@ -488,6 +579,8 @@ async def create_message_as_task(
             stop_sequences=stop_sequences,
             metadata=metadata,
             model_preferences=model_preferences,
+            tools=tools,
+            tool_choice=tool_choice,
             related_task_id=self.task_id,
             task=TaskMetadata(ttl=ttl),
         )

src/mcp/server/session.py

@@ -48,7 +48,7 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 import mcp.types as types
 from mcp.server.experimental.session_features import ExperimentalServerSessionFeatures
 from mcp.server.models import InitializationOptions
-from mcp.shared.exceptions import McpError
+from mcp.server.validation import validate_sampling_tools, validate_tool_use_result_messages
 from mcp.shared.experimental.tasks.capabilities import check_tasks_capability
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.response_router import ResponseRouter
@@ -293,47 +293,12 @@ async def create_message(
             The sampling result from the client.
 
         Raises:
-            McpError: If tool_use or tool_result blocks are misused when tools are provided.
+            McpError: If tools are provided but client doesn't support them.
+            ValueError: If tool_use or tool_result message structure is invalid.
         """
-
-        if tools is not None or tool_choice is not None:
-            has_tools_cap = self.check_client_capability(
-                types.ClientCapabilities(sampling=types.SamplingCapability(tools=types.SamplingToolsCapability()))
-            )
-            if not has_tools_cap:
-                raise McpError(
-                    types.ErrorData(
-                        code=types.INVALID_PARAMS,
-                        message="Client does not support sampling tools capability",
-                    )
-                )
-
-        # Validate tool_use/tool_result message structure per SEP-1577:
-        # https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577
-        # This validation runs regardless of whether `tools` is in this request,
-        # since a tool loop continuation may omit `tools` while still containing
-        # tool_result content that must match previous tool_use.
-        if messages:
-            last_content = messages[-1].content_as_list
-            has_tool_results = any(c.type == "tool_result" for c in last_content)
-
-            previous_content = messages[-2].content_as_list if len(messages) >= 2 else None
-            has_previous_tool_use = previous_content and any(c.type == "tool_use" for c in previous_content)
-
-            if has_tool_results:
-                # Per spec: "SamplingMessage with tool result content blocks
-                # MUST NOT contain other content types."
-                if any(c.type != "tool_result" for c in last_content):
-                    raise ValueError("The last message must contain only tool_result content if any is present")
-                if previous_content is None:
-                    raise ValueError("tool_result requires a previous message containing tool_use")
-                if not has_previous_tool_use:
-                    raise ValueError("tool_result blocks do not match any tool_use in the previous message")
-            if has_previous_tool_use and previous_content:
-                tool_use_ids = {c.id for c in previous_content if c.type == "tool_use"}
-                tool_result_ids = {c.toolUseId for c in last_content if c.type == "tool_result"}
-                if tool_use_ids != tool_result_ids:
-                    raise ValueError("ids of tool_result blocks and tool_use blocks from previous message do not match")
+        client_caps = self._client_params.capabilities if self._client_params else None
+        validate_sampling_tools(client_caps, tools, tool_choice)
+        validate_tool_use_result_messages(messages)
 
         return await self.send_request(
             request=types.ServerRequest(
@@ -525,14 +490,14 @@ async def send_elicit_complete(
     # by TaskContext to construct requests that will be queued instead of sent
     # directly, avoiding code duplication between ServerSession and TaskContext.
 
-    def _build_elicit_request(
+    def _build_elicit_form_request(
         self,
         message: str,
         requestedSchema: types.ElicitRequestedSchema,
         related_task_id: str | None = None,
         task: types.TaskMetadata | None = None,
     ) -> types.JSONRPCRequest:
-        """Build an elicitation request without sending it.
+        """Build a form mode elicitation request without sending it.
 
         Args:
             message: The message to present to the user
@@ -567,6 +532,48 @@ def _build_elicit_request(
             params=params_data,
         )
 
+    def _build_elicit_url_request(
+        self,
+        message: str,
+        url: str,
+        elicitation_id: str,
+        related_task_id: str | None = None,
+    ) -> types.JSONRPCRequest:
+        """Build a URL mode elicitation request without sending it.
+
+        Args:
+            message: Human-readable explanation of why the interaction is needed
+            url: The URL the user should navigate to
+            elicitation_id: Unique identifier for tracking this elicitation
+            related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata
+
+        Returns:
+            A JSONRPCRequest ready to be sent or queued
+        """
+        params = types.ElicitRequestURLParams(
+            message=message,
+            url=url,
+            elicitationId=elicitation_id,
+        )
+        params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True)
+
+        # Add related-task metadata if associated with a parent task
+        if related_task_id is not None:
+            if "_meta" not in params_data:
+                params_data["_meta"] = {}
+            params_data["_meta"]["io.modelcontextprotocol/related-task"] = {"taskId": related_task_id}
+
+        request_id = f"task-{related_task_id}-{id(params)}" if related_task_id else self._request_id
+        if related_task_id is None:
+            self._request_id += 1
+
+        return types.JSONRPCRequest(
+            jsonrpc="2.0",
+            id=request_id,
+            method="elicitation/create",
+            params=params_data,
+        )
+
     def _build_create_message_request(
         self,
         messages: list[types.SamplingMessage],
@@ -578,6 +585,8 @@ def _build_create_message_request(
         stop_sequences: list[str] | None = None,
         metadata: dict[str, Any] | None = None,
         model_preferences: types.ModelPreferences | None = None,
+        tools: list[types.Tool] | None = None,
+        tool_choice: types.ToolChoice | None = None,
         related_task_id: str | None = None,
         task: types.TaskMetadata | None = None,
     ) -> types.JSONRPCRequest:
@@ -592,6 +601,8 @@ def _build_create_message_request(
             stop_sequences: Optional stop sequences
             metadata: Optional metadata to pass through to the LLM provider
             model_preferences: Optional model selection preferences
+            tools: Optional list of tools the LLM can use during sampling
+            tool_choice: Optional control over tool usage behavior
             related_task_id: If provided, adds io.modelcontextprotocol/related-task metadata
             task: If provided, makes this a task-augmented request
 
@@ -607,6 +618,8 @@ def _build_create_message_request(
             stopSequences=stop_sequences,
             metadata=metadata,
             modelPreferences=model_preferences,
+            tools=tools,
+            toolChoice=tool_choice,
             task=task,
         )
         params_data = params.model_dump(by_alias=True, mode="json", exclude_none=True)