Drive task-augmented tool calls transparently in the client

A client that declares the tasks extension now handles CreateTaskResult
everywhere SEP-2663 requires it:

- ClientSession widens its tools/call adapter per session when the
  extension is declared, and call_tool gains allow_create_task overloads
  mirroring allow_input_required (the guarded default raises with
  guidance instead of a pydantic ValidationError).
- Client.call_tool composes a polling driver after the input_required
  driver: a CreateTaskResult is polled via tasks/get (honoring
  pollIntervalMs with a 1s fallback) until terminal, the inlined result
  re-enters output-schema validation, and the public contract stays
  CallToolResult. failed, cancelled, and input_required tasks surface as
  typed TaskFailedError, TaskCancelledError, and TaskInputRequiredError.
- The tasks story's modern path becomes the plain typed call_tool (the
  same line demonstrates legacy degradation), with a compact manual leg
  over the mcp.shared.tasks wrappers.

Author: maxisbey

docs/migration.md

@@ -470,10 +470,18 @@ Two reference extensions ship in their own modules:
   keeps completed tasks in a pluggable `TaskStore` (`Tasks(store=...)`,
   in-memory default) that enforces `default_ttl_ms`. A `tasks/*` call from a
   non-declaring modern client is rejected with `-32021` (missing required
-  client capability); legacy calls get `METHOD_NOT_FOUND`. This is the core
-  SEP-2663 surface; background execution (`working` tasks), the in-task
-  `input_required` loop over `tasks/update`, `notifications/tasks`, and task
-  routing headers are deferred.
+  client capability); legacy calls get `METHOD_NOT_FOUND`. On the client side,
+  a `Client` that declares the extension gets transparent polling:
+  `Client.call_tool` recognises the `CreateTaskResult`, polls `tasks/get`
+  (honoring `pollIntervalMs`), and returns the final `CallToolResult`
+  unchanged, while `failed`/`cancelled` tasks surface as the typed
+  `TaskFailedError`/`TaskCancelledError`. Manual driving stays available —
+  `client.session.call_tool(..., allow_create_task=True)` returns the typed
+  `CreateTaskResult`, and the `mcp.shared.tasks` request wrappers drive
+  `tasks/get`/`tasks/update`/`tasks/cancel` over `session.send_request`. This
+  is the core SEP-2663 surface; background execution (`working` tasks), the
+  in-task `input_required` loop over `tasks/update`, `notifications/tasks`,
+  and task routing headers are deferred.
 
 Extension methods are strictly additive: a `MethodBinding` cannot name a
 spec-defined request method, and registering one whose method collides with

examples/stories/tasks/README.md

@@ -3,18 +3,20 @@
 Task-augmented execution (SEP-2663). A client declares the
 `io.modelcontextprotocol/tasks` extension; the server may then answer a
 `tools/call` with a `CreateTaskResult` (carrying a task id) instead of the
-`CallToolResult`, and the client fetches the result via `tasks/get`.
+`CallToolResult`. `Client.call_tool` drives the polling transparently and
+surfaces only the final result — the SEP's recommended client shape.
 
 ## Run it
 
 ```bash
 # stdio (default) — today's stdio negotiates the legacy wire, which cannot carry
 # the extension capability, so this leg demonstrates graceful degradation: the
-# same tools/call returns a plain CallToolResult, never a task.
+# same call_tool returns a plain CallToolResult, never a task.
 uv run python -m stories.tasks.client
 
 # HTTP — the modern wire negotiates the extension; the server defers the call as
-# a task and the client reads the result back via tasks/get
+# a task, Client.call_tool polls it to completion, and a manual leg shows the
+# raw CreateTaskResult -> tasks/get wire flow
 uv run python -m stories.tasks.client --http
 ```
 
@@ -29,9 +31,14 @@ uv run python -m stories.tasks.client --http
   declared the extension on the request, returning a flat `CreateTaskResult`
   (`resultType: "task"`).
 - `client.py` `Client(target, extensions={EXTENSION_ID: {}})` — declaring the
-  extension is what lets the server defer; `main` then reads the `CreateTaskResult`
-  and fetches `tasks/get`, whose completed envelope inlines the original
-  `CallToolResult`.
+  extension is what lets the server defer. The transparent path is then just
+  `await client.call_tool(...)`: `Client` recognises the `CreateTaskResult`,
+  polls `tasks/get` (honoring `pollIntervalMs`), and returns the final
+  `CallToolResult`; a `failed` task raises `TaskFailedError`.
+- The manual leg — `session.call_tool(..., allow_create_task=True)` returns the
+  typed `CreateTaskResult` (mirroring `allow_input_required`), and the shared
+  `mcp.shared.tasks` wrappers (`GetTaskRequest`/`GetTaskResult`) drive `tasks/get`
+  by hand over `session.send_request`.
 
 ## Scope
 

examples/stories/tasks/client.py

@@ -1,59 +1,58 @@
-"""Declare the tasks extension, let the server defer a tool call, then fetch the result via tasks/get.
+"""Declare the tasks extension and let `Client.call_tool` drive the task transparently.
 
 The client declares `io.modelcontextprotocol/tasks` (via `Client(extensions=...)`),
-so the server is free to answer `tools/call` with a `CreateTaskResult`. `Client`
-exposes only spec verbs, so the augmented call and `tasks/get` drop to
-`client.session`; the thin `_send` helper keeps that out of the story below.
+so the server is free to answer `tools/call` with a `CreateTaskResult`. SEP-2663
+advises clients to keep a fixed public contract and drive the polling internally —
+`Client.call_tool` does exactly that, so the modern path is the same typed call a
+task-less server would get. A compact manual leg then shows the raw wire flow:
+`session.call_tool(allow_create_task=True)` for the typed `CreateTaskResult`, and
+the shared `mcp.shared.tasks` wrappers over `session.send_request` for `tasks/get`.
 """
 
-from typing import Any, Literal, cast
+from typing import cast
 
 import mcp_types as types
-from pydantic import TypeAdapter
 
-from mcp.client import Client, ClientSession
-from mcp.server.tasks import EXTENSION_ID, GetTaskRequestParams
+from mcp.client import Client
+from mcp.server.tasks import EXTENSION_ID
+from mcp.shared.tasks import CreateTaskResult, GetTaskRequest, GetTaskRequestParams, GetTaskResult
 from stories._harness import Target, run_client
 
-_RAW: TypeAdapter[dict[str, Any]] = TypeAdapter(dict)
-
-
-class _GetTaskRequest(types.Request[GetTaskRequestParams, Literal["tasks/get"]]):
-    method: Literal["tasks/get"] = "tasks/get"
-    params: GetTaskRequestParams
-
-
-async def _send(session: ClientSession, request: types.Request[Any, Any]) -> dict[str, Any]:
-    """Send a request whose result has a non-spec (extension) shape; return the raw dict."""
-    return await session.send_request(cast("types.ClientRequest", request), cast("Any", _RAW))
-
 
 async def main(target: Target, *, mode: str = "auto") -> None:
     async with Client(target, mode=mode, extensions={EXTENSION_ID: {}}) as client:
-        # The extension is a modern-only capability negotiated over server/discover.
-        # A legacy connection (today's stdio) cannot carry it, and the server then
-        # must not augment: the same tools/call degrades to a plain CallToolResult.
+        # The transparent path. On the modern wire the server augments this
+        # tools/call into a task (we declared the extension) and Client.call_tool
+        # polls tasks/get to the final result; on a legacy connection (today's
+        # stdio) the extension cannot be negotiated, the server must not augment,
+        # and the very same call simply returns the plain CallToolResult.
+        result = await client.call_tool("render_report", {"title": "Q3", "sections": 2})
+        assert isinstance(result.content[0], types.TextContent), result
+        assert result.content[0].text.startswith("# Q3"), result
+        # No 2025-style related-task _meta either; the task plumbing never leaks
+        # into the surfaced result.
+        assert result.meta is None, result
+
         if client.server_capabilities.extensions is None:
-            result = await client.call_tool("render_report", {"title": "Q3", "sections": 2})
-            assert isinstance(result.content[0], types.TextContent), result
-            assert result.content[0].text.startswith("# Q3"), result
-            # No 2025-style related-task _meta either; SEP-2663 augmentation would
-            # have replaced the whole result, failing CallToolResult parsing above.
-            assert result.meta is None, result
+            # Legacy wire: nothing more to show — the degradation above is the point.
             return
         assert client.server_capabilities.extensions == {EXTENSION_ID: {}}
 
-        # The server augments this tools/call into a task because we declared the extension.
-        call = types.CallToolRequest(
-            params=types.CallToolRequestParams(name="render_report", arguments={"title": "Q3", "sections": 2})
+        # The manual leg: the same flow driven by hand on the raw wire.
+        # allow_create_task=True hands back the typed CreateTaskResult instead of
+        # polling, and the shared SEP-2663 request wrappers fetch the outcome.
+        created = await client.session.call_tool(
+            "render_report", {"title": "Q3", "sections": 1}, allow_create_task=True
         )
-        created = await _send(client.session, call)
-        assert created["resultType"] == "task", created
-        task_id = created["taskId"]
+        assert isinstance(created, CreateTaskResult), created
 
-        task = await _send(client.session, _GetTaskRequest(params=GetTaskRequestParams(task_id=task_id)))
-        assert task["status"] == "completed", task
-        assert task["result"]["content"][0]["text"].startswith("# Q3"), task
+        task = await client.session.send_request(
+            cast("types.ClientRequest", GetTaskRequest(params=GetTaskRequestParams(task_id=created.task_id))),
+            GetTaskResult,
+        )
+        assert task.status == "completed", task
+        assert task.result is not None, task
+        assert task.result["content"][0]["text"].startswith("# Q3"), task
 
 
 if __name__ == "__main__":

src/mcp/__init__.py

@@ -59,6 +59,7 @@
 from mcp_types import Role as SamplingRole
 
 from .client._input_required import InputRequiredRoundsExceededError
+from .client._tasks import TaskCancelledError, TaskFailedError, TaskInputRequiredError
 from .client.client import Client
 from .client.session import ClientSession
 from .client.session_group import ClientSessionGroup
@@ -128,6 +129,9 @@
     "StdioServerParameters",
     "StopReason",
     "SubscribeRequest",
+    "TaskCancelledError",
+    "TaskFailedError",
+    "TaskInputRequiredError",
     "Tool",
     "ToolChoice",
     "ToolResultContent",

src/mcp/client/__init__.py

@@ -1,6 +1,7 @@
 """MCP Client module."""
 
 from mcp.client._input_required import InputRequiredRoundsExceededError
+from mcp.client._tasks import TaskCancelledError, TaskFailedError, TaskInputRequiredError
 from mcp.client._transport import Transport
 from mcp.client.caching import (
     CacheConfig,
@@ -25,5 +26,8 @@
     "InMemoryResponseCacheStore",
     "InputRequiredRoundsExceededError",
     "ResponseCacheStore",
+    "TaskCancelledError",
+    "TaskFailedError",
+    "TaskInputRequiredError",
     "Transport",
 ]

src/mcp/client/_tasks.py

@@ -0,0 +1,122 @@
+"""SEP-2663 client-side task polling driver.
+
+When a server augments a `tools/call` into a task — a `CreateTaskResult` in
+place of the `CallToolResult` — the client polls `tasks/get` until the task
+reaches a terminal status and surfaces only the final result. SEP-2663 advises
+exactly this shape: "existing code returning a fixed shape ... can transparently
+drive the polling flow internally and surface only the final, completed result".
+This module implements that loop as a pure function so it stays testable with
+plain closures; `Client` builds the `get_task` closure over its session,
+`ClientSession` stays mechanics-only (mirroring `_input_required`).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+
+import anyio
+from mcp_types import CallToolResult, ErrorData
+
+from mcp.shared.exceptions import MCPError
+from mcp.shared.tasks import CreateTaskResult, GetTaskResult
+
+DEFAULT_POLL_INTERVAL_SECONDS = 1.0
+"""Poll cadence when neither the snapshot nor the `CreateTaskResult` carries `pollIntervalMs`.
+
+SEP-2663 makes the hint optional and only says clients SHOULD honor it when
+present; one second is the SDK's conservative default in its absence.
+"""
+
+
+class TaskFailedError(MCPError):
+    """The task reached `failed`: a JSON-RPC error occurred during execution (SEP-2663).
+
+    Carries the JSON-RPC error inlined on `tasks/get` as `code`/`message`/`data`,
+    plus the snapshot's optional `statusMessage` diagnostic.
+    """
+
+    def __init__(self, error: ErrorData, status_message: str | None = None) -> None:
+        super().__init__(code=error.code, message=error.message, data=error.data)
+        self.status_message = status_message
+
+
+class TaskCancelledError(RuntimeError):
+    """The task reached `cancelled` before producing a result (SEP-2663)."""
+
+    def __init__(self, task_id: str, status_message: str | None = None) -> None:
+        detail = f": {status_message}" if status_message is not None else ""
+        super().__init__(f"Task {task_id!r} was cancelled{detail}")
+        self.task_id = task_id
+        self.status_message = status_message
+
+
+class TaskInputRequiredError(RuntimeError):
+    """The task reached `input_required`, which this driver does not drive yet.
+
+    SEP-2663's in-task input loop (fulfil `inputRequests` via `tasks/update`) is
+    a deferred follow-up in this SDK. Drive it manually: poll with
+    `mcp.shared.tasks.GetTaskRequest` and answer with
+    `mcp.shared.tasks.UpdateTaskRequest` over `session.send_request`.
+    """
+
+    def __init__(self, task_id: str) -> None:
+        super().__init__(
+            f"Task {task_id!r} requires in-task input (status `input_required`); the SDK's automatic "
+            "in-task input loop is not implemented yet. Drive it manually with the `mcp.shared.tasks` "
+            "request wrappers (`GetTaskRequest`/`UpdateTaskRequest`) over `session.send_request`."
+        )
+        self.task_id = task_id
+
+
+async def run_task_driver(
+    created: CreateTaskResult,
+    *,
+    get_task: Callable[[str], Awaitable[GetTaskResult]],
+    sleep: Callable[[float], Awaitable[None]] = anyio.sleep,
+) -> CallToolResult:
+    """Poll a `CreateTaskResult` to its final `CallToolResult`.
+
+    Polls `tasks/get` (via `get_task`) until the task reaches a terminal status.
+    Between polls it honors the SEP-2663 `pollIntervalMs` hint: each non-terminal
+    snapshot sleeps its own `poll_interval_ms`, falling back to the
+    `CreateTaskResult`'s, then to `DEFAULT_POLL_INTERVAL_SECONDS`.
+
+    The loop deliberately imposes no round cap or deadline of its own: SEP-2663
+    tasks represent unbounded server-side work, so how long to wait is the
+    caller's policy — cancel via an enclosing anyio cancel scope, or bound each
+    `tasks/get` round with the session read timeout the `get_task` closure
+    carries.
+
+    Args:
+        created: The `CreateTaskResult` the augmented request returned.
+        get_task: Sends one `tasks/get` for the given task id and returns the
+            parsed `GetTaskResult` snapshot.
+        sleep: Awaits the given number of seconds between polls (injectable for
+            deterministic tests).
+
+    Raises:
+        TaskFailedError: The task reached `failed`; carries the inlined JSON-RPC error.
+        TaskCancelledError: The task reached `cancelled`.
+        TaskInputRequiredError: The task reached `input_required` (the in-task
+            input loop is not implemented yet).
+        RuntimeError: The server violated SEP-2663 — a `completed` snapshot
+            without `result`, or a `failed` snapshot without `error`.
+    """
+    while True:
+        snapshot = await get_task(created.task_id)
+        if snapshot.status == "completed":
+            if snapshot.result is None:
+                raise RuntimeError(
+                    f"Task {created.task_id!r} is `completed` but carries no `result` (SEP-2663 violation)"
+                )
+            return CallToolResult.model_validate(snapshot.result, by_name=False)
+        if snapshot.status == "failed":
+            if snapshot.error is None:
+                raise RuntimeError(f"Task {created.task_id!r} is `failed` but carries no `error` (SEP-2663 violation)")
+            raise TaskFailedError(ErrorData.model_validate(snapshot.error), snapshot.status_message)
+        if snapshot.status == "cancelled":
+            raise TaskCancelledError(created.task_id, snapshot.status_message)
+        if snapshot.status == "input_required":
+            raise TaskInputRequiredError(created.task_id)
+        interval_ms = snapshot.poll_interval_ms if snapshot.poll_interval_ms is not None else created.poll_interval_ms
+        await sleep(DEFAULT_POLL_INTERVAL_SECONDS if interval_ms is None else interval_ms / 1000)