Validate tool_use / tool_result blocks in create_message

Author: ochafik

src/mcp/server/session.py

@@ -246,7 +246,41 @@ async def create_message(
 
         Returns:
             The sampling result from the client.
+
+        Raises:
+            McpError: If tool_use or tool_result blocks are misused
         """
+
+        if messages and tools:
+            last_msg_content = messages[-1].content
+            last_content = last_msg_content if isinstance(last_msg_content, list) else [last_msg_content]
+            has_tool_results = any(c.type == "tool_result" for c in last_content)
+
+            previous_content: list[types.SamplingMessageContentBlock] | None = None
+            if len(messages) >= 2:
+                prev_msg_content = messages[-2].content
+                previous_content = prev_msg_content if isinstance(prev_msg_content, list) else [prev_msg_content]
+            has_previous_tool_use = previous_content and any(c.type == "tool_use" for c in previous_content)
+
+            if has_tool_results:
+                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 len(messages) == 1:
+                    raise ValueError(
+                        "The last message cannot contain tool_result content if there is no previous message"
+                    )
+                if not has_previous_tool_use:
+                    raise ValueError(
+                        "The previous message must contain tool_use content if the last message contains tool_result content"
+                    )
+            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(
+                        "The tool_result content in the last message must correspond to all tool_use content in the previous message"
+                    )
+
         return await self.send_request(
             request=types.ServerRequest(
                 types.CreateMessageRequest(

tests/server/test_session.py

@@ -288,6 +288,80 @@ async def mock_client():
     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,
+    ):
+        session = ServerSession(
+            client_to_server_receive,
+            server_to_client_send,
+            InitializationOptions(
+                server_name="test",
+                server_version="0.1.0",
+                capabilities=ServerCapabilities(),
+            ),
+        )
+
+        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="no 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="previous message must contain 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="must correspond to all tool_use"):
+            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],
+            )
+
+
 @pytest.mark.anyio
 async def test_other_requests_blocked_before_initialization():
     """Test that non-ping requests are still blocked before initialization."""