Refactor client task handlers into ExperimentalTaskHandlers dataclass

- Replace 6 individual task handler parameters with single
  `experimental_task_handlers: ExperimentalTaskHandlers` (keyword-only)
- ExperimentalTaskHandlers dataclass groups all handlers and provides:
  - `build_capability()` - auto-builds ClientTasksCapability from handlers
  - `handles_request()` - checks if request is task-related
  - `handle_request()` - dispatches to appropriate handler
- Simplify ClientSession._received_request by delegating task requests
- Update tests to use new ExperimentalTaskHandlers API

Author: maxisbey

src/mcp/client/experimental/task_handlers.py

@@ -12,10 +12,12 @@
 - Server polls client's task status via tasks/get, tasks/result, etc.
 """
 
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Protocol
 
 import mcp.types as types
 from mcp.shared.context import RequestContext
+from mcp.shared.session import RequestResponder
 
 if TYPE_CHECKING:
     from mcp.client.session import ClientSession
@@ -109,7 +111,11 @@ async def __call__(
     ) -> types.CreateTaskResult | types.ErrorData: ...  # pragma: no branch
 
 
-# Default handlers for experimental task requests (return "not supported" errors)
+# =============================================================================
+# Default Handlers (return "not supported" errors)
+# =============================================================================
+
+
 async def default_get_task_handler(
     context: RequestContext["ClientSession", Any],
     params: types.GetTaskRequestParams,
@@ -150,7 +156,7 @@ async def default_cancel_task_handler(
     )
 
 
-async def default_task_augmented_sampling_callback(
+async def default_task_augmented_sampling(
     context: RequestContext["ClientSession", Any],
     params: types.CreateMessageRequestParams,
     task_metadata: types.TaskMetadata,
@@ -161,7 +167,7 @@ async def default_task_augmented_sampling_callback(
     )
 
 
-async def default_task_augmented_elicitation_callback(
+async def default_task_augmented_elicitation(
     context: RequestContext["ClientSession", Any],
     params: types.ElicitRequestParams,
     task_metadata: types.TaskMetadata,
@@ -172,58 +178,118 @@ async def default_task_augmented_elicitation_callback(
     )
 
 
-def build_client_tasks_capability(
-    *,
-    list_tasks_handler: ListTasksHandlerFnT | None = None,
-    cancel_task_handler: CancelTaskHandlerFnT | None = None,
-    task_augmented_sampling_callback: TaskAugmentedSamplingFnT | None = None,
-    task_augmented_elicitation_callback: TaskAugmentedElicitationFnT | None = None,
-) -> types.ClientTasksCapability | None:
-    """Build ClientTasksCapability from the provided handlers.
-
-    This helper builds the appropriate capability object based on which
-    handlers are provided (non-None and not the default handlers).
+@dataclass
+class ExperimentalTaskHandlers:
+    """Container for experimental task handlers.
 
-    WARNING: This is experimental and may change without notice.
+    Groups all task-related handlers that handle server -> client requests.
+    This includes both pure task requests (get, list, cancel, result) and
+    task-augmented request handlers (sampling, elicitation with task field).
 
-    Args:
-        list_tasks_handler: Handler for tasks/list requests
-        cancel_task_handler: Handler for tasks/cancel requests
-        task_augmented_sampling_callback: Handler for task-augmented sampling
-        task_augmented_elicitation_callback: Handler for task-augmented elicitation
+    WARNING: These APIs are experimental and may change without notice.
 
-    Returns:
-        ClientTasksCapability if any handlers are provided, None otherwise
+    Example:
+        handlers = ExperimentalTaskHandlers(
+            get_task=my_get_task_handler,
+            list_tasks=my_list_tasks_handler,
+        )
+        session = ClientSession(..., experimental_task_handlers=handlers)
     """
-    has_list = list_tasks_handler is not None and list_tasks_handler is not default_list_tasks_handler
-    has_cancel = cancel_task_handler is not None and cancel_task_handler is not default_cancel_task_handler
-    has_sampling = (
-        task_augmented_sampling_callback is not None
-        and task_augmented_sampling_callback is not default_task_augmented_sampling_callback
-    )
-    has_elicitation = (
-        task_augmented_elicitation_callback is not None
-        and task_augmented_elicitation_callback is not default_task_augmented_elicitation_callback
-    )
 
-    # If no handlers are provided, return None
-    if not any([has_list, has_cancel, has_sampling, has_elicitation]):
-        return None
-
-    # Build requests capability if any request handlers are provided
-    requests_capability: types.ClientTasksRequestsCapability | None = None
-    if has_sampling or has_elicitation:
-        requests_capability = types.ClientTasksRequestsCapability(
-            sampling=types.TasksSamplingCapability(createMessage=types.TasksCreateMessageCapability())
-            if has_sampling
-            else None,
-            elicitation=types.TasksElicitationCapability(create=types.TasksCreateElicitationCapability())
-            if has_elicitation
-            else None,
+    # Pure task request handlers
+    get_task: GetTaskHandlerFnT = field(default=default_get_task_handler)
+    get_task_result: GetTaskResultHandlerFnT = field(default=default_get_task_result_handler)
+    list_tasks: ListTasksHandlerFnT = field(default=default_list_tasks_handler)
+    cancel_task: CancelTaskHandlerFnT = field(default=default_cancel_task_handler)
+
+    # Task-augmented request handlers
+    augmented_sampling: TaskAugmentedSamplingFnT = field(default=default_task_augmented_sampling)
+    augmented_elicitation: TaskAugmentedElicitationFnT = field(default=default_task_augmented_elicitation)
+
+    def build_capability(self) -> types.ClientTasksCapability | None:
+        """Build ClientTasksCapability from the configured handlers.
+
+        Returns a capability object that reflects which handlers are configured
+        (i.e., not using the default "not supported" handlers).
+
+        Returns:
+            ClientTasksCapability if any handlers are provided, None otherwise
+        """
+        has_list = self.list_tasks is not default_list_tasks_handler
+        has_cancel = self.cancel_task is not default_cancel_task_handler
+        has_sampling = self.augmented_sampling is not default_task_augmented_sampling
+        has_elicitation = self.augmented_elicitation is not default_task_augmented_elicitation
+
+        # If no handlers are provided, return None
+        if not any([has_list, has_cancel, has_sampling, has_elicitation]):
+            return None
+
+        # Build requests capability if any request handlers are provided
+        requests_capability: types.ClientTasksRequestsCapability | None = None
+        if has_sampling or has_elicitation:
+            requests_capability = types.ClientTasksRequestsCapability(
+                sampling=types.TasksSamplingCapability(createMessage=types.TasksCreateMessageCapability())
+                if has_sampling
+                else None,
+                elicitation=types.TasksElicitationCapability(create=types.TasksCreateElicitationCapability())
+                if has_elicitation
+                else None,
+            )
+
+        return types.ClientTasksCapability(
+            list=types.TasksListCapability() if has_list else None,
+            cancel=types.TasksCancelCapability() if has_cancel else None,
+            requests=requests_capability,
         )
 
-    return types.ClientTasksCapability(
-        list=types.TasksListCapability() if has_list else None,
-        cancel=types.TasksCancelCapability() if has_cancel else None,
-        requests=requests_capability,
-    )
+    @staticmethod
+    def handles_request(request: types.ServerRequest) -> bool:
+        """Check if this handler handles the given request type."""
+        return isinstance(
+            request.root,
+            types.GetTaskRequest | types.GetTaskPayloadRequest | types.ListTasksRequest | types.CancelTaskRequest,
+        )
+
+    async def handle_request(
+        self,
+        ctx: RequestContext["ClientSession", Any],
+        responder: RequestResponder[types.ServerRequest, types.ClientResult],
+    ) -> None:
+        """Handle a task-related request from the server.
+
+        Call handles_request() first to check if this handler can handle the request.
+        """
+        from pydantic import TypeAdapter
+
+        client_response_type: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
+            types.ClientResult | types.ErrorData
+        )
+
+        match responder.request.root:
+            case types.GetTaskRequest(params=params):
+                response = await self.get_task(ctx, params)
+                client_response = client_response_type.validate_python(response)
+                await responder.respond(client_response)
+
+            case types.GetTaskPayloadRequest(params=params):
+                response = await self.get_task_result(ctx, params)
+                client_response = client_response_type.validate_python(response)
+                await responder.respond(client_response)
+
+            case types.ListTasksRequest(params=params):
+                response = await self.list_tasks(ctx, params)
+                client_response = client_response_type.validate_python(response)
+                await responder.respond(client_response)
+
+            case types.CancelTaskRequest(params=params):
+                response = await self.cancel_task(ctx, params)
+                client_response = client_response_type.validate_python(response)
+                await responder.respond(client_response)
+
+            case _:  # pragma: no cover
+                raise ValueError(f"Unhandled request type: {type(responder.request.root)}")
+
+
+# Backwards compatibility aliases
+default_task_augmented_sampling_callback = default_task_augmented_sampling
+default_task_augmented_elicitation_callback = default_task_augmented_elicitation

src/mcp/client/session.py

@@ -9,21 +9,7 @@
 
 import mcp.types as types
 from mcp.client.experimental import ExperimentalClientFeatures
-from mcp.client.experimental.task_handlers import (
-    CancelTaskHandlerFnT,
-    GetTaskHandlerFnT,
-    GetTaskResultHandlerFnT,
-    ListTasksHandlerFnT,
-    TaskAugmentedElicitationFnT,
-    TaskAugmentedSamplingFnT,
-    build_client_tasks_capability,
-    default_cancel_task_handler,
-    default_get_task_handler,
-    default_get_task_result_handler,
-    default_list_tasks_handler,
-    default_task_augmented_elicitation_callback,
-    default_task_augmented_sampling_callback,
-)
+from mcp.client.experimental.task_handlers import ExperimentalTaskHandlers
 from mcp.shared.context import RequestContext
 from mcp.shared.message import SessionMessage
 from mcp.shared.session import BaseSession, ProgressFnT, RequestResponder
@@ -134,14 +120,8 @@ def __init__(
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
-        tasks_capability: types.ClientTasksCapability | None = None,
-        # Experimental: Task handlers for server -> client requests
-        get_task_handler: GetTaskHandlerFnT | None = None,
-        get_task_result_handler: GetTaskResultHandlerFnT | None = None,
-        list_tasks_handler: ListTasksHandlerFnT | None = None,
-        cancel_task_handler: CancelTaskHandlerFnT | None = None,
-        task_augmented_sampling_callback: TaskAugmentedSamplingFnT | None = None,
-        task_augmented_elicitation_callback: TaskAugmentedElicitationFnT | None = None,
+        *,
+        experimental_task_handlers: ExperimentalTaskHandlers | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -158,25 +138,10 @@ def __init__(
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
         self._server_capabilities: types.ServerCapabilities | None = None
-        self._experimental: ExperimentalClientFeatures | None = None
-        # Experimental: Task handlers
-        self._get_task_handler = get_task_handler or default_get_task_handler
-        self._get_task_result_handler = get_task_result_handler or default_get_task_result_handler
-        self._list_tasks_handler = list_tasks_handler or default_list_tasks_handler
-        self._cancel_task_handler = cancel_task_handler or default_cancel_task_handler
-        self._task_augmented_sampling_callback = (
-            task_augmented_sampling_callback or default_task_augmented_sampling_callback
-        )
-        self._task_augmented_elicitation_callback = (
-            task_augmented_elicitation_callback or default_task_augmented_elicitation_callback
-        )
-        # Build tasks capability from handlers if not explicitly provided
-        self._tasks_capability = tasks_capability or build_client_tasks_capability(
-            list_tasks_handler=list_tasks_handler,
-            cancel_task_handler=cancel_task_handler,
-            task_augmented_sampling_callback=task_augmented_sampling_callback,
-            task_augmented_elicitation_callback=task_augmented_elicitation_callback,
-        )
+        self._experimental_features: ExperimentalClientFeatures | None = None
+
+        # Experimental: Task handlers (use defaults if not provided)
+        self._task_handlers = experimental_task_handlers or ExperimentalTaskHandlers()
 
     async def initialize(self) -> types.InitializeResult:
         sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
@@ -207,7 +172,7 @@ async def initialize(self) -> types.InitializeResult:
                             elicitation=elicitation,
                             experimental=None,
                             roots=roots,
-                            tasks=self._tasks_capability,
+                            tasks=self._task_handlers.build_capability(),
                         ),
                         clientInfo=self._client_info,
                     ),
@@ -242,9 +207,9 @@ def experimental(self) -> ExperimentalClientFeatures:
             status = await session.experimental.get_task(task_id)
             result = await session.experimental.get_task_result(task_id, CallToolResult)
         """
-        if self._experimental is None:
-            self._experimental = ExperimentalClientFeatures(self)
-        return self._experimental
+        if self._experimental_features is None:
+            self._experimental_features = ExperimentalClientFeatures(self)
+        return self._experimental_features
 
     async def send_ping(self) -> types.EmptyResult:
         """Send a ping request."""
@@ -579,12 +544,19 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
             lifespan_context=None,
         )
 
+        # Delegate to experimental task handler if applicable
+        if self._task_handlers.handles_request(responder.request):
+            with responder:
+                await self._task_handlers.handle_request(ctx, responder)
+            return None
+
+        # Core request handling
         match responder.request.root:
             case types.CreateMessageRequest(params=params):
                 with responder:
                     # Check if this is a task-augmented request
                     if params.task is not None:
-                        response = await self._task_augmented_sampling_callback(ctx, params, params.task)
+                        response = await self._task_handlers.augmented_sampling(ctx, params, params.task)
                     else:
                         response = await self._sampling_callback(ctx, params)
                     client_response = ClientResponse.validate_python(response)
@@ -594,7 +566,7 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
                 with responder:
                     # Check if this is a task-augmented request
                     if params.task is not None:
-                        response = await self._task_augmented_elicitation_callback(ctx, params, params.task)
+                        response = await self._task_handlers.augmented_elicitation(ctx, params, params.task)
                     else:
                         response = await self._elicitation_callback(ctx, params)
                     client_response = ClientResponse.validate_python(response)
@@ -610,33 +582,9 @@ async def _received_request(self, responder: RequestResponder[types.ServerReques
                 with responder:
                     return await responder.respond(types.ClientResult(root=types.EmptyResult()))
 
-            # Experimental: Task management requests from server
-            case types.GetTaskRequest(params=params):
-                with responder:
-                    response = await self._get_task_handler(ctx, params)
-                    client_response = ClientResponse.validate_python(response)
-                    await responder.respond(client_response)
-
-            case types.GetTaskPayloadRequest(params=params):
-                with responder:
-                    response = await self._get_task_result_handler(ctx, params)
-                    client_response = ClientResponse.validate_python(response)
-                    await responder.respond(client_response)
-
-            case types.ListTasksRequest(params=params):
-                with responder:
-                    response = await self._list_tasks_handler(ctx, params)
-                    client_response = ClientResponse.validate_python(response)
-                    await responder.respond(client_response)
-
-            case types.CancelTaskRequest(params=params):
-                with responder:
-                    response = await self._cancel_task_handler(ctx, params)
-                    client_response = ClientResponse.validate_python(response)
-                    await responder.respond(client_response)
-
             case _:  # pragma: no cover
                 raise NotImplementedError()
+        return None
 
     async def _handle_incoming(
         self,