Clean up task code: fix error codes, imports, type annotations, and add test fixtures

Source code:
- Use INVALID_REQUEST constant instead of hardcoded -32600
- Move inline imports to module level in task_context.py
- Use cast() instead of type: ignore in resolver.py

Test improvements:
- Add ClientTestStreams dataclass and client_streams fixture
- Add store fixture for test_store.py with automatic cleanup
- Remove unused ClientTaskContext dataclass
- Add STREAM_BUFFER_SIZE constant for magic number

Author: maxisbey

src/mcp/server/experimental/task_context.py

@@ -15,10 +15,12 @@
 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
 from mcp.types import (
+    INVALID_REQUEST,
     TASK_STATUS_INPUT_REQUIRED,
     TASK_STATUS_WORKING,
     ClientCapabilities,
@@ -35,6 +37,7 @@
     SamplingMessage,
     ServerNotification,
     Task,
+    TaskMetadata,
     TaskStatusNotification,
     TaskStatusNotificationParams,
 )
@@ -90,14 +93,9 @@ def __init__(
         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, we need to get the task from the store synchronously
-        # This is a limitation - for async task lookup, use task= parameter
+        # If task_id provided, create a minimal task object
+        # This is for backwards compatibility with tests that pass task_id
         if task is None:
-            # Create a minimal task object - the real task state comes from the store
-            # This is for backwards compatibility with tests that pass task_id
-            from mcp.shared.experimental.tasks.helpers import create_task_state
-            from mcp.types import TaskMetadata
-
             task = create_task_state(TaskMetadata(ttl=None), task_id=task_id)
 
         self._ctx = TaskContext(task=task, store=store)
@@ -191,7 +189,7 @@ def _check_elicitation_capability(self) -> None:
         if not self._session.check_client_capability(ClientCapabilities(elicitation=ElicitationCapability())):
             raise McpError(
                 ErrorData(
-                    code=-32600,  # INVALID_REQUEST
+                    code=INVALID_REQUEST,
                     message="Client does not support elicitation capability",
                 )
             )
@@ -201,7 +199,7 @@ def _check_sampling_capability(self) -> None:
         if not self._session.check_client_capability(ClientCapabilities(sampling=SamplingCapability())):
             raise McpError(
                 ErrorData(
-                    code=-32600,  # INVALID_REQUEST
+                    code=INVALID_REQUEST,
                     message="Client does not support sampling capability",
                 )
             )

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

@@ -5,7 +5,7 @@
 to another without depending on asyncio.Future.
 """
 
-from typing import Generic, TypeVar
+from typing import Generic, TypeVar, cast
 
 import anyio
 
@@ -52,7 +52,8 @@ async def wait(self) -> T:
         await self._event.wait()
         if self._exception is not None:
             raise self._exception
-        return self._value  # type: ignore[return-value]
+        # If we reach here, set_result() was called, so _value is set
+        return cast(T, self._value)
 
     def done(self) -> bool:
         """Return True if the resolver has been completed."""

tests/experimental/tasks/client/test_handlers.py

@@ -10,12 +10,14 @@
 client -> server task requests.
 """
 
-from dataclasses import dataclass, field
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
 
 import anyio
 import pytest
 from anyio import Event
 from anyio.abc import TaskGroup
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
 import mcp.types as types
 from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
@@ -43,14 +45,47 @@
     TextContent,
 )
 
+# Buffer size for test streams
+STREAM_BUFFER_SIZE = 10
+
 
 @dataclass
-class ClientTaskContext:
-    """Context for managing client-side tasks during tests."""
+class ClientTestStreams:
+    """Bidirectional message streams for client/server communication in tests."""
+
+    server_send: MemoryObjectSendStream[SessionMessage]
+    server_receive: MemoryObjectReceiveStream[SessionMessage]
+    client_send: MemoryObjectSendStream[SessionMessage]
+    client_receive: MemoryObjectReceiveStream[SessionMessage]
+
 
-    task_group: TaskGroup
-    store: InMemoryTaskStore
-    task_done_events: dict[str, Event] = field(default_factory=lambda: {})
+@pytest.fixture
+async def client_streams() -> AsyncIterator[ClientTestStreams]:
+    """Create bidirectional message streams for client tests.
+
+    Automatically closes all streams after the test completes.
+    """
+    server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](
+        STREAM_BUFFER_SIZE
+    )
+    client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](
+        STREAM_BUFFER_SIZE
+    )
+
+    streams = ClientTestStreams(
+        server_send=server_to_client_send,
+        server_receive=client_to_server_receive,
+        client_send=client_to_server_send,
+        client_receive=server_to_client_receive,
+    )
+
+    yield streams
+
+    # Cleanup
+    await server_to_client_send.aclose()
+    await server_to_client_receive.aclose()
+    await client_to_server_send.aclose()
+    await client_to_server_receive.aclose()
 
 
 @pytest.mark.anyio