|
| 1 | +"""SEP-2663 client-side task polling driver. |
| 2 | +
|
| 3 | +When a server augments a `tools/call` into a task — a `CreateTaskResult` in |
| 4 | +place of the `CallToolResult` — the client polls `tasks/get` until the task |
| 5 | +reaches a terminal status and surfaces only the final result. SEP-2663 advises |
| 6 | +exactly this shape: "existing code returning a fixed shape ... can transparently |
| 7 | +drive the polling flow internally and surface only the final, completed result". |
| 8 | +This module implements that loop as a pure function so it stays testable with |
| 9 | +plain closures; `Client` builds the `get_task` closure over its session, |
| 10 | +`ClientSession` stays mechanics-only (mirroring `_input_required`). |
| 11 | +""" |
| 12 | + |
| 13 | +from __future__ import annotations |
| 14 | + |
| 15 | +from collections.abc import Awaitable, Callable, Sequence |
| 16 | +from typing import Any, cast |
| 17 | + |
| 18 | +import anyio |
| 19 | +from mcp_types import CallToolResult, ClientRequest, ErrorData |
| 20 | + |
| 21 | +from mcp.client.extension import ClaimContext, ClientExtension, ResultClaim |
| 22 | +from mcp.shared.exceptions import MCPError |
| 23 | +from mcp.shared.tasks import ( |
| 24 | + EXTENSION_ID, |
| 25 | + CreateTaskResult, |
| 26 | + GetTaskRequest, |
| 27 | + GetTaskRequestParams, |
| 28 | + GetTaskResult, |
| 29 | +) |
| 30 | + |
| 31 | +DEFAULT_POLL_INTERVAL_SECONDS = 1.0 |
| 32 | +"""Poll cadence when neither the snapshot nor the `CreateTaskResult` carries `pollIntervalMs`. |
| 33 | +
|
| 34 | +SEP-2663 makes the hint optional and only says clients SHOULD honor it when |
| 35 | +present; one second is the SDK's conservative default in its absence. |
| 36 | +""" |
| 37 | + |
| 38 | + |
| 39 | +class TaskFailedError(MCPError): |
| 40 | + """The task reached `failed`: a JSON-RPC error occurred during execution (SEP-2663). |
| 41 | +
|
| 42 | + Carries the JSON-RPC error inlined on `tasks/get` as `code`/`message`/`data`, |
| 43 | + plus the snapshot's optional `statusMessage` diagnostic. |
| 44 | + """ |
| 45 | + |
| 46 | + def __init__(self, error: ErrorData, status_message: str | None = None) -> None: |
| 47 | + super().__init__(code=error.code, message=error.message, data=error.data) |
| 48 | + self.status_message = status_message |
| 49 | + |
| 50 | + |
| 51 | +class TaskCancelledError(RuntimeError): |
| 52 | + """The task reached `cancelled` before producing a result (SEP-2663).""" |
| 53 | + |
| 54 | + def __init__(self, task_id: str, status_message: str | None = None) -> None: |
| 55 | + detail = f": {status_message}" if status_message is not None else "" |
| 56 | + super().__init__(f"Task {task_id!r} was cancelled{detail}") |
| 57 | + self.task_id = task_id |
| 58 | + self.status_message = status_message |
| 59 | + |
| 60 | + |
| 61 | +class TaskInputRequiredError(RuntimeError): |
| 62 | + """The task reached `input_required`, which this driver does not drive yet. |
| 63 | +
|
| 64 | + SEP-2663's in-task input loop (fulfil `inputRequests` via `tasks/update`) is |
| 65 | + a deferred follow-up in this SDK. Drive it manually: poll with |
| 66 | + `mcp.shared.tasks.GetTaskRequest` and answer with |
| 67 | + `mcp.shared.tasks.UpdateTaskRequest` over `session.send_request`. |
| 68 | + """ |
| 69 | + |
| 70 | + def __init__(self, task_id: str) -> None: |
| 71 | + super().__init__( |
| 72 | + f"Task {task_id!r} requires in-task input (status `input_required`); the SDK's automatic " |
| 73 | + "in-task input loop is not implemented yet. Drive it manually with the `mcp.shared.tasks` " |
| 74 | + "request wrappers (`GetTaskRequest`/`UpdateTaskRequest`) over `session.send_request`." |
| 75 | + ) |
| 76 | + self.task_id = task_id |
| 77 | + |
| 78 | + |
| 79 | +async def run_task_driver( |
| 80 | + created: CreateTaskResult, |
| 81 | + *, |
| 82 | + get_task: Callable[[str], Awaitable[GetTaskResult]], |
| 83 | + sleep: Callable[[float], Awaitable[None]] = anyio.sleep, |
| 84 | +) -> CallToolResult: |
| 85 | + """Poll a `CreateTaskResult` to its final `CallToolResult`. |
| 86 | +
|
| 87 | + Polls `tasks/get` (via `get_task`) until the task reaches a terminal status. |
| 88 | + Between polls it honors the SEP-2663 `pollIntervalMs` hint: each non-terminal |
| 89 | + snapshot sleeps its own `poll_interval_ms`, falling back to the |
| 90 | + `CreateTaskResult`'s, then to `DEFAULT_POLL_INTERVAL_SECONDS`. |
| 91 | +
|
| 92 | + The loop deliberately imposes no round cap or deadline of its own: SEP-2663 |
| 93 | + tasks represent unbounded server-side work, so how long to wait is the |
| 94 | + caller's policy — cancel via an enclosing anyio cancel scope, or bound each |
| 95 | + `tasks/get` round with the session read timeout the `get_task` closure |
| 96 | + carries. |
| 97 | +
|
| 98 | + Args: |
| 99 | + created: The `CreateTaskResult` the augmented request returned. |
| 100 | + get_task: Sends one `tasks/get` for the given task id and returns the |
| 101 | + parsed `GetTaskResult` snapshot. |
| 102 | + sleep: Awaits the given number of seconds between polls (injectable for |
| 103 | + deterministic tests). |
| 104 | +
|
| 105 | + Raises: |
| 106 | + TaskFailedError: The task reached `failed`; carries the inlined JSON-RPC error. |
| 107 | + TaskCancelledError: The task reached `cancelled`. |
| 108 | + TaskInputRequiredError: The task reached `input_required` (the in-task |
| 109 | + input loop is not implemented yet). |
| 110 | + RuntimeError: The server violated SEP-2663 — a `completed` snapshot |
| 111 | + without `result`, or a `failed` snapshot without `error`. |
| 112 | + """ |
| 113 | + while True: |
| 114 | + snapshot = await get_task(created.task_id) |
| 115 | + if snapshot.status == "completed": |
| 116 | + if snapshot.result is None: |
| 117 | + raise RuntimeError( |
| 118 | + f"Task {created.task_id!r} is `completed` but carries no `result` (SEP-2663 violation)" |
| 119 | + ) |
| 120 | + return CallToolResult.model_validate(snapshot.result, by_name=False) |
| 121 | + if snapshot.status == "failed": |
| 122 | + if snapshot.error is None: |
| 123 | + raise RuntimeError(f"Task {created.task_id!r} is `failed` but carries no `error` (SEP-2663 violation)") |
| 124 | + raise TaskFailedError(ErrorData.model_validate(snapshot.error), snapshot.status_message) |
| 125 | + if snapshot.status == "cancelled": |
| 126 | + raise TaskCancelledError(created.task_id, snapshot.status_message) |
| 127 | + if snapshot.status == "input_required": |
| 128 | + raise TaskInputRequiredError(created.task_id) |
| 129 | + interval_ms = snapshot.poll_interval_ms if snapshot.poll_interval_ms is not None else created.poll_interval_ms |
| 130 | + await sleep(DEFAULT_POLL_INTERVAL_SECONDS if interval_ms is None else interval_ms / 1000) |
| 131 | + |
| 132 | + |
| 133 | +class TasksExtension(ClientExtension): |
| 134 | + """SEP-2663 Tasks as a client extension. |
| 135 | +
|
| 136 | + Declares `io.modelcontextprotocol/tasks` and claims the `task` resultType on |
| 137 | + `tools/call`: a `CreateTaskResult` is resolved by polling `tasks/get` to the |
| 138 | + final `CallToolResult` via `run_task_driver`. |
| 139 | + """ |
| 140 | + |
| 141 | + identifier = EXTENSION_ID |
| 142 | + |
| 143 | + def claims(self) -> Sequence[ResultClaim[Any]]: |
| 144 | + return (ResultClaim(result_type="task", model=CreateTaskResult, resolve=_resolve_created_task),) |
| 145 | + |
| 146 | + |
| 147 | +async def _resolve_created_task(created: CreateTaskResult, ctx: ClaimContext) -> CallToolResult: |
| 148 | + """Poll an SEP-2663 task to its final `CallToolResult` (the transparent flow). |
| 149 | +
|
| 150 | + Each `tasks/get` round goes through `ctx.session.send_request`, so it carries |
| 151 | + the caller's per-request read timeout (falling back to the session read |
| 152 | + timeout) and the `Mcp-Name` routing header. `Client.call_tool` re-validates |
| 153 | + the returned result against the tool's output schema, exactly as on the |
| 154 | + direct path. |
| 155 | + """ |
| 156 | + session = ctx.session |
| 157 | + |
| 158 | + async def get_task(task_id: str) -> GetTaskResult: |
| 159 | + request = GetTaskRequest(params=GetTaskRequestParams(task_id=task_id)) |
| 160 | + return await session.send_request( |
| 161 | + cast("ClientRequest", request), GetTaskResult, request_read_timeout_seconds=ctx.read_timeout_seconds |
| 162 | + ) |
| 163 | + |
| 164 | + return await run_task_driver(created, get_task=get_task) |
0 commit comments