Add typed client task functions in mcp.client.tasks

The private driver module becomes the public mcp.client.tasks: get_task,
wait_task, update_task, and cancel_task are typed free functions over
ClientSession, so manually driving a task no longer means hand-building
wire requests. wait_task accepts a bare task id (the persist-and-resume
shape -- task ids are bearer capabilities that survive reconnects) or
the CreateTaskResult, which additionally seeds the poll-interval
fallback; the polling loop now exists once, shared with the
TasksExtension claim resolver. update_task and cancel_task hide the
empty acknowledgement and return None. The task errors gain a common
TaskError base (one except arm for any non-completion) and pickle
support, since they are public API. The manual-driving docs, story leg,
and tutorial are rewritten onto the functions.

Author: maxisbey

docs/advanced/tasks.md

@@ -29,7 +29,8 @@ then come back.
   `TaskFailedError` carrying the inlined JSON-RPC error; a `cancelled` one raises
   `TaskCancelledError`; an `input_required` one raises `TaskInputRequiredError` —
   the automatic in-task input loop is not implemented yet, so drive that task
-  manually (below).
+  manually (below). All three subclass `TaskError`, so one `except TaskError`
+  catches any non-completion.
 
 Degradation is built in. A modern client that does not declare the extension is
 never augmented: it keeps getting plain `CallToolResult`s. And a legacy
@@ -105,20 +106,30 @@ Background execution is on the roadmap (below).
 ## Driving the task yourself
 
 The transparent flow is a convenience, not a requirement. Drop one layer to get the
-`CreateTaskResult` and poll on your own schedule:
+`CreateTaskResult`, then drive `tasks/*` with the typed functions in
+`mcp.client.tasks`:
 
-```python title="client.py" hl_lines="22 27-28"
+```python title="client.py" hl_lines="19 24 29"
 --8<-- "docs_src/tasks/tutorial003.py"
 ```
 
 * `session.call_tool(..., allow_claimed=True)` returns the typed
   `CreateTaskResult` instead of polling. Without the flag, an unexpected
   `CreateTaskResult` raises `RuntimeError` rather than leaking the widened union
   into code that expected a `CallToolResult`.
-* The `mcp.shared.tasks` wrappers — `GetTaskRequest`, `UpdateTaskRequest`,
-  `CancelTaskRequest` — drive the `tasks/*` methods over `session.send_request`,
-  and `GetTaskResult` parses the snapshot: `result` is set on a `completed` task,
-  `error` on a `failed` one, never both.
+* `get_task` is one `tasks/get`: it returns the `GetTaskResult` snapshot —
+  `result` is set on a `completed` task, `error` on a `failed` one, never both.
+* `wait_task` polls to a terminal status and returns the final `CallToolResult`,
+  raising the same typed errors as the transparent flow. Pass the
+  `CreateTaskResult` and its `pollIntervalMs` hint seeds the cadence — or pass a
+  bare task id: task ids are bearer capabilities (above), so a client that
+  reconnected, or restarted with nothing but the persisted id, can resume a task
+  it no longer holds the `CreateTaskResult` for.
+* `update_task` answers a task's in-task `inputRequests`, and `cancel_task` asks
+  the server to stop one. Both hide the empty acknowledgement and return `None`.
+  Cancellation is cooperative in SEP-2663 — it may never take effect, and in this
+  SDK the work has always finished already — so follow with `get_task` for the
+  status that actually resulted.
 
 ## Who sees what
 

docs/migration.md

@@ -484,11 +484,12 @@ Two reference extensions ship in their own modules:
   resultType on `tools/call`, so `Client.call_tool` admits 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 —
+  tasks surface as the typed `TaskFailedError`/`TaskCancelledError` (all task
+  errors share the `TaskError` base). Manual driving stays available —
   `client.session.call_tool(..., allow_claimed=True)` returns the typed
-  `CreateTaskResult`, and the `mcp.shared.tasks` request wrappers drive
-  `tasks/get`/`tasks/update`/`tasks/cancel` over `session.send_request`, with
+  `CreateTaskResult`, and the typed `mcp.client.tasks` functions
+  (`get_task`/`wait_task`/`update_task`/`cancel_task`) drive
+  `tasks/get`/`tasks/update`/`tasks/cancel` over the session, with
   the `Mcp-Name` routing header stamped automatically over Streamable HTTP.
   This is the core SEP-2663 surface — see [Tasks](advanced/tasks.md);
   background execution (`working` tasks), the in-task `input_required` loop

docs_src/tasks/tutorial003.py

@@ -1,12 +1,9 @@
-from typing import cast
-
-import mcp_types as types
-
 from mcp import Client
 from mcp.client import TasksExtension
+from mcp.client.tasks import get_task, wait_task
 from mcp.server.mcpserver import MCPServer
-from mcp.server.tasks import CreateTaskResult, Tasks
-from mcp.shared.tasks import GetTaskRequest, GetTaskRequestParams, GetTaskResult
+from mcp.server.tasks import Tasks
+from mcp.shared.tasks import CreateTaskResult
 
 mcp = MCPServer("bakery", extensions=[Tasks()])
 
@@ -24,8 +21,11 @@ async def main() -> None:
         print(created.status)
         # completed
 
-        request = GetTaskRequest(params=GetTaskRequestParams(task_id=created.task_id))
-        polled = await client.session.send_request(cast("types.ClientRequest", request), GetTaskResult)
+        polled = await get_task(client.session, created.task_id)
         assert polled.result is not None
         print(polled.result["content"])
         # [{'text': 'One mocha cake, ready.', 'type': 'text'}]
+
+        result = await wait_task(client.session, created)
+        print(result.content)
+        # [TextContent(text='One mocha cake, ready.')]

examples/stories/tasks/README.md

@@ -39,9 +39,11 @@ uv run python -m stories.tasks.client --http
   call returns the plain `CallToolResult` — the story guards its manual leg
   on the negotiated capability.
 - The manual leg — `session.call_tool(..., allow_claimed=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`.
+  typed `CreateTaskResult` (mirroring `allow_input_required`), and the typed
+  `mcp.client.tasks` functions drive it by hand: `get_task` fetches one
+  `tasks/get` snapshot with the outcome inlined, then `wait_task` polls the
+  same task to its final `CallToolResult` — from the bare persisted id, the
+  resume-after-reconnect shape.
 
 ## Scope
 

examples/stories/tasks/client.py

@@ -6,16 +6,15 @@
 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_claimed=True)` for the typed `CreateTaskResult`, and
-the shared `mcp.shared.tasks` wrappers over `session.send_request` for `tasks/get`.
+`session.call_tool(allow_claimed=True)` for the typed `CreateTaskResult`, and the
+typed `mcp.client.tasks` functions (`get_task`, `wait_task`) to drive `tasks/get`.
 """
 
-from typing import cast
-
 import mcp_types as types
 
 from mcp.client import Client, TasksExtension
-from mcp.shared.tasks import EXTENSION_ID, CreateTaskResult, GetTaskRequest, GetTaskRequestParams, GetTaskResult
+from mcp.client.tasks import get_task, wait_task
+from mcp.shared.tasks import EXTENSION_ID, CreateTaskResult
 from stories._harness import Target, run_client
 
 
@@ -38,20 +37,23 @@ async def main(target: Target, *, mode: str = "auto") -> None:
             return
         assert client.server_capabilities.extensions == {EXTENSION_ID: {}}
 
-        # The manual leg: the same flow driven by hand on the raw wire.
-        # allow_claimed=True hands back the typed CreateTaskResult instead of
-        # polling, and the shared SEP-2663 request wrappers fetch the outcome.
+        # The manual leg: the same flow driven by hand. allow_claimed=True hands
+        # back the typed CreateTaskResult instead of polling, and get_task fetches
+        # one tasks/get snapshot with the outcome inlined.
         created = await client.session.call_tool("render_report", {"title": "Q3", "sections": 1}, allow_claimed=True)
         assert isinstance(created, CreateTaskResult), created
 
-        task = await client.session.send_request(
-            cast("types.ClientRequest", GetTaskRequest(params=GetTaskRequestParams(task_id=created.task_id))),
-            GetTaskResult,
-        )
+        task = await get_task(client.session, created.task_id)
         assert task.status == "completed", task
         assert task.result is not None, task
         assert task.result["content"][0]["text"].startswith("# Q3"), task
 
+        # wait_task polls the same task to its final CallToolResult — from the
+        # bare persisted id here, the resume-after-reconnect shape.
+        final = await wait_task(client.session, created.task_id)
+        assert isinstance(final.content[0], types.TextContent), final
+        assert final.content[0].text.startswith("# Q3"), final
+
 
 if __name__ == "__main__":
     run_client(main)

src/mcp/__init__.py

@@ -59,11 +59,11 @@
 from mcp_types import Role as SamplingRole
 
 from .client._input_required import InputRequiredRoundsExceededError
-from .client._tasks import TaskCancelledError, TaskFailedError, TaskInputRequiredError, TasksExtension
 from .client.client import Client
 from .client.session import ClientSession
 from .client.session_group import ClientSessionGroup
 from .client.stdio import StdioServerParameters, stdio_client
+from .client.tasks import TaskCancelledError, TaskError, TaskFailedError, TaskInputRequiredError, TasksExtension
 from .server.session import ServerSession
 from .server.stdio import stdio_server
 from .shared.exceptions import MCPDeprecationWarning, MCPError, UrlElicitationRequiredError
@@ -130,6 +130,7 @@
     "StopReason",
     "SubscribeRequest",
     "TaskCancelledError",
+    "TaskError",
     "TaskFailedError",
     "TaskInputRequiredError",
     "TasksExtension",

src/mcp/client/__init__.py

@@ -1,7 +1,6 @@
 """MCP Client module."""
 
 from mcp.client._input_required import InputRequiredRoundsExceededError
-from mcp.client._tasks import TaskCancelledError, TaskFailedError, TaskInputRequiredError, TasksExtension
 from mcp.client._transport import Transport
 from mcp.client.caching import (
     CacheConfig,
@@ -22,6 +21,7 @@
     advertise,
 )
 from mcp.client.session import ClientSession
+from mcp.client.tasks import TaskCancelledError, TaskError, TaskFailedError, TaskInputRequiredError, TasksExtension
 
 __all__ = [
     "CacheConfig",
@@ -39,6 +39,7 @@
     "ResponseCacheStore",
     "ResultClaim",
     "TaskCancelledError",
+    "TaskError",
     "TaskFailedError",
     "TaskInputRequiredError",
     "TasksExtension",

src/mcp/client/_tasks.py

@@ -723,7 +723,8 @@ async def call_tool(
                 reached `input_required`; the SDK's automatic in-task input
                 loop is not implemented yet — drive the task manually via
                 `session.call_tool(..., allow_claimed=True)` and the
-                `mcp.shared.tasks` wrappers.
+                `mcp.client.tasks` functions (`get_task`, `update_task`,
+                `wait_task`). The task errors share the `TaskError` base.
         """
 
         async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult | Result: