apply review suggestions

Author: ochafik

src/mcp/server/session.py

@@ -47,6 +47,7 @@ async def handle_list_prompts(ctx: RequestContext) -> list[types.Prompt]:
 
 import mcp.types as types
 from mcp.server.models import InitializationOptions
+from mcp.shared.exceptions import McpError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import (
     BaseSession,
@@ -264,27 +265,34 @@ async def create_message(
                 types.ClientCapabilities(sampling=types.SamplingCapability(tools=types.SamplingToolsCapability()))
             )
             if not has_tools_cap:
-                from mcp.shared.exceptions import McpError
-
                 raise McpError(
                     types.ErrorData(
                         code=types.INVALID_PARAMS,
                         message="Client does not support sampling tools capability",
                     )
                 )
 
-        if messages and tools:
+        # Validate tool_use/tool_result message structure per SEP-1577:
+        # https://github.com/modelcontextprotocol/modelcontextprotocol/issues/1577
+        # This validation runs regardless of whether `tools` is in this request,
+        # since a tool loop continuation may omit `tools` while still containing
+        # tool_result content that must match previous tool_use.
+        if messages:
             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 are not matching any tool_use from the previous message")
+                    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.toolUseId for c in last_content if c.type == "tool_result"}

tests/server/test_session.py

@@ -337,15 +337,15 @@ async def test_create_message_tool_result_validation():
                 )
 
             # Case 2: tool_result without previous message
-            with pytest.raises(ValueError, match="not matching any tool_use"):
+            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="not matching any tool_use"):
+            with pytest.raises(ValueError, match="do not match any tool_use"):
                 await session.create_message(
                     messages=[
                         types.SamplingMessage(role="user", content=text),
@@ -371,8 +371,10 @@ async def test_create_message_tool_result_validation():
                 )
 
             # Case 5: text-only message with tools (no tool_results) - passes validation
-            # This covers branch 261->266 (has_tool_results=False) and 266->272
-            # We use move_on_after since send_request will block waiting for response
+            # 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)],
@@ -381,7 +383,8 @@ async def test_create_message_tool_result_validation():
                 )
 
             # Case 6: valid matching tool_result/tool_use IDs - passes validation
-            # This covers branch 269->272 (IDs match, no error raised)
+            # 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=[
@@ -393,6 +396,18 @@ async def test_create_message_tool_result_validation():
                     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
+                )
+
 
 @pytest.mark.anyio
 async def test_create_message_without_tools_capability():