Add coverage tests for experimental tasks code

- Remove unused _get_event method from InMemoryTaskMessageQueue
- Add Resolver error case tests (set_result/set_exception when completed)
- Add message_queue tests for peek empty, double-check race condition
- Add in_memory_task_store test for wait_for_update with nonexistent task
- Add task_support test for accessing task_group before run()
- Add task_result_handler tests for exception handling and missing original_id

Author: maxisbey

src/mcp/shared/experimental/tasks/message_queue.py

@@ -171,12 +171,6 @@ def _get_queue(self, task_id: str) -> list[QueuedMessage]:
             self._queues[task_id] = []
         return self._queues[task_id]
 
-    def _get_event(self, task_id: str) -> anyio.Event:
-        """Get or create the wait event for a task."""
-        if task_id not in self._events:
-            self._events[task_id] = anyio.Event()
-        return self._events[task_id]
-
     async def enqueue(self, task_id: str, message: QueuedMessage) -> None:
         """Add a message to the queue."""
         queue = self._get_queue(task_id)

tests/experimental/tasks/server/test_run_task_flow.py

@@ -233,6 +233,17 @@ async def work(task: ServerTaskContext) -> CallToolResult:
         await experimental.run_task(work)
 
 
+@pytest.mark.anyio
+async def test_task_support_task_group_before_run_raises() -> None:
+    """Test that accessing task_group before run() raises RuntimeError."""
+    from mcp.server.experimental.task_support import TaskSupport
+
+    task_support = TaskSupport.in_memory()
+
+    with pytest.raises(RuntimeError, match="TaskSupport not running"):
+        _ = task_support.task_group
+
+
 @pytest.mark.anyio
 async def test_run_task_without_session_raises() -> None:
     """Test that run_task raises when session is not available."""

tests/experimental/tasks/server/test_store.py

@@ -321,9 +321,11 @@ async def test_terminal_status_allows_same_status(store: InMemoryTaskStore) -> N
     assert updated.statusMessage == "Updated message"
 
 
-# =============================================================================
-# cancel_task helper function tests
-# =============================================================================
+@pytest.mark.anyio
+async def test_wait_for_update_nonexistent_raises(store: InMemoryTaskStore) -> None:
+    """Test that wait_for_update raises for nonexistent task."""
+    with pytest.raises(ValueError, match="not found"):
+        await store.wait_for_update("nonexistent-task-id")
 
 
 @pytest.mark.anyio

tests/experimental/tasks/server/test_task_result_handler.py

@@ -253,3 +253,94 @@ async def test_deliver_registers_resolver_for_request_messages(
 
     assert "inner-req-1" in handler._pending_requests
     assert handler._pending_requests["inner-req-1"] is resolver
+
+
+@pytest.mark.anyio
+async def test_deliver_skips_resolver_registration_when_no_original_id(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that _deliver_queued_messages skips resolver registration when original_request_id is None."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    resolver: Resolver[dict[str, Any]] = Resolver()
+    queued_msg = QueuedMessage(
+        type="request",
+        message=JSONRPCRequest(
+            jsonrpc="2.0",
+            id="inner-req-1",
+            method="elicitation/create",
+            params={},
+        ),
+        resolver=resolver,
+        original_request_id=None,  # No original request ID
+    )
+    await queue.enqueue(task.taskId, queued_msg)
+
+    mock_session = Mock()
+    mock_session.send_message = AsyncMock()
+
+    await handler._deliver_queued_messages(task.taskId, mock_session, "outer-req-1")
+
+    # Resolver should NOT be registered since original_request_id is None
+    assert len(handler._pending_requests) == 0
+    # But the message should still be sent
+    mock_session.send_message.assert_called_once()
+
+
+@pytest.mark.anyio
+async def test_wait_for_task_update_handles_store_exception(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that _wait_for_task_update handles store exception gracefully."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    # Make wait_for_update raise an exception
+    async def failing_wait(task_id: str) -> None:
+        raise RuntimeError("Store error")
+
+    store.wait_for_update = failing_wait  # type: ignore[method-assign]
+
+    # Queue a message to unblock the race via the queue path
+    async def enqueue_later() -> None:
+        await anyio.sleep(0.01)
+        await queue.enqueue(
+            task.taskId,
+            QueuedMessage(
+                type="notification",
+                message=JSONRPCRequest(
+                    jsonrpc="2.0",
+                    id="notif-1",
+                    method="test/notification",
+                    params={},
+                ),
+            ),
+        )
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(enqueue_later)
+        # This should complete via the queue path even though store raises
+        await handler._wait_for_task_update(task.taskId)
+
+
+@pytest.mark.anyio
+async def test_wait_for_task_update_handles_queue_exception(
+    store: InMemoryTaskStore, queue: InMemoryTaskMessageQueue, handler: TaskResultHandler
+) -> None:
+    """Test that _wait_for_task_update handles queue exception gracefully."""
+    task = await store.create_task(TaskMetadata(ttl=60000), task_id="test-task")
+
+    # Make wait_for_message raise an exception
+    async def failing_wait(task_id: str) -> None:
+        raise RuntimeError("Queue error")
+
+    queue.wait_for_message = failing_wait  # type: ignore[method-assign]
+
+    # Update the store to unblock the race via the store path
+    async def update_later() -> None:
+        await anyio.sleep(0.01)
+        await store.update_task(task.taskId, status="completed")
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(update_later)
+        # This should complete via the store path even though queue raises
+        await handler._wait_for_task_update(task.taskId)

tests/experimental/tasks/test_message_queue.py

@@ -247,3 +247,85 @@ async def wait_for_notification() -> None:
             tg.start_soon(notify_when_ready)
 
         assert notified is True
+
+    @pytest.mark.anyio
+    async def test_peek_empty_queue_returns_none(self, queue: InMemoryTaskMessageQueue) -> None:
+        """Peek on empty queue returns None."""
+        result = await queue.peek("nonexistent-task")
+        assert result is None
+
+    @pytest.mark.anyio
+    async def test_wait_for_message_double_check_race_condition(self, queue: InMemoryTaskMessageQueue) -> None:
+        """wait_for_message returns early if message arrives after event creation but before wait."""
+        task_id = "task-1"
+
+        # To test the double-check path (lines 223-225), we need a message to arrive
+        # after the event is created (line 220) but before event.wait() (line 228).
+        # We simulate this by injecting a message before is_empty is called the second time.
+
+        original_is_empty = queue.is_empty
+        call_count = 0
+
+        async def is_empty_with_injection(tid: str) -> bool:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 2 and tid == task_id:
+                # Before second check, inject a message - this simulates a message
+                # arriving between event creation and the double-check
+                queue._queues[task_id] = [QueuedMessage(type="request", message=make_request())]
+            return await original_is_empty(tid)
+
+        queue.is_empty = is_empty_with_injection  # type: ignore[method-assign]
+
+        # Should return immediately due to double-check finding the message
+        with anyio.fail_after(1):
+            await queue.wait_for_message(task_id)
+
+
+class TestResolver:
+    @pytest.mark.anyio
+    async def test_set_result_and_wait(self) -> None:
+        """Test basic set_result and wait flow."""
+        resolver: Resolver[str] = Resolver()
+
+        resolver.set_result("hello")
+        result = await resolver.wait()
+
+        assert result == "hello"
+        assert resolver.done()
+
+    @pytest.mark.anyio
+    async def test_set_exception_and_wait(self) -> None:
+        """Test set_exception raises on wait."""
+        resolver: Resolver[str] = Resolver()
+
+        resolver.set_exception(ValueError("test error"))
+
+        with pytest.raises(ValueError, match="test error"):
+            await resolver.wait()
+
+        assert resolver.done()
+
+    @pytest.mark.anyio
+    async def test_set_result_when_already_completed_raises(self) -> None:
+        """Test that set_result raises if resolver already completed."""
+        resolver: Resolver[str] = Resolver()
+        resolver.set_result("first")
+
+        with pytest.raises(RuntimeError, match="already completed"):
+            resolver.set_result("second")
+
+    @pytest.mark.anyio
+    async def test_set_exception_when_already_completed_raises(self) -> None:
+        """Test that set_exception raises if resolver already completed."""
+        resolver: Resolver[str] = Resolver()
+        resolver.set_result("done")
+
+        with pytest.raises(RuntimeError, match="already completed"):
+            resolver.set_exception(ValueError("too late"))
+
+    @pytest.mark.anyio
+    async def test_done_returns_false_before_completion(self) -> None:
+        """Test done() returns False before any result is set."""
+        resolver: Resolver[str] = Resolver()
+        assert resolver.done() is False