Simplify sampling types to match draft spec

- Replace UserMessage/AssistantMessage with unified SamplingMessage
- Add SamplingMessageContentBlock type alias for content blocks
- Update CreateMessageResult to use SamplingMessageContentBlock
- Remove disable_parallel_tool_use from ToolChoice (not in spec)
- Add ResourceLink to ToolResultContent.content for consistency

This aligns with the draft spec where SamplingMessage uses a single
content type union rather than separate UserMessage/AssistantMessage
classes. The CreateMessageResult now extends SamplingMessage's pattern
and supports both single content blocks and arrays.

Author: ochafik

src/mcp/__init__.py

@@ -5,7 +5,6 @@
 from .server.stdio import stdio_server
 from .shared.exceptions import McpError
 from .types import (
-    AssistantMessage,
     CallToolRequest,
     ClientCapabilities,
     ClientNotification,
@@ -45,6 +44,7 @@
     SamplingCapability,
     SamplingContextCapability,
     SamplingMessage,
+    SamplingMessageContentBlock,
     SamplingToolsCapability,
     ServerCapabilities,
     ServerNotification,
@@ -59,14 +59,12 @@
     ToolsCapability,
     ToolUseContent,
     UnsubscribeRequest,
-    UserMessage,
 )
 from .types import (
     Role as SamplingRole,
 )
 
 __all__ = [
-    "AssistantMessage",
     "CallToolRequest",
     "ClientCapabilities",
     "ClientNotification",
@@ -109,6 +107,7 @@
     "SamplingCapability",
     "SamplingContextCapability",
     "SamplingMessage",
+    "SamplingMessageContentBlock",
     "SamplingRole",
     "SamplingToolsCapability",
     "ServerCapabilities",
@@ -126,7 +125,6 @@
     "ToolsCapability",
     "ToolUseContent",
     "UnsubscribeRequest",
-    "UserMessage",
     "stdio_client",
     "stdio_server",
 ]

src/mcp/types.py

@@ -822,7 +822,7 @@ class ToolResultContent(BaseModel):
     toolUseId: str
     """The unique identifier that corresponds to the tool call's id field."""
 
-    content: list[Union[TextContent, ImageContent, AudioContent, "EmbeddedResource"]] = []
+    content: list[Union[TextContent, ImageContent, AudioContent, "ResourceLink", "EmbeddedResource"]] = []
     """
     A list of content objects representing the tool result.
     Defaults to empty list if not provided.
@@ -844,23 +844,17 @@ class ToolResultContent(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
-class SamplingMessage(BaseModel):
-    """
-    Describes a message issued to or received from an LLM API.
+SamplingMessageContentBlock: TypeAlias = (
+    TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
+)
+"""Content block types allowed in sampling messages."""
 
-    For backward compatibility, this class accepts any role and any content type.
-    For type-safe usage with tool calling, use UserMessage or AssistantMessage instead.
-    """
+
+class SamplingMessage(BaseModel):
+    """Describes a message issued to or received from an LLM API."""
 
     role: Role
-    content: (
-        TextContent
-        | ImageContent
-        | AudioContent
-        | ToolUseContent
-        | ToolResultContent
-        | list[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
-    )
+    content: SamplingMessageContentBlock | list[SamplingMessageContentBlock]
     """
     Message content. Can be a single content block or an array of content blocks
     for multi-modal messages and tool interactions.
@@ -873,50 +867,6 @@ class SamplingMessage(BaseModel):
     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")
-
-
 class EmbeddedResource(BaseModel):
     """
     The contents of a resource, embedded into a prompt or tool call result.
@@ -1218,12 +1168,6 @@ class ToolChoice(BaseModel):
     - "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")
 
 
@@ -1275,11 +1219,11 @@ class CreateMessageRequest(Request[CreateMessageRequestParams, Literal["sampling
 class CreateMessageResult(Result):
     """The client's response to a sampling/create_message request from the server."""
 
-    role: Literal["assistant"]
-    """The role is always 'assistant' in responses from the LLM."""
-    content: AssistantMessageContent | list[AssistantMessageContent]
+    role: Role
+    """The role of the message sender (typically 'assistant' for LLM responses)."""
+    content: SamplingMessageContentBlock | list[SamplingMessageContentBlock]
     """
-    Response content from the assistant. May be a single content block or an array.
+    Response content. May be a single content block or an array.
     May include ToolUseContent if stopReason is 'toolUse'.
     """
     model: str

tests/test_types.py

@@ -2,7 +2,6 @@
 
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
-    AssistantMessage,
     ClientCapabilities,
     ClientRequest,
     CreateMessageRequestParams,
@@ -19,7 +18,6 @@
     ToolChoice,
     ToolResultContent,
     ToolUseContent,
-    UserMessage,
 )
 
 
@@ -117,23 +115,28 @@ async def test_tool_result_content():
 async def test_tool_choice():
     """Test ToolChoice type for SEP-1577."""
     # Test with mode
-    tool_choice_data = {"mode": "required", "disable_parallel_tool_use": True}
+    tool_choice_data = {"mode": "required"}
     tool_choice = ToolChoice.model_validate(tool_choice_data)
     assert tool_choice.mode == "required"
-    assert tool_choice.disable_parallel_tool_use is True
 
     # Test with minimal data (all fields optional)
     minimal_choice = ToolChoice.model_validate({})
     assert minimal_choice.mode is None
-    assert minimal_choice.disable_parallel_tool_use is None
+
+    # Test different modes
+    auto_choice = ToolChoice.model_validate({"mode": "auto"})
+    assert auto_choice.mode == "auto"
+
+    none_choice = ToolChoice.model_validate({"mode": "none"})
+    assert none_choice.mode == "none"
 
 
 @pytest.mark.anyio
-async def test_user_message():
-    """Test UserMessage type for SEP-1577."""
+async def test_sampling_message_with_user_role():
+    """Test SamplingMessage with user role for SEP-1577."""
     # Test with single content
     user_msg_data = {"role": "user", "content": {"type": "text", "text": "Hello"}}
-    user_msg = UserMessage.model_validate(user_msg_data)
+    user_msg = SamplingMessage.model_validate(user_msg_data)
     assert user_msg.role == "user"
     assert isinstance(user_msg.content, TextContent)
 
@@ -145,15 +148,15 @@ async def test_user_message():
             {"type": "tool_result", "toolUseId": "call_123", "content": []},
         ],
     }
-    multi_msg = UserMessage.model_validate(multi_content_data)
+    multi_msg = SamplingMessage.model_validate(multi_content_data)
     assert multi_msg.role == "user"
     assert isinstance(multi_msg.content, list)
     assert len(multi_msg.content) == 2
 
 
 @pytest.mark.anyio
-async def test_assistant_message():
-    """Test AssistantMessage type for SEP-1577."""
+async def test_sampling_message_with_assistant_role():
+    """Test SamplingMessage with assistant role for SEP-1577."""
     # Test with tool use content
     assistant_msg_data = {
         "role": "assistant",
@@ -164,7 +167,7 @@ async def test_assistant_message():
             "input": {"query": "MCP protocol"},
         },
     }
-    assistant_msg = AssistantMessage.model_validate(assistant_msg_data)
+    assistant_msg = SamplingMessage.model_validate(assistant_msg_data)
     assert assistant_msg.role == "assistant"
     assert isinstance(assistant_msg.content, ToolUseContent)
 
@@ -176,7 +179,7 @@ async def test_assistant_message():
             {"type": "tool_use", "name": "search", "id": "call_789", "input": {}},
         ],
     }
-    multi_msg = AssistantMessage.model_validate(multi_content_data)
+    multi_msg = SamplingMessage.model_validate(multi_content_data)
     assert isinstance(multi_msg.content, list)
     assert len(multi_msg.content) == 2