Achieve 100% test coverage for tasks code

- Add tests for validation.py (check_sampling_tools_capability,
  validate_sampling_tools, validate_tool_use_result_messages)
- Add flow test for elicit_url() in ServerTaskContext
- Add pragma no cover comments to defensive _meta checks in builder
  methods (model_dump never includes _meta with current types)
- Fix test code to use assertions instead of conditional branches
- Add pragma no branch to polling loops in test scenarios

Author: maxisbey

src/mcp/server/session.py

@@ -517,7 +517,8 @@ def _build_elicit_form_request(
 
         # Add related-task metadata if associated with a parent task
         if related_task_id is not None:
-            if "_meta" not in params_data:
+            # Defensive: model_dump() never includes _meta, but guard against future changes
+            if "_meta" not in params_data:  # pragma: no cover
                 params_data["_meta"] = {}
             params_data["_meta"]["io.modelcontextprotocol/related-task"] = {"taskId": related_task_id}
 
@@ -559,7 +560,8 @@ def _build_elicit_url_request(
 
         # Add related-task metadata if associated with a parent task
         if related_task_id is not None:
-            if "_meta" not in params_data:
+            # Defensive: model_dump() never includes _meta, but guard against future changes
+            if "_meta" not in params_data:  # pragma: no cover
                 params_data["_meta"] = {}
             params_data["_meta"]["io.modelcontextprotocol/related-task"] = {"taskId": related_task_id}
 
@@ -626,7 +628,8 @@ def _build_create_message_request(
 
         # Add related-task metadata if associated with a parent task
         if related_task_id is not None:
-            if "_meta" not in params_data:
+            # Defensive: model_dump() never includes _meta, but guard against future changes
+            if "_meta" not in params_data:  # pragma: no cover
                 params_data["_meta"] = {}
             params_data["_meta"]["io.modelcontextprotocol/related-task"] = {"taskId": related_task_id}
 

tests/experimental/tasks/server/test_server_task_context.py

@@ -370,6 +370,78 @@ async def run_elicit() -> None:
     store.cleanup()
 
 
+@pytest.mark.anyio
+async def test_elicit_url_queues_request_and_waits_for_response() -> None:
+    """Test that elicit_url() queues request and waits for response."""
+    import anyio
+
+    from mcp.types import JSONRPCRequest
+
+    store = InMemoryTaskStore()
+    queue = InMemoryTaskMessageQueue()
+    handler = TaskResultHandler(store, queue)
+    task = await store.create_task(TaskMetadata(ttl=60000))
+
+    mock_session = Mock()
+    mock_session.check_client_capability = Mock(return_value=True)
+    mock_session._build_elicit_url_request = Mock(
+        return_value=JSONRPCRequest(
+            jsonrpc="2.0",
+            id="test-url-req-1",
+            method="elicitation/create",
+            params={"message": "Authorize", "url": "https://example.com", "elicitationId": "123", "mode": "url"},
+        )
+    )
+
+    ctx = ServerTaskContext(
+        task=task,
+        store=store,
+        session=mock_session,
+        queue=queue,
+        handler=handler,
+    )
+
+    elicit_result = None
+
+    async def run_elicit_url() -> None:
+        nonlocal elicit_result
+        elicit_result = await ctx.elicit_url(
+            message="Authorize",
+            url="https://example.com/oauth",
+            elicitation_id="oauth-123",
+        )
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(run_elicit_url)
+
+        # Wait for request to be queued
+        await queue.wait_for_message(task.taskId)
+
+        # Verify task is in input_required status
+        updated_task = await store.get_task(task.taskId)
+        assert updated_task is not None
+        assert updated_task.status == "input_required"
+
+        # Dequeue and simulate response
+        msg = await queue.dequeue(task.taskId)
+        assert msg is not None
+        assert msg.resolver is not None
+
+        # Resolve with mock elicitation response (URL mode just returns action)
+        msg.resolver.set_result({"action": "accept"})
+
+    # Verify result
+    assert elicit_result is not None
+    assert elicit_result.action == "accept"
+
+    # Verify task is back to working
+    final_task = await store.get_task(task.taskId)
+    assert final_task is not None
+    assert final_task.status == "working"
+
+    store.cleanup()
+
+
 @pytest.mark.anyio
 async def test_create_message_queues_request_and_waits_for_response() -> None:
     """Test that create_message() queues request and waits for response."""

tests/experimental/tasks/test_elicitation_scenarios.py

@@ -95,8 +95,8 @@ async def handle_get_task_result(
     ) -> GetTaskPayloadResult | ErrorData:
         """Handle tasks/result from server."""
         event = task_complete_events.get(params.taskId)
-        if event:
-            await event.wait()
+        assert event is not None, f"No completion event for task: {params.taskId}"
+        await event.wait()
         result = await client_task_store.get_result(params.taskId)
         assert result is not None, f"Result not found for task: {params.taskId}"
         return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True))
@@ -163,8 +163,8 @@ async def handle_get_task_result(
     ) -> GetTaskPayloadResult | ErrorData:
         """Handle tasks/result from server."""
         event = task_complete_events.get(params.taskId)
-        if event:
-            await event.wait()
+        assert event is not None, f"No completion event for task: {params.taskId}"
+        await event.wait()
         result = await client_task_store.get_result(params.taskId)
         assert result is not None, f"Result not found for task: {params.taskId}"
         return GetTaskPayloadResult.model_validate(result.model_dump(by_alias=True))
@@ -417,9 +417,12 @@ async def run_client() -> None:
             assert create_result.task.status == "working"
 
             # Poll until input_required, then call tasks/result
-            async for status in client_session.experimental.poll_task(task_id):
-                if status.status == "input_required":
+            found_input_required = False
+            async for status in client_session.experimental.poll_task(task_id):  # pragma: no branch
+                if status.status == "input_required":  # pragma: no branch
+                    found_input_required = True
                     break
+            assert found_input_required, "Expected to see input_required status"
 
             # This will deliver the elicitation and get the response
             final_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
@@ -523,9 +526,12 @@ async def run_client() -> None:
             assert create_result.task.status == "working"
 
             # Poll until input_required or terminal, then call tasks/result
-            async for status in client_session.experimental.poll_task(task_id):
-                if status.status == "input_required" or is_terminal(status.status):
+            found_expected_status = False
+            async for status in client_session.experimental.poll_task(task_id):  # pragma: no branch
+                if status.status == "input_required" or is_terminal(status.status):  # pragma: no branch
+                    found_expected_status = True
                     break
+            assert found_expected_status, "Expected to see input_required or terminal status"
 
             # This will deliver the task-augmented elicitation,
             # server will poll client, and eventually return the tool result
@@ -581,9 +587,8 @@ async def handle_call_tool(name: str, arguments: dict[str, Any]) -> CallToolResu
             ttl=60000,
         )
 
-        response_text = ""
-        if isinstance(result.content, TextContent):
-            response_text = result.content.text
+        assert isinstance(result.content, TextContent), "Expected TextContent response"
+        response_text = result.content.text
 
         tool_result.append(response_text)
         return CallToolResult(content=[TextContent(type="text", text=response_text)])
@@ -671,9 +676,8 @@ async def work(task: ServerTaskContext) -> CallToolResult:
                 ttl=60000,
             )
 
-            response_text = ""
-            if isinstance(result.content, TextContent):
-                response_text = result.content.text
+            assert isinstance(result.content, TextContent), "Expected TextContent response"
+            response_text = result.content.text
 
             work_completed.set()
             return CallToolResult(content=[TextContent(type="text", text=response_text)])
@@ -710,9 +714,12 @@ async def run_client() -> None:
             assert create_result.task.status == "working"
 
             # Poll until input_required or terminal
-            async for status in client_session.experimental.poll_task(task_id):
-                if status.status == "input_required" or is_terminal(status.status):
+            found_expected_status = False
+            async for status in client_session.experimental.poll_task(task_id):  # pragma: no branch
+                if status.status == "input_required" or is_terminal(status.status):  # pragma: no branch
+                    found_expected_status = True
                     break
+            assert found_expected_status, "Expected to see input_required or terminal status"
 
             final_result = await client_session.experimental.get_task_result(task_id, CallToolResult)
 

tests/server/test_validation.py

@@ -0,0 +1,141 @@
+"""Tests for server validation functions."""
+
+import pytest
+
+from mcp.server.validation import (
+    check_sampling_tools_capability,
+    validate_sampling_tools,
+    validate_tool_use_result_messages,
+)
+from mcp.shared.exceptions import McpError
+from mcp.types import (
+    ClientCapabilities,
+    SamplingCapability,
+    SamplingMessage,
+    SamplingToolsCapability,
+    TextContent,
+    Tool,
+    ToolChoice,
+    ToolResultContent,
+    ToolUseContent,
+)
+
+
+class TestCheckSamplingToolsCapability:
+    """Tests for check_sampling_tools_capability function."""
+
+    def test_returns_false_when_caps_none(self) -> None:
+        """Returns False when client_caps is None."""
+        assert check_sampling_tools_capability(None) is False
+
+    def test_returns_false_when_sampling_none(self) -> None:
+        """Returns False when client_caps.sampling is None."""
+        caps = ClientCapabilities()
+        assert check_sampling_tools_capability(caps) is False
+
+    def test_returns_false_when_tools_none(self) -> None:
+        """Returns False when client_caps.sampling.tools is None."""
+        caps = ClientCapabilities(sampling=SamplingCapability())
+        assert check_sampling_tools_capability(caps) is False
+
+    def test_returns_true_when_tools_present(self) -> None:
+        """Returns True when sampling.tools is present."""
+        caps = ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
+        assert check_sampling_tools_capability(caps) is True
+
+
+class TestValidateSamplingTools:
+    """Tests for validate_sampling_tools function."""
+
+    def test_no_error_when_tools_none(self) -> None:
+        """No error when tools and tool_choice are None."""
+        validate_sampling_tools(None, None, None)  # Should not raise
+
+    def test_raises_when_tools_provided_but_no_capability(self) -> None:
+        """Raises McpError when tools provided but client doesn't support."""
+        tool = Tool(name="test", inputSchema={"type": "object"})
+        with pytest.raises(McpError) as exc_info:
+            validate_sampling_tools(None, [tool], None)
+        assert "sampling tools capability" in str(exc_info.value)
+
+    def test_raises_when_tool_choice_provided_but_no_capability(self) -> None:
+        """Raises McpError when tool_choice provided but client doesn't support."""
+        with pytest.raises(McpError) as exc_info:
+            validate_sampling_tools(None, None, ToolChoice(mode="auto"))
+        assert "sampling tools capability" in str(exc_info.value)
+
+    def test_no_error_when_capability_present(self) -> None:
+        """No error when client has sampling.tools capability."""
+        caps = ClientCapabilities(sampling=SamplingCapability(tools=SamplingToolsCapability()))
+        tool = Tool(name="test", inputSchema={"type": "object"})
+        validate_sampling_tools(caps, [tool], ToolChoice(mode="auto"))  # Should not raise
+
+
+class TestValidateToolUseResultMessages:
+    """Tests for validate_tool_use_result_messages function."""
+
+    def test_no_error_for_empty_messages(self) -> None:
+        """No error when messages list is empty."""
+        validate_tool_use_result_messages([])  # Should not raise
+
+    def test_no_error_for_simple_text_messages(self) -> None:
+        """No error for simple text messages."""
+        messages = [
+            SamplingMessage(role="user", content=TextContent(type="text", text="Hello")),
+            SamplingMessage(role="assistant", content=TextContent(type="text", text="Hi")),
+        ]
+        validate_tool_use_result_messages(messages)  # Should not raise
+
+    def test_raises_when_tool_result_mixed_with_other_content(self) -> None:
+        """Raises when tool_result is mixed with other content types."""
+        messages = [
+            SamplingMessage(
+                role="user",
+                content=[
+                    ToolResultContent(type="tool_result", toolUseId="123"),
+                    TextContent(type="text", text="also this"),
+                ],
+            ),
+        ]
+        with pytest.raises(ValueError, match="only tool_result content"):
+            validate_tool_use_result_messages(messages)
+
+    def test_raises_when_tool_result_without_previous_tool_use(self) -> None:
+        """Raises when tool_result appears without preceding tool_use."""
+        messages = [
+            SamplingMessage(
+                role="user",
+                content=ToolResultContent(type="tool_result", toolUseId="123"),
+            ),
+        ]
+        with pytest.raises(ValueError, match="previous message containing tool_use"):
+            validate_tool_use_result_messages(messages)
+
+    def test_raises_when_tool_result_ids_dont_match_tool_use(self) -> None:
+        """Raises when tool_result IDs don't match tool_use IDs."""
+        messages = [
+            SamplingMessage(
+                role="assistant",
+                content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+            ),
+            SamplingMessage(
+                role="user",
+                content=ToolResultContent(type="tool_result", toolUseId="tool-2"),
+            ),
+        ]
+        with pytest.raises(ValueError, match="do not match"):
+            validate_tool_use_result_messages(messages)
+
+    def test_no_error_when_tool_result_matches_tool_use(self) -> None:
+        """No error when tool_result IDs match tool_use IDs."""
+        messages = [
+            SamplingMessage(
+                role="assistant",
+                content=ToolUseContent(type="tool_use", id="tool-1", name="test", input={}),
+            ),
+            SamplingMessage(
+                role="user",
+                content=ToolResultContent(type="tool_result", toolUseId="tool-1"),
+            ),
+        ]
+        validate_tool_use_result_messages(messages)  # Should not raise