Client auto-resolves InputRequiredResult via existing callbacks (SEP-2322) (#2998)

src/mcp/client/_input_required.py

@@ -0,0 +1,127 @@
+"""SEP-2322 client-side multi-round-trip driver.
+
+When a server returns `InputRequiredResult` instead of the normal result of a
+`tools/call` / `prompts/get` / `resources/read`, the client fulfils the
+embedded `input_requests` (sampling, elicitation, roots) and retries the
+original request carrying the responses and the echoed opaque `request_state`.
+This module implements that retry loop as a pure function so it can drive any
+of the three methods identically; `Client` builds the `dispatch` and `retry`
+closures, `ClientSession` stays mechanics-only.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TypeVar
+
+import anyio
+import anyio.abc
+from mcp_types import ErrorData, InputRequest, InputRequiredResult, InputResponse, InputResponses
+
+from mcp.shared.exceptions import MCPError
+
+DEFAULT_INPUT_REQUIRED_MAX_ROUNDS = 10
+"""Default cap on `InputRequiredResult` retry rounds before the driver gives up.
+
+Matches the typescript-sdk default; csharp-sdk and go-sdk use the same value
+as a hard constant.
+"""
+
+_STATE_ONLY_BACKOFF_INITIAL_SECONDS = 0.05
+"""First sleep when an `InputRequiredResult` carries only `request_state` (no input requests)."""
+
+_STATE_ONLY_BACKOFF_CAP_SECONDS = 0.25
+"""Upper bound on the state-only backoff sleep; reached after three consecutive state-only legs."""
+
+
+ResultT = TypeVar("ResultT")
+
+
+class InputRequiredRoundsExceededError(RuntimeError):
+    """The server kept returning `InputRequiredResult` past the configured `max_rounds`."""
+
+    def __init__(self, max_rounds: int) -> None:
+        super().__init__(
+            f"Server returned InputRequiredResult for more than {max_rounds} rounds; "
+            "raise input_required_max_rounds on the Client, or use "
+            "client.session.<method>(..., allow_input_required=True) to drive the loop manually."
+        )
+        self.max_rounds = max_rounds
+
+
+async def run_input_required_driver(
+    first: InputRequiredResult,
+    *,
+    dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]],
+    retry: Callable[[InputResponses | None, str | None], Awaitable[ResultT | InputRequiredResult]],
+    max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS,
+) -> ResultT:
+    """Resolve an `InputRequiredResult` to its terminal result.
+
+    Loops until `retry` returns a non-`InputRequiredResult`, or `max_rounds` is
+    exhausted. Each round either dispatches all `input_requests` concurrently
+    and retries with the collected responses, or — when the server sent only
+    `request_state` — sleeps with exponential backoff (50ms doubling to a 250ms
+    cap, reset by any leg that carries input requests) and retries empty.
+    `request_state` is passed through byte-exact and never inspected.
+
+    Args:
+        first: The `InputRequiredResult` the original call returned.
+        dispatch: Runs one embedded `InputRequest` through the client's
+            sampling / elicitation / roots callbacks. Called concurrently per
+            request key. An `ErrorData` return aborts the loop as an `MCPError`.
+        retry: Re-issues the original request with the collected responses and
+            the latest `request_state`. Each call mints a fresh JSON-RPC id.
+        max_rounds: Cap on retry rounds.
+
+    Raises:
+        InputRequiredRoundsExceededError: `max_rounds` exhausted.
+        MCPError: A `dispatch` call returned `ErrorData`.
+    """
+    rounds = 0
+    state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS
+    current: ResultT | InputRequiredResult = first
+    while isinstance(current, InputRequiredResult):
+        rounds += 1
+        if rounds > max_rounds:
+            raise InputRequiredRoundsExceededError(max_rounds)
+        if current.input_requests:
+            state_only_delay = _STATE_ONLY_BACKOFF_INITIAL_SECONDS
+            responses: InputResponses | None = await _dispatch_all(current.input_requests, dispatch)
+        else:
+            await anyio.sleep(state_only_delay)
+            state_only_delay = min(state_only_delay * 2, _STATE_ONLY_BACKOFF_CAP_SECONDS)
+            responses = None
+        current = await retry(responses, current.request_state)
+    return current
+
+
+async def _dispatch_all(
+    requests: dict[str, InputRequest],
+    dispatch: Callable[[str, InputRequest], Awaitable[InputResponse | ErrorData]],
+) -> InputResponses:
+    """Run `dispatch` concurrently for every key, raising `MCPError` on the first `ErrorData`.
+
+    The first task to return `ErrorData` cancels its siblings via the task
+    group's cancel scope, so a refused input does not wait on a slow peer.
+    A callback that *raises* propagates as an `ExceptionGroup` like any other
+    task-group failure.
+    """
+    responses: InputResponses = {}
+    refused: ErrorData | None = None
+
+    async def run_one(tg: anyio.abc.TaskGroup, key: str, req: InputRequest) -> None:
+        nonlocal refused
+        result = await dispatch(key, req)
+        if isinstance(result, ErrorData):
+            refused = result
+            tg.cancel_scope.cancel()
+        else:
+            responses[key] = result
+
+    async with anyio.create_task_group() as tg:
+        for key, req in requests.items():
+            tg.start_soon(run_one, tg, key, req)
+    if refused is not None:
+        raise MCPError.from_error_data(refused)
+    return responses