Validate full sampling tool result history

Author: Epochex

src/mcp/server/validation.py

@@ -5,7 +5,7 @@
 """
 
 from mcp.shared.exceptions import MCPError
-from mcp.types import INVALID_PARAMS, ClientCapabilities, SamplingMessage, Tool, ToolChoice
+from mcp.types import INVALID_PARAMS, ClientCapabilities, SamplingMessage, SamplingMessageContentBlock, Tool, ToolChoice
 
 
 def check_sampling_tools_capability(client_caps: ClientCapabilities | None) -> bool:
@@ -53,6 +53,7 @@ def validate_tool_use_result_messages(messages: list[SamplingMessage]) -> None:
     1. Messages with tool_result content contain ONLY tool_result content
     2. tool_result messages are preceded by a message with tool_use
     3. tool_result IDs match the tool_use IDs from the previous message
+    4. Every tool_use message in the history is followed by matching tool_result content
 
     See: https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577
 
@@ -65,24 +66,26 @@ def validate_tool_use_result_messages(messages: list[SamplingMessage]) -> None:
     if not messages:
         return
 
-    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.tool_use_id 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")
+    previous_content: list[SamplingMessageContentBlock] | None = None
+    for content in (message.content_as_list for message in messages):
+        has_tool_results = any(c.type == "tool_result" for c in content)
+        previous_tool_use_ids: set[str] = set()
+        if previous_content is not None:
+            previous_tool_use_ids = {c.id for c in previous_content if c.type == "tool_use"}
+
+        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 content):
+                raise ValueError("A 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 previous_tool_use_ids:
+                raise ValueError("tool_result blocks do not match any tool_use in the previous message")
+
+        if previous_tool_use_ids:
+            tool_result_ids = {c.tool_use_id for c in content if c.type == "tool_result"}
+            if previous_tool_use_ids != tool_result_ids:
+                raise ValueError("ids of tool_result blocks and tool_use blocks from previous message do not match")
+
+        previous_content = content

tests/server/test_session.py

@@ -347,6 +347,21 @@ async def test_create_message_tool_result_validation():
                     tools=[tool],
                 )
 
+            # Case 4b: earlier mismatched tool result with a later plain message
+            with pytest.raises(ValueError, match="ids of tool_result blocks and tool_use blocks"):
+                await session.create_message(
+                    messages=[
+                        types.SamplingMessage(role="assistant", content=tool_use),
+                        types.SamplingMessage(
+                            role="user",
+                            content=types.ToolResultContent(type="tool_result", tool_use_id="wrong_id", content=[]),
+                        ),
+                        types.SamplingMessage(role="assistant", content=text),
+                    ],
+                    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

tests/server/test_validation.py

@@ -108,6 +108,27 @@ def test_validate_tool_use_result_messages_raises_when_tool_result_mixed_with_ot
         validate_tool_use_result_messages(messages)
 
 
+def test_validate_tool_use_result_messages_raises_for_earlier_mixed_tool_result() -> None:
+    """Raises when an earlier message mixes tool_result with other content."""
+    messages = [
+        SamplingMessage(
+            role="assistant",
+            content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+        ),
+        SamplingMessage(
+            role="user",
+            content=[
+                ToolResultContent(type="tool_result", tool_use_id="tool-1"),
+                TextContent(type="text", text="also this"),
+            ],
+        ),
+        SamplingMessage(role="assistant", content=TextContent(type="text", text="done")),
+    ]
+
+    with pytest.raises(ValueError, match="only tool_result content"):
+        validate_tool_use_result_messages(messages)
+
+
 def test_validate_tool_use_result_messages_raises_when_tool_result_without_previous_tool_use() -> None:
     """Raises when tool_result appears without preceding tool_use."""
     messages = [
@@ -136,6 +157,39 @@ def test_validate_tool_use_result_messages_raises_when_tool_result_ids_dont_matc
         validate_tool_use_result_messages(messages)
 
 
+def test_validate_tool_use_result_messages_raises_when_earlier_tool_result_ids_dont_match_tool_use() -> None:
+    """Raises when an earlier tool_result does not match the previous tool_use."""
+    messages = [
+        SamplingMessage(
+            role="assistant",
+            content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+        ),
+        SamplingMessage(
+            role="user",
+            content=ToolResultContent(type="tool_result", tool_use_id="tool-2"),
+        ),
+        SamplingMessage(role="assistant", content=TextContent(type="text", text="done")),
+    ]
+
+    with pytest.raises(ValueError, match="do not match"):
+        validate_tool_use_result_messages(messages)
+
+
+def test_validate_tool_use_result_messages_raises_when_tool_use_is_not_answered() -> None:
+    """Raises when a tool_use is followed by a non-tool_result message."""
+    messages = [
+        SamplingMessage(
+            role="assistant",
+            content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+        ),
+        SamplingMessage(role="user", content=TextContent(type="text", text="not a result")),
+        SamplingMessage(role="assistant", content=TextContent(type="text", text="done")),
+    ]
+
+    with pytest.raises(ValueError, match="do not match"):
+        validate_tool_use_result_messages(messages)
+
+
 def test_validate_tool_use_result_messages_no_error_when_tool_result_matches_tool_use() -> None:
     """No error when tool_result IDs match tool_use IDs."""
     messages = [
@@ -149,3 +203,34 @@ def test_validate_tool_use_result_messages_no_error_when_tool_result_matches_too
         ),
     ]
     validate_tool_use_result_messages(messages)  # Should not raise
+
+
+def test_validate_tool_use_result_messages_no_error_for_multiple_tool_pairs() -> None:
+    """No error when every tool_use in the history has a matching tool_result."""
+    messages = [
+        SamplingMessage(role="user", content=TextContent(type="text", text="first")),
+        SamplingMessage(
+            role="assistant",
+            content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+        ),
+        SamplingMessage(
+            role="user",
+            content=ToolResultContent(type="tool_result", tool_use_id="tool-1"),
+        ),
+        SamplingMessage(
+            role="assistant",
+            content=[
+                ToolUseContent(type="tool_use", id="tool-2", name="test", input={}),
+                ToolUseContent(type="tool_use", id="tool-3", name="test", input={}),
+            ],
+        ),
+        SamplingMessage(
+            role="user",
+            content=[
+                ToolResultContent(type="tool_result", tool_use_id="tool-3"),
+                ToolResultContent(type="tool_result", tool_use_id="tool-2"),
+            ],
+        ),
+    ]
+
+    validate_tool_use_result_messages(messages)