Implement SEP-1577: Sampling With Tools

Add support for tool calling during sampling requests, enabling MCP servers
to execute agentic workflows using client LLM capabilities.

Key changes:
- Add ToolUseContent type for assistant tool invocation requests
- Add ToolResultContent type for tool execution results
- Add ToolChoice type to control tool usage behavior
- Add UserMessage and AssistantMessage types for role-specific messages
- Extend SamplingMessage to support tool content (backward compatible)
- Add SamplingToolsCapability for capability negotiation
- Update CreateMessageRequestParams with tools and toolChoice fields
- Update CreateMessageResult to support tool use content
- Update StopReason to include "toolUse" value
- Add comprehensive unit tests for all new types

The implementation maintains backward compatibility by keeping SamplingMessage
as a flexible BaseModel while adding more specific UserMessage and
AssistantMessage types for type-safe tool interactions.

All new types follow existing patterns:
- Use Pydantic V2 BaseModel
- Allow extra fields with ConfigDict(extra="allow")
- Include proper docstrings and field descriptions
- Support optional fields where appropriate

Github-Issue: #1577

Author: ochafik

src/mcp/__init__.py

@@ -5,6 +5,7 @@
 from .server.stdio import stdio_server
 from .shared.exceptions import McpError
 from .types import (
+    AssistantMessage,
     CallToolRequest,
     ClientCapabilities,
     ClientNotification,
@@ -42,6 +43,7 @@
     ResourceUpdatedNotification,
     RootsCapability,
     SamplingMessage,
+    SamplingToolsCapability,
     ServerCapabilities,
     ServerNotification,
     ServerRequest,
@@ -50,21 +52,27 @@
     StopReason,
     SubscribeRequest,
     Tool,
+    ToolChoice,
+    ToolResultContent,
     ToolsCapability,
+    ToolUseContent,
     UnsubscribeRequest,
+    UserMessage,
 )
 from .types import (
     Role as SamplingRole,
 )
 
 __all__ = [
+    "AssistantMessage",
     "CallToolRequest",
     "ClientCapabilities",
     "ClientNotification",
     "ClientRequest",
     "ClientResult",
     "ClientSession",
     "ClientSessionGroup",
+    "CompleteRequest",
     "CreateMessageRequest",
     "CreateMessageResult",
     "ErrorData",
@@ -77,6 +85,7 @@
     "InitializedNotification",
     "JSONRPCError",
     "JSONRPCRequest",
+    "JSONRPCResponse",
     "ListPromptsRequest",
     "ListPromptsResult",
     "ListResourcesRequest",
@@ -91,12 +100,13 @@
     "PromptsCapability",
     "ReadResourceRequest",
     "ReadResourceResult",
+    "Resource",
     "ResourcesCapability",
     "ResourceUpdatedNotification",
-    "Resource",
     "RootsCapability",
     "SamplingMessage",
     "SamplingRole",
+    "SamplingToolsCapability",
     "ServerCapabilities",
     "ServerNotification",
     "ServerRequest",
@@ -107,10 +117,12 @@
     "StopReason",
     "SubscribeRequest",
     "Tool",
+    "ToolChoice",
+    "ToolResultContent",
     "ToolsCapability",
+    "ToolUseContent",
     "UnsubscribeRequest",
+    "UserMessage",
     "stdio_client",
     "stdio_server",
-    "CompleteRequest",
-    "JSONRPCResponse",
 ]

src/mcp/types.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar
+from typing import Annotated, Any, Generic, Literal, TypeAlias, TypeVar, Union
 
 from pydantic import BaseModel, ConfigDict, Field, FileUrl, RootModel
 from pydantic.networks import AnyUrl, UrlConstraints
@@ -256,19 +256,52 @@ class SamplingCapability(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
+class SamplingToolsCapability(BaseModel):
+    """
+    Capability indicating support for tool calling during sampling.
+
+    When present in ClientCapabilities.sampling, indicates that the client
+    supports the tools and toolChoice parameters in sampling requests.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+
 class ElicitationCapability(BaseModel):
     """Capability for elicitation operations."""
 
     model_config = ConfigDict(extra="allow")
 
 
+class SamplingCapabilityNested(BaseModel):
+    """
+    Nested structure for sampling capabilities, allowing fine-grained capability advertisement.
+    """
+
+    context: SamplingCapability | None = None
+    """
+    Present if the client supports non-'none' values for includeContext parameter.
+    SOFT-DEPRECATED: New implementations should use tools parameter instead.
+    """
+    tools: SamplingToolsCapability | None = None
+    """
+    Present if the client supports tools and toolChoice parameters in sampling requests.
+    Presence indicates full tool calling support during sampling.
+    """
+    model_config = ConfigDict(extra="allow")
+
+
 class ClientCapabilities(BaseModel):
     """Capabilities a client may support."""
 
     experimental: dict[str, dict[str, Any]] | None = None
     """Experimental, non-standard capabilities that the client supports."""
-    sampling: SamplingCapability | None = None
-    """Present if the client supports sampling from an LLM."""
+    sampling: SamplingCapabilityNested | SamplingCapability | None = None
+    """
+    Present if the client supports sampling from an LLM.
+    Can be a structured object (SamplingCapabilityNested) with fine-grained capabilities,
+    or a simple marker object (SamplingCapability) for backward compatibility.
+    """
     elicitation: ElicitationCapability | None = None
     """Present if the client supports elicitation from the user."""
     roots: RootsCapability | None = None
@@ -742,11 +775,141 @@ class AudioContent(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
+class ToolUseContent(BaseModel):
+    """
+    Content representing an assistant's request to invoke a tool.
+
+    This content type appears in assistant messages when the LLM wants to call a tool
+    during sampling. The server should execute the tool and return a ToolResultContent
+    in the next user message.
+    """
+
+    type: Literal["tool_use"]
+    """Discriminator for tool use content."""
+
+    name: str
+    """The name of the tool to invoke. Must match a tool name from the request's tools array."""
+
+    id: str
+    """Unique identifier for this tool call, used to correlate with ToolResultContent."""
+
+    input: dict[str, Any]
+    """Arguments to pass to the tool. Must conform to the tool's inputSchema."""
+
+    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
+    """
+    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
+    for notes on _meta usage.
+    """
+    model_config = ConfigDict(extra="allow")
+
+
+class ToolResultContent(BaseModel):
+    """
+    Content representing the result of a tool execution.
+
+    This content type appears in user messages as a response to a ToolUseContent
+    from the assistant. It contains the output of executing the requested tool.
+    """
+
+    type: Literal["tool_result"]
+    """Discriminator for tool result content."""
+
+    toolUseId: str
+    """The unique identifier that corresponds to the tool call's id field."""
+
+    content: list[Union[TextContent, ImageContent, AudioContent, "EmbeddedResource"]] = []
+    """
+    A list of content objects representing the tool result.
+    Defaults to empty list if not provided.
+    """
+
+    structuredContent: dict[str, Any] | None = None
+    """
+    Optional structured tool output that matches the tool's outputSchema (if defined).
+    """
+
+    isError: bool | None = None
+    """Whether the tool execution resulted in an error."""
+
+    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
+    """
+    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
+    for notes on _meta usage.
+    """
+    model_config = ConfigDict(extra="allow")
+
+
 class SamplingMessage(BaseModel):
-    """Describes a message issued to or received from an LLM API."""
+    """
+    Describes a message issued to or received from an LLM API.
+
+    For backward compatibility, this class accepts any role and any content type.
+    For type-safe usage with tool calling, use UserMessage or AssistantMessage instead.
+    """
 
     role: Role
-    content: TextContent | ImageContent | AudioContent
+    content: (
+        TextContent
+        | ImageContent
+        | AudioContent
+        | ToolUseContent
+        | ToolResultContent
+        | list[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
+    )
+    """
+    Message content. Can be a single content block or an array of content blocks
+    for multi-modal messages and tool interactions.
+    """
+    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
+    """
+    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
+    for notes on _meta usage.
+    """
+    model_config = ConfigDict(extra="allow")
+
+
+# Type aliases for role-specific messages
+UserMessageContent: TypeAlias = TextContent | ImageContent | AudioContent | ToolResultContent
+"""Content types allowed in user messages during sampling."""
+
+AssistantMessageContent: TypeAlias = TextContent | ImageContent | AudioContent | ToolUseContent
+"""Content types allowed in assistant messages during sampling."""
+
+
+class UserMessage(BaseModel):
+    """
+    A message from the user (server) in a sampling conversation.
+
+    User messages can include tool results in response to assistant tool use requests.
+    """
+
+    role: Literal["user"]
+    content: UserMessageContent | list[UserMessageContent]
+    """Message content. Can be a single content block or an array for multi-modal messages."""
+    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
+    """
+    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
+    for notes on _meta usage.
+    """
+    model_config = ConfigDict(extra="allow")
+
+
+class AssistantMessage(BaseModel):
+    """
+    A message from the assistant (LLM) in a sampling conversation.
+
+    Assistant messages can include tool use requests when the LLM wants to call tools.
+    """
+
+    role: Literal["assistant"]
+    content: AssistantMessageContent | list[AssistantMessageContent]
+    """Message content. Can be a single content block or an array for multi-modal messages."""
+    meta: dict[str, Any] | None = Field(alias="_meta", default=None)
+    """
+    See [MCP specification](https://github.com/modelcontextprotocol/modelcontextprotocol/blob/47339c03c143bb4ec01a26e721a1b8fe66634ebe/docs/specification/draft/basic/index.mdx#general-fields)
+    for notes on _meta usage.
+    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -1035,6 +1198,31 @@ class ModelPreferences(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
+class ToolChoice(BaseModel):
+    """
+    Controls tool usage behavior during sampling.
+
+    Allows the server to specify whether and how the LLM should use tools
+    in its response.
+    """
+
+    mode: Literal["auto", "required", "none"] | None = None
+    """
+    Controls when tools are used:
+    - "auto": Model decides whether to use tools (default)
+    - "required": Model MUST use at least one tool before completing
+    - "none": Model should not use tools
+    """
+
+    disable_parallel_tool_use: bool | None = None
+    """
+    If true, the model should not use multiple tools in parallel.
+    Some models may ignore this hint. Default: false (parallel use enabled).
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+
 class CreateMessageRequestParams(RequestParams):
     """Parameters for creating a message."""
 
@@ -1057,6 +1245,16 @@ class CreateMessageRequestParams(RequestParams):
     stopSequences: list[str] | None = None
     metadata: dict[str, Any] | None = None
     """Optional metadata to pass through to the LLM provider."""
+    tools: list["Tool"] | None = None
+    """
+    Tool definitions for the LLM to use during sampling.
+    Requires clientCapabilities.sampling.tools to be present.
+    """
+    toolChoice: ToolChoice | None = None
+    """
+    Controls tool usage behavior.
+    Requires clientCapabilities.sampling.tools and the tools parameter to be present.
+    """
     model_config = ConfigDict(extra="allow")
 
 
@@ -1067,18 +1265,26 @@ class CreateMessageRequest(Request[CreateMessageRequestParams, Literal["sampling
     params: CreateMessageRequestParams
 
 
-StopReason = Literal["endTurn", "stopSequence", "maxTokens"] | str
+StopReason = Literal["endTurn", "stopSequence", "maxTokens", "toolUse"] | str
 
 
 class CreateMessageResult(Result):
     """The client's response to a sampling/create_message request from the server."""
 
-    role: Role
-    content: TextContent | ImageContent | AudioContent
+    role: Literal["assistant"]
+    """The role is always 'assistant' in responses from the LLM."""
+    content: AssistantMessageContent | list[AssistantMessageContent]
+    """
+    Response content from the assistant. May be a single content block or an array.
+    May include ToolUseContent if stopReason is 'toolUse'.
+    """
     model: str
     """The name of the model that generated the message."""
     stopReason: StopReason | None = None
-    """The reason why sampling stopped, if known."""
+    """
+    The reason why sampling stopped, if known.
+    'toolUse' indicates the model wants to use a tool.
+    """
 
 
 class ResourceTemplateReference(BaseModel):