Add coverage tests and fix gaps

Coverage improvements:
- Add test_server_task_context.py with tests for ServerTaskContext
- Add tests for default task handlers (list, cancel, result, augmented sampling/elicitation)
- Add test for task-augmented handler capability building
- Add tests for elicit/create_message flows including cancellation
- Add test for meta parameter in call_tool_as_task
- Add test for model_immediate_response parameter
- Add tests for run_task error cases

Code changes:
- Remove backwards compat task_id parameter from ServerTaskContext (require task directly)
- Remove type: ignore from request_context.py
- Add pragma: no cover for async exception handlers (coverage limitation with async)

Author: maxisbey

src/mcp/server/experimental/request_context.py

@@ -220,7 +220,7 @@ async def work(task: ServerTaskContext) -> CallToolResult:
             session=self._session,
             queue=support.queue,
             handler=support.handler,
-        )  # type: ignore[call-arg]
+        )
 
         # Spawn the work
         async def execute() -> None:

src/mcp/server/experimental/task_context.py

@@ -15,7 +15,6 @@
 from mcp.server.session import ServerSession
 from mcp.shared.exceptions import McpError
 from mcp.shared.experimental.tasks.context import TaskContext
-from mcp.shared.experimental.tasks.helpers import create_task_state
 from mcp.shared.experimental.tasks.message_queue import QueuedMessage, TaskMessageQueue
 from mcp.shared.experimental.tasks.resolver import Resolver
 from mcp.shared.experimental.tasks.store import TaskStore
@@ -37,7 +36,6 @@
     SamplingMessage,
     ServerNotification,
     Task,
-    TaskMetadata,
     TaskStatusNotification,
     TaskStatusNotificationParams,
 )
@@ -70,8 +68,7 @@ async def my_task_work(task: ServerTaskContext) -> CallToolResult:
     def __init__(
         self,
         *,
-        task: Task | None = None,
-        task_id: str | None = None,
+        task: Task,
         store: TaskStore,
         session: ServerSession,
         queue: TaskMessageQueue,
@@ -81,23 +78,12 @@ def __init__(
         Create a ServerTaskContext.
 
         Args:
-            task: The Task object (provide either task or task_id)
-            task_id: The task ID to look up (provide either task or task_id)
+            task: The Task object
             store: The task store
             session: The server session
             queue: The message queue for elicitation/sampling
             handler: The result handler for response routing (required for elicit/create_message)
         """
-        if task is None and task_id is None:
-            raise ValueError("Must provide either task or task_id")
-        if task is not None and task_id is not None:
-            raise ValueError("Provide either task or task_id, not both")
-
-        # If task_id provided, create a minimal task object
-        # This is for backwards compatibility with tests that pass task_id
-        if task is None:
-            task = create_task_state(TaskMetadata(ttl=None), task_id=task_id)
-
         self._ctx = TaskContext(task=task, store=store)
         self._session = session
         self._queue = queue
@@ -264,7 +250,10 @@ async def elicit(
             response_data = await resolver.wait()
             await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
             return ElicitResult.model_validate(response_data)
-        except anyio.get_cancelled_exc_class():
+        except anyio.get_cancelled_exc_class():  # pragma: no cover
+            # Coverage can't track async exception handlers reliably.
+            # This path is tested in test_elicit_restores_status_on_cancellation
+            # which verifies status is restored to "working" after cancellation.
             await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
             raise
 
@@ -347,6 +336,9 @@ async def create_message(
             response_data = await resolver.wait()
             await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
             return CreateMessageResult.model_validate(response_data)
-        except anyio.get_cancelled_exc_class():
+        except anyio.get_cancelled_exc_class():  # pragma: no cover
+            # Coverage can't track async exception handlers reliably.
+            # This path is tested in test_create_message_restores_status_on_cancellation
+            # which verifies status is restored to "working" after cancellation.
             await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
             raise

tests/experimental/tasks/client/test_capabilities.py

@@ -251,3 +251,81 @@ async def mock_server():
     assert received_capabilities.tasks.cancel is not None
     # requests should be None since we didn't provide task-augmented handlers
     assert received_capabilities.tasks.requests is None
+
+
+@pytest.mark.anyio
+async def test_client_capabilities_with_task_augmented_handlers():
+    """Test that requests capability is built when augmented handlers are provided."""
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](1)
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](1)
+
+    received_capabilities: ClientCapabilities | None = None
+
+    # Define task-augmented handler
+    async def my_augmented_sampling_handler(
+        context: RequestContext[ClientSession, None],
+        params: types.CreateMessageRequestParams,
+        task_metadata: types.TaskMetadata,
+    ) -> types.CreateTaskResult | types.ErrorData:
+        return types.ErrorData(code=types.INVALID_REQUEST, message="Not implemented")
+
+    async def mock_server():
+        nonlocal received_capabilities
+
+        session_message = await client_to_server_receive.receive()
+        jsonrpc_request = session_message.message
+        assert isinstance(jsonrpc_request.root, JSONRPCRequest)
+        request = ClientRequest.model_validate(
+            jsonrpc_request.model_dump(by_alias=True, mode="json", exclude_none=True)
+        )
+        assert isinstance(request.root, InitializeRequest)
+        received_capabilities = request.root.params.capabilities
+
+        result = ServerResult(
+            InitializeResult(
+                protocolVersion=LATEST_PROTOCOL_VERSION,
+                capabilities=ServerCapabilities(),
+                serverInfo=Implementation(name="mock-server", version="0.1.0"),
+            )
+        )
+
+        async with server_to_client_send:
+            await server_to_client_send.send(
+                SessionMessage(
+                    JSONRPCMessage(
+                        JSONRPCResponse(
+                            jsonrpc="2.0",
+                            id=jsonrpc_request.root.id,
+                            result=result.model_dump(by_alias=True, mode="json", exclude_none=True),
+                        )
+                    )
+                )
+            )
+            await client_to_server_receive.receive()
+
+    # Provide task-augmented sampling handler
+    task_handlers = ExperimentalTaskHandlers(
+        augmented_sampling=my_augmented_sampling_handler,
+    )
+
+    async with (
+        ClientSession(
+            server_to_client_receive,
+            client_to_server_send,
+            experimental_task_handlers=task_handlers,
+        ) as session,
+        anyio.create_task_group() as tg,
+        client_to_server_send,
+        client_to_server_receive,
+        server_to_client_send,
+        server_to_client_receive,
+    ):
+        tg.start_soon(mock_server)
+        await session.initialize()
+
+    # Assert that tasks capability includes requests.sampling
+    assert received_capabilities is not None
+    assert received_capabilities.tasks is not None
+    assert received_capabilities.tasks.requests is not None
+    assert received_capabilities.tasks.requests.sampling is not None
+    assert received_capabilities.tasks.requests.elicitation is None  # Not provided

tests/experimental/tasks/client/test_handlers.py

@@ -686,3 +686,201 @@ async def run_client() -> None:
             )
 
             tg.cancel_scope.cancel()
+
+
+@pytest.mark.anyio
+async def test_client_returns_error_for_unhandled_task_result_request(client_streams: ClientTestStreams) -> None:
+    """Test that client returns error for unhandled tasks/result request."""
+    with anyio.fail_after(10):
+        client_ready = anyio.Event()
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_client() -> None:
+                async with ClientSession(
+                    client_streams.client_receive,
+                    client_streams.client_send,
+                    message_handler=_default_message_handler,
+                ):
+                    client_ready.set()
+                    await anyio.sleep_forever()
+
+            tg.start_soon(run_client)
+            await client_ready.wait()
+
+            typed_request = GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId="nonexistent"))
+            request = types.JSONRPCRequest(
+                jsonrpc="2.0",
+                id="req-result",
+                **typed_request.model_dump(by_alias=True),
+            )
+            await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
+
+            response_msg = await client_streams.server_receive.receive()
+            response = response_msg.message.root
+            assert isinstance(response, types.JSONRPCError)
+            assert "not supported" in response.error.message.lower()
+
+            tg.cancel_scope.cancel()
+
+
+@pytest.mark.anyio
+async def test_client_returns_error_for_unhandled_list_tasks_request(client_streams: ClientTestStreams) -> None:
+    """Test that client returns error for unhandled tasks/list request."""
+    with anyio.fail_after(10):
+        client_ready = anyio.Event()
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_client() -> None:
+                async with ClientSession(
+                    client_streams.client_receive,
+                    client_streams.client_send,
+                    message_handler=_default_message_handler,
+                ):
+                    client_ready.set()
+                    await anyio.sleep_forever()
+
+            tg.start_soon(run_client)
+            await client_ready.wait()
+
+            typed_request = ListTasksRequest()
+            request = types.JSONRPCRequest(
+                jsonrpc="2.0",
+                id="req-list",
+                **typed_request.model_dump(by_alias=True),
+            )
+            await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
+
+            response_msg = await client_streams.server_receive.receive()
+            response = response_msg.message.root
+            assert isinstance(response, types.JSONRPCError)
+            assert "not supported" in response.error.message.lower()
+
+            tg.cancel_scope.cancel()
+
+
+@pytest.mark.anyio
+async def test_client_returns_error_for_unhandled_cancel_task_request(client_streams: ClientTestStreams) -> None:
+    """Test that client returns error for unhandled tasks/cancel request."""
+    with anyio.fail_after(10):
+        client_ready = anyio.Event()
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_client() -> None:
+                async with ClientSession(
+                    client_streams.client_receive,
+                    client_streams.client_send,
+                    message_handler=_default_message_handler,
+                ):
+                    client_ready.set()
+                    await anyio.sleep_forever()
+
+            tg.start_soon(run_client)
+            await client_ready.wait()
+
+            typed_request = CancelTaskRequest(params=CancelTaskRequestParams(taskId="nonexistent"))
+            request = types.JSONRPCRequest(
+                jsonrpc="2.0",
+                id="req-cancel",
+                **typed_request.model_dump(by_alias=True),
+            )
+            await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
+
+            response_msg = await client_streams.server_receive.receive()
+            response = response_msg.message.root
+            assert isinstance(response, types.JSONRPCError)
+            assert "not supported" in response.error.message.lower()
+
+            tg.cancel_scope.cancel()
+
+
+@pytest.mark.anyio
+async def test_client_returns_error_for_unhandled_task_augmented_sampling(client_streams: ClientTestStreams) -> None:
+    """Test that client returns error for task-augmented sampling without handler."""
+    with anyio.fail_after(10):
+        client_ready = anyio.Event()
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_client() -> None:
+                # No task handlers provided - uses defaults
+                async with ClientSession(
+                    client_streams.client_receive,
+                    client_streams.client_send,
+                    message_handler=_default_message_handler,
+                ):
+                    client_ready.set()
+                    await anyio.sleep_forever()
+
+            tg.start_soon(run_client)
+            await client_ready.wait()
+
+            # Send task-augmented sampling request
+            typed_request = CreateMessageRequest(
+                params=CreateMessageRequestParams(
+                    messages=[SamplingMessage(role="user", content=TextContent(type="text", text="Hello"))],
+                    maxTokens=100,
+                    task=TaskMetadata(ttl=60000),
+                )
+            )
+            request = types.JSONRPCRequest(
+                jsonrpc="2.0",
+                id="req-sampling",
+                **typed_request.model_dump(by_alias=True),
+            )
+            await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
+
+            response_msg = await client_streams.server_receive.receive()
+            response = response_msg.message.root
+            assert isinstance(response, types.JSONRPCError)
+            assert "not supported" in response.error.message.lower()
+
+            tg.cancel_scope.cancel()
+
+
+@pytest.mark.anyio
+async def test_client_returns_error_for_unhandled_task_augmented_elicitation(
+    client_streams: ClientTestStreams,
+) -> None:
+    """Test that client returns error for task-augmented elicitation without handler."""
+    with anyio.fail_after(10):
+        client_ready = anyio.Event()
+
+        async with anyio.create_task_group() as tg:
+
+            async def run_client() -> None:
+                # No task handlers provided - uses defaults
+                async with ClientSession(
+                    client_streams.client_receive,
+                    client_streams.client_send,
+                    message_handler=_default_message_handler,
+                ):
+                    client_ready.set()
+                    await anyio.sleep_forever()
+
+            tg.start_soon(run_client)
+            await client_ready.wait()
+
+            # Send task-augmented elicitation request
+            typed_request = ElicitRequest(
+                params=ElicitRequestFormParams(
+                    message="What is your name?",
+                    requestedSchema={"type": "object", "properties": {"name": {"type": "string"}}},
+                    task=TaskMetadata(ttl=60000),
+                )
+            )
+            request = types.JSONRPCRequest(
+                jsonrpc="2.0",
+                id="req-elicit",
+                **typed_request.model_dump(by_alias=True),
+            )
+            await client_streams.server_send.send(SessionMessage(types.JSONRPCMessage(request)))
+
+            response_msg = await client_streams.server_receive.receive()
+            response = response_msg.message.root
+            assert isinstance(response, types.JSONRPCError)
+            assert "not supported" in response.error.message.lower()
+
+            tg.cancel_scope.cancel()