SEP-1686: Tasks (#1645)

src/mcp/client/session.py

@@ -8,6 +8,8 @@ from pydantic import AnyUrl, TypeAdapter
 from typing_extensions import deprecated
 
 import mcp.types as types
+from mcp.client.experimental import ExperimentalClientFeatures
+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
@@ -118,6 +120,8 @@ class ClientSession(
         logging_callback: LoggingFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         client_info: types.Implementation | None = None,
+        *,
+        experimental_task_handlers: ExperimentalTaskHandlers | None = None,
     ) -> None:
         super().__init__(
             read_stream,
@@ -134,6 +138,10 @@ class ClientSession(
         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_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
@@ -164,6 +172,7 @@ class ClientSession(
                             elicitation=elicitation,
                             experimental=None,
                             roots=roots,
+                            tasks=self._task_handlers.build_capability(),
                         ),
                         clientInfo=self._client_info,
                     ),
@@ -188,6 +197,20 @@ class ClientSession(
         """
         return self._server_capabilities
 
+    @property
+    def experimental(self) -> ExperimentalClientFeatures:
+        """Experimental APIs for tasks and other features.
+
+        WARNING: These APIs are experimental and may change without notice.
+
+        Example:
+            status = await session.experimental.get_task(task_id)
+            result = await session.experimental.get_task_result(task_id, CallToolResult)
+        """
+        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."""
         return await self.send_request(
@@ -521,16 +544,31 @@ class ClientSession(
             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:
-                    response = await self._sampling_callback(ctx, params)
+                    # Check if this is a task-augmented request
+                    if params.task is not None:
+                        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)
                     await responder.respond(client_response)
 
             case types.ElicitRequest(params=params):
                 with responder:
-                    response = await self._elicitation_callback(ctx, params)
+                    # Check if this is a task-augmented request
+                    if params.task is not None:
+                        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)
                     await responder.respond(client_response)
 
@@ -544,6 +582,11 @@ class ClientSession(
                 with responder:
                     return await responder.respond(types.ClientResult(root=types.EmptyResult()))
 
+            case _:  # pragma: no cover
+                pass  # Task requests handled above by _task_handlers
+
+        return None
+
     async def _handle_incoming(
         self,
         req: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,