Implement SEP-1577 - Sampling With Tools (#1594)

Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>
Co-authored-by: Claude <noreply@anthropic.com>

tests/server/test_session.py

@@ -9,6 +9,7 @@ from mcp.server import Server
 from mcp.server.lowlevel import NotificationOptions
 from mcp.server.models import InitializationOptions
 from mcp.server.session import ServerSession
+from mcp.shared.exceptions import McpError
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import RequestResponder
 from mcp.types import (
@@ -288,6 +289,184 @@ async def test_ping_request_before_initialization():
     assert ping_response_id == 42
 
 
+@pytest.mark.anyio
+async def test_create_message_tool_result_validation():
+    """Test tool_use/tool_result validation in create_message."""
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+
+    async with (
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        async with ServerSession(
+            client_to_server_receive,
+            server_to_client_send,
+            InitializationOptions(
+                server_name="test",
+                server_version="0.1.0",
+                capabilities=ServerCapabilities(),
+            ),
+        ) as session:
+            # Set up client params with sampling.tools capability for the test
+            session._client_params = types.InitializeRequestParams(
+                protocolVersion=types.LATEST_PROTOCOL_VERSION,
+                capabilities=types.ClientCapabilities(
+                    sampling=types.SamplingCapability(tools=types.SamplingToolsCapability())
+                ),
+                clientInfo=types.Implementation(name="test", version="1.0"),
+            )
+
+            tool = types.Tool(name="test_tool", inputSchema={"type": "object"})
+            text = types.TextContent(type="text", text="hello")
+            tool_use = types.ToolUseContent(type="tool_use", id="call_1", name="test_tool", input={})
+            tool_result = types.ToolResultContent(type="tool_result", toolUseId="call_1", content=[])
+
+            # Case 1: tool_result mixed with other content
+            with pytest.raises(ValueError, match="only tool_result content"):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="user", content=text),
+                        types.SamplingMessage(role="assistant", content=tool_use),
+                        types.SamplingMessage(role="user", content=[tool_result, text]),  # mixed!
+                    ],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 2: tool_result without previous message
+            with pytest.raises(ValueError, match="requires a previous message"):
+                await session.create_message(
+                    messages=[types.SamplingMessage(role="user", content=tool_result)],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 3: tool_result without previous tool_use
+            with pytest.raises(ValueError, match="do not match any tool_use"):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="user", content=text),
+                        types.SamplingMessage(role="user", content=tool_result),
+                    ],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 4: mismatched tool IDs
+            with pytest.raises(ValueError, match="ids of tool_result blocks and tool_use blocks"):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="user", content=text),
+                        types.SamplingMessage(role="assistant", content=tool_use),
+                        types.SamplingMessage(
+                            role="user",
+                            content=types.ToolResultContent(type="tool_result", toolUseId="wrong_id", content=[]),
+                        ),
+                    ],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 5: text-only message with tools (no tool_results) - passes validation
+            # Covers has_tool_results=False branch.
+            # We use move_on_after because validation happens synchronously before
+            # send_request, which would block indefinitely waiting for a response.
+            # The timeout lets validation pass, then cancels the blocked send.
+            with anyio.move_on_after(0.01):
+                await session.create_message(
+                    messages=[types.SamplingMessage(role="user", content=text)],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 6: valid matching tool_result/tool_use IDs - passes validation
+            # Covers tool_use_ids == tool_result_ids branch.
+            # (see Case 5 comment for move_on_after explanation)
+            with anyio.move_on_after(0.01):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="user", content=text),
+                        types.SamplingMessage(role="assistant", content=tool_use),
+                        types.SamplingMessage(role="user", content=tool_result),
+                    ],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+
+            # Case 7: validation runs even without `tools` parameter
+            # (tool loop continuation may omit tools while containing tool_result)
+            with pytest.raises(ValueError, match="do not match any tool_use"):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="user", content=text),
+                        types.SamplingMessage(role="user", content=tool_result),
+                    ],
+                    max_tokens=100,
+                    # Note: no tools parameter
+                )
+
+            # Case 8: empty messages list - skips validation entirely
+            # Covers the `if messages:` branch (line 280->302)
+            with anyio.move_on_after(0.01):
+                await session.create_message(
+                    messages=[],
+                    max_tokens=100,
+                )
+
+
+@pytest.mark.anyio
+async def test_create_message_without_tools_capability():
+    """Test that create_message raises McpError when tools are provided without capability."""
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage | Exception](1)
+
+    async with (
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        async with ServerSession(
+            client_to_server_receive,
+            server_to_client_send,
+            InitializationOptions(
+                server_name="test",
+                server_version="0.1.0",
+                capabilities=ServerCapabilities(),
+            ),
+        ) as session:
+            # Set up client params WITHOUT sampling.tools capability
+            session._client_params = types.InitializeRequestParams(
+                protocolVersion=types.LATEST_PROTOCOL_VERSION,
+                capabilities=types.ClientCapabilities(sampling=types.SamplingCapability()),
+                clientInfo=types.Implementation(name="test", version="1.0"),
+            )
+
+            tool = types.Tool(name="test_tool", inputSchema={"type": "object"})
+            text = types.TextContent(type="text", text="hello")
+
+            # Should raise McpError when tools are provided but client lacks capability
+            with pytest.raises(McpError) as exc_info:
+                await session.create_message(
+                    messages=[types.SamplingMessage(role="user", content=text)],
+                    max_tokens=100,
+                    tools=[tool],
+                )
+            assert "does not support sampling tools capability" in exc_info.value.error.message
+
+            # Should also raise McpError when tool_choice is provided
+            with pytest.raises(McpError) as exc_info:
+                await session.create_message(
+                    messages=[types.SamplingMessage(role="user", content=text)],
+                    max_tokens=100,
+                    tool_choice=types.ToolChoice(mode="auto"),
+                )
+            assert "does not support sampling tools capability" in exc_info.value.error.message
+
+
 @pytest.mark.anyio
 async def test_other_requests_blocked_before_initialization():
     """Test that non-ping requests are still blocked before initialization."""

tests/shared/test_streamable_http.py

@@ -210,7 +210,10 @@ class ServerTest(Server):  # pragma: no cover
                 )
 
                 # Return the sampling result in the tool response
-                response = sampling_result.content.text if sampling_result.content.type == "text" else None
+                if all(c.type == "text" for c in sampling_result.content_as_list):
+                    response = "\n".join(c.text for c in sampling_result.content_as_list if c.type == "text")
+                else:
+                    response = str(sampling_result.content)
                 return [
                     TextContent(
                         type="text",
@@ -1239,7 +1242,8 @@ async def test_streamablehttp_server_sampling(basic_server: None, basic_server_u
         nonlocal sampling_callback_invoked, captured_message_params
         sampling_callback_invoked = True
         captured_message_params = params
-        message_received = params.messages[0].content.text if params.messages[0].content.type == "text" else None
+        msg_content = params.messages[0].content_as_list[0]
+        message_received = msg_content.text if msg_content.type == "text" else None
 
         return types.CreateMessageResult(
             role="assistant",

tests/test_types.py

@@ -1,16 +1,26 @@
+from typing import Any
+
 import pytest
 
 from mcp.types import (
     LATEST_PROTOCOL_VERSION,
     ClientCapabilities,
     ClientRequest,
+    CreateMessageRequestParams,
+    CreateMessageResult,
     Implementation,
     InitializeRequest,
     InitializeRequestParams,
     JSONRPCMessage,
     JSONRPCRequest,
     ListToolsResult,
+    SamplingCapability,
+    SamplingMessage,
+    TextContent,
     Tool,
+    ToolChoice,
+    ToolResultContent,
+    ToolUseContent,
 )
 
 
@@ -60,6 +70,215 @@ async def test_method_initialization():
     assert initialize_request.params.protocolVersion == LATEST_PROTOCOL_VERSION
 
 
+@pytest.mark.anyio
+async def test_tool_use_content():
+    """Test ToolUseContent type for SEP-1577."""
+    tool_use_data = {
+        "type": "tool_use",
+        "name": "get_weather",
+        "id": "call_abc123",
+        "input": {"location": "San Francisco", "unit": "celsius"},
+    }
+
+    tool_use = ToolUseContent.model_validate(tool_use_data)
+    assert tool_use.type == "tool_use"
+    assert tool_use.name == "get_weather"
+    assert tool_use.id == "call_abc123"
+    assert tool_use.input == {"location": "San Francisco", "unit": "celsius"}
+
+    # Test serialization
+    serialized = tool_use.model_dump(by_alias=True, exclude_none=True)
+    assert serialized["type"] == "tool_use"
+    assert serialized["name"] == "get_weather"
+
+
+@pytest.mark.anyio
+async def test_tool_result_content():
+    """Test ToolResultContent type for SEP-1577."""
+    tool_result_data = {
+        "type": "tool_result",
+        "toolUseId": "call_abc123",
+        "content": [{"type": "text", "text": "It's 72°F in San Francisco"}],
+        "isError": False,
+    }
+
+    tool_result = ToolResultContent.model_validate(tool_result_data)
+    assert tool_result.type == "tool_result"
+    assert tool_result.toolUseId == "call_abc123"
+    assert len(tool_result.content) == 1
+    assert tool_result.isError is False
+
+    # Test with empty content (should default to [])
+    minimal_result_data = {"type": "tool_result", "toolUseId": "call_xyz"}
+    minimal_result = ToolResultContent.model_validate(minimal_result_data)
+    assert minimal_result.content == []
+
+
+@pytest.mark.anyio
+async def test_tool_choice():
+    """Test ToolChoice type for SEP-1577."""
+    # Test with mode
+    tool_choice_data = {"mode": "required"}
+    tool_choice = ToolChoice.model_validate(tool_choice_data)
+    assert tool_choice.mode == "required"
+
+    # Test with minimal data (all fields optional)
+    minimal_choice = ToolChoice.model_validate({})
+    assert minimal_choice.mode 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_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 = SamplingMessage.model_validate(user_msg_data)
+    assert user_msg.role == "user"
+    assert isinstance(user_msg.content, TextContent)
+
+    # Test with array of content including tool result
+    multi_content_data: dict[str, Any] = {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Here's the result:"},
+            {"type": "tool_result", "toolUseId": "call_123", "content": []},
+        ],
+    }
+    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_sampling_message_with_assistant_role():
+    """Test SamplingMessage with assistant role for SEP-1577."""
+    # Test with tool use content
+    assistant_msg_data = {
+        "role": "assistant",
+        "content": {
+            "type": "tool_use",
+            "name": "search",
+            "id": "call_456",
+            "input": {"query": "MCP protocol"},
+        },
+    }
+    assistant_msg = SamplingMessage.model_validate(assistant_msg_data)
+    assert assistant_msg.role == "assistant"
+    assert isinstance(assistant_msg.content, ToolUseContent)
+
+    # Test with array of mixed content
+    multi_content_data: dict[str, Any] = {
+        "role": "assistant",
+        "content": [
+            {"type": "text", "text": "Let me search for that..."},
+            {"type": "tool_use", "name": "search", "id": "call_789", "input": {}},
+        ],
+    }
+    multi_msg = SamplingMessage.model_validate(multi_content_data)
+    assert isinstance(multi_msg.content, list)
+    assert len(multi_msg.content) == 2
+
+
+@pytest.mark.anyio
+async def test_sampling_message_backward_compatibility():
+    """Test that SamplingMessage maintains backward compatibility."""
+    # Old-style message (single content, no tools)
+    old_style_data = {"role": "user", "content": {"type": "text", "text": "Hello"}}
+    old_msg = SamplingMessage.model_validate(old_style_data)
+    assert old_msg.role == "user"
+    assert isinstance(old_msg.content, TextContent)
+
+    # New-style message with tool content
+    new_style_data: dict[str, Any] = {
+        "role": "assistant",
+        "content": {"type": "tool_use", "name": "test", "id": "call_1", "input": {}},
+    }
+    new_msg = SamplingMessage.model_validate(new_style_data)
+    assert new_msg.role == "assistant"
+    assert isinstance(new_msg.content, ToolUseContent)
+
+    # Array content
+    array_style_data: dict[str, Any] = {
+        "role": "user",
+        "content": [{"type": "text", "text": "Result:"}, {"type": "tool_result", "toolUseId": "call_1", "content": []}],
+    }
+    array_msg = SamplingMessage.model_validate(array_style_data)
+    assert isinstance(array_msg.content, list)
+
+
+@pytest.mark.anyio
+async def test_create_message_request_params_with_tools():
+    """Test CreateMessageRequestParams with tools for SEP-1577."""
+    tool = Tool(
+        name="get_weather",
+        description="Get weather information",
+        inputSchema={"type": "object", "properties": {"location": {"type": "string"}}},
+    )
+
+    params = CreateMessageRequestParams(
+        messages=[SamplingMessage(role="user", content=TextContent(type="text", text="What's the weather?"))],
+        maxTokens=1000,
+        tools=[tool],
+        toolChoice=ToolChoice(mode="auto"),
+    )
+
+    assert params.tools is not None
+    assert len(params.tools) == 1
+    assert params.tools[0].name == "get_weather"
+    assert params.toolChoice is not None
+    assert params.toolChoice.mode == "auto"
+
+
+@pytest.mark.anyio
+async def test_create_message_result_with_tool_use():
+    """Test CreateMessageResult with tool use content for SEP-1577."""
+    result_data = {
+        "role": "assistant",
+        "content": {"type": "tool_use", "name": "search", "id": "call_123", "input": {"query": "test"}},
+        "model": "claude-3",
+        "stopReason": "toolUse",
+    }
+
+    result = CreateMessageResult.model_validate(result_data)
+    assert result.role == "assistant"
+    assert isinstance(result.content, ToolUseContent)
+    assert result.stopReason == "toolUse"
+    assert result.model == "claude-3"
+
+    # Test content_as_list with single content (covers else branch)
+    content_list = result.content_as_list
+    assert len(content_list) == 1
+    assert content_list[0] == result.content
+
+
+@pytest.mark.anyio
+async def test_client_capabilities_with_sampling_tools():
+    """Test ClientCapabilities with nested sampling capabilities for SEP-1577."""
+    # New structured format
+    capabilities_data: dict[str, Any] = {
+        "sampling": {"tools": {}},
+    }
+    capabilities = ClientCapabilities.model_validate(capabilities_data)
+    assert capabilities.sampling is not None
+    assert isinstance(capabilities.sampling, SamplingCapability)
+    assert capabilities.sampling.tools is not None
+
+    # With both context and tools
+    full_capabilities_data: dict[str, Any] = {"sampling": {"context": {}, "tools": {}}}
+    full_caps = ClientCapabilities.model_validate(full_capabilities_data)
+    assert isinstance(full_caps.sampling, SamplingCapability)
+    assert full_caps.sampling.context is not None
+    assert full_caps.sampling.tools is not None
+
+
 def test_tool_preserves_json_schema_2020_12_fields():
     """Verify that JSON Schema 2020-12 keywords are preserved in Tool.inputSchema.