chore: import upstream snapshot with attribution

tests/client/test_input_required.py

@@ -0,0 +1,271 @@
+"""Unit tests for the SEP-2322 client-side multi-round-trip driver.
+
+`run_input_required_driver` is pure: it takes the first `InputRequiredResult`
+plus `dispatch` / `retry` closures and loops until a terminal result. These
+tests build those closures by hand (scripted lists, recording lists) so the
+driver is exercised without a `ClientSession`. Integration against a real
+server lives in `test_client.py`.
+"""
+
+import anyio
+import pytest
+from inline_snapshot import snapshot
+from mcp_types import (
+    INVALID_REQUEST,
+    CallToolResult,
+    ElicitRequest,
+    ElicitRequestFormParams,
+    ElicitResult,
+    ErrorData,
+    InputRequest,
+    InputRequiredResult,
+    InputResponse,
+    InputResponses,
+    TextContent,
+)
+from trio.testing import MockClock
+
+from mcp import MCPError
+from mcp.client._input_required import (
+    _STATE_ONLY_BACKOFF_CAP_SECONDS,
+    _STATE_ONLY_BACKOFF_INITIAL_SECONDS,
+    DEFAULT_INPUT_REQUIRED_MAX_ROUNDS,
+    InputRequiredRoundsExceededError,
+    run_input_required_driver,
+)
+
+pytestmark = pytest.mark.anyio
+
+
+@pytest.fixture(autouse=True)
+def _module_runner_lease() -> None:
+    """Opt out of the shared per-module event loop: this module parametrizes `anyio_backend`."""
+
+
+def _elicit(message: str = "What is your name?") -> ElicitRequest:
+    schema = {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
+    return ElicitRequest(params=ElicitRequestFormParams(message=message, requested_schema=schema))
+
+
+async def _never_dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+    """Dispatch closure for tests whose script never carries `input_requests`."""
+    raise NotImplementedError
+
+
+async def test_single_round_dispatches_then_retries_to_terminal_result() -> None:
+    """One `InputRequiredResult` with one elicit request: dispatch runs once,
+    retry runs once with the collected response, and the terminal result is returned."""
+    first = InputRequiredResult(input_requests={"ask": _elicit()})
+    terminal = CallToolResult(content=[TextContent(text="done")])
+    dispatched: list[tuple[str, InputRequest]] = []
+    retried: list[tuple[InputResponses | None, str | None]] = []
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        dispatched.append((key, req))
+        return ElicitResult(action="accept", content={"name": "Ada"})
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        retried.append((responses, state))
+        return terminal
+
+    with anyio.fail_after(5):
+        result = await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=3)
+
+    assert result is terminal
+    assert first.input_requests is not None
+    assert dispatched == [("ask", first.input_requests["ask"])]
+    assert retried == [({"ask": ElicitResult(action="accept", content={"name": "Ada"})}, None)]
+
+
+async def test_multi_round_loops_until_retry_returns_non_input_required() -> None:
+    """Two consecutive `InputRequiredResult` legs followed by a terminal result:
+    the driver dispatches and retries each leg in order."""
+    terminal = CallToolResult(content=[TextContent(text="done")])
+    script: list[CallToolResult | InputRequiredResult] = [
+        InputRequiredResult(input_requests={"b": _elicit("second?")}),
+        terminal,
+    ]
+    retried: list[tuple[InputResponses | None, str | None]] = []
+    dispatched_keys: list[str] = []
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        dispatched_keys.append(key)
+        return ElicitResult(action="decline")
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        retried.append((responses, state))
+        return script.pop(0)
+
+    first = InputRequiredResult(input_requests={"a": _elicit("first?")})
+    with anyio.fail_after(5):
+        result = await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=5)
+
+    assert result is terminal
+    assert dispatched_keys == ["a", "b"]
+    assert retried == snapshot(
+        [
+            ({"a": ElicitResult(action="decline")}, None),
+            ({"b": ElicitResult(action="decline")}, None),
+        ]
+    )
+
+
+async def test_exceeding_max_rounds_raises_with_the_configured_cap() -> None:
+    """When every retry returns another `InputRequiredResult`, the driver gives
+    up after `max_rounds` retries with `InputRequiredRoundsExceededError`."""
+    rounds: list[int] = []
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        return ElicitResult(action="decline")
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        rounds.append(len(rounds))
+        return InputRequiredResult(input_requests={"again": _elicit()})
+
+    first = InputRequiredResult(input_requests={"again": _elicit()})
+    with anyio.fail_after(5):
+        with pytest.raises(InputRequiredRoundsExceededError) as exc:
+            await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=3)
+
+    assert exc.value.max_rounds == 3
+    # `first` counts as round 1; rounds 1-3 each retry, round 4 trips the cap before dispatching.
+    assert len(rounds) == 3
+
+
+async def test_dispatch_returning_error_data_aborts_the_loop_as_mcp_error() -> None:
+    """SDK-defined: a callback that refuses an embedded request returns
+    `ErrorData`; the driver surfaces it as `MCPError` rather than retrying."""
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        return ErrorData(code=INVALID_REQUEST, message="not supported")
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        raise NotImplementedError  # unreachable: dispatch errored before any retry
+
+    first = InputRequiredResult(input_requests={"ask": _elicit()})
+    with anyio.fail_after(5):
+        with pytest.raises(MCPError) as exc:
+            await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=3)
+    assert exc.value.error.code == INVALID_REQUEST
+
+
+async def test_request_state_passes_through_byte_identical() -> None:
+    """`request_state` is opaque to the driver: each leg's value reaches `retry`
+    as the same object the server sent, never parsed or rebuilt."""
+    states = ['{"round": 1, "tag": "héllo"}', '{"round": 2, "tag": "wörld"}']
+    received_states: list[str | None] = []
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        return ElicitResult(action="decline")
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        received_states.append(state)
+        if len(received_states) < 2:
+            return InputRequiredResult(input_requests={"k": _elicit()}, request_state=states[1])
+        return CallToolResult(content=[])
+
+    first = InputRequiredResult(input_requests={"k": _elicit()}, request_state=states[0])
+    with anyio.fail_after(5):
+        await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=3)
+
+    assert received_states[0] is states[0]
+    assert received_states[1] is states[1]
+
+
+# Runs on trio's autojumping virtual clock so the backoff sleeps add zero
+# wall-clock and the recorded deltas are exact: `anyio.sleep` advances the
+# MockClock by precisely the requested duration once every task is idle.
+@pytest.mark.parametrize(
+    "anyio_backend",
+    [pytest.param(("trio", {"clock": MockClock(autojump_threshold=0)}), id="trio-mockclock")],
+)
+async def test_state_only_legs_back_off_exponentially_to_the_cap() -> None:
+    """SDK-defined pacing: state-only legs sleep 50ms, 100ms, 200ms, then cap at
+    250ms. Six state-only rounds → deltas `[0.05, 0.1, 0.2, 0.25, 0.25, 0.25]`."""
+    retry_times: list[float] = []
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        retry_times.append(anyio.current_time())
+        assert responses is None
+        if len(retry_times) == 6:
+            return CallToolResult(content=[])
+        return InputRequiredResult(request_state="poll")
+
+    start = anyio.current_time()
+    first = InputRequiredResult(request_state="poll")
+    await run_input_required_driver(first, dispatch=_never_dispatch, retry=retry, max_rounds=10)
+
+    deltas = [round(retry_times[0] - start, 9)] + [
+        round(retry_times[i] - retry_times[i - 1], 9) for i in range(1, len(retry_times))
+    ]
+    assert deltas == snapshot([0.05, 0.1, 0.2, 0.25, 0.25, 0.25])
+    assert _STATE_ONLY_BACKOFF_INITIAL_SECONDS == 0.05
+    assert _STATE_ONLY_BACKOFF_CAP_SECONDS == 0.25
+
+
+@pytest.mark.parametrize(
+    "anyio_backend",
+    [pytest.param(("trio", {"clock": MockClock(autojump_threshold=0)}), id="trio-mockclock")],
+)
+async def test_backoff_counter_resets_after_a_leg_with_input_requests() -> None:
+    """A leg carrying `input_requests` resets `consecutive_state_only`: the
+    next state-only leg sleeps the initial 50ms again, not the prior position."""
+    # state-only, state-only, dispatch leg (no sleep), state-only, terminal.
+    script: list[CallToolResult | InputRequiredResult] = [
+        InputRequiredResult(request_state="s"),
+        InputRequiredResult(input_requests={"k": _elicit()}),
+        InputRequiredResult(request_state="s"),
+        CallToolResult(content=[]),
+    ]
+    retry_times: list[float] = []
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        return ElicitResult(action="decline")
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        retry_times.append(anyio.current_time())
+        return script.pop(0)
+
+    start = anyio.current_time()
+    first = InputRequiredResult(request_state="s")
+    await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=10)
+
+    deltas = [round(retry_times[0] - start, 9)] + [
+        round(retry_times[i] - retry_times[i - 1], 9) for i in range(1, len(retry_times))
+    ]
+    # 0.05, 0.1 (two state-only), 0.0 (dispatch leg has no sleep), 0.05 (reset).
+    assert deltas == snapshot([0.05, 0.1, 0.0, 0.05])
+
+
+async def test_input_requests_are_dispatched_concurrently() -> None:
+    """All `input_requests` in a round are dispatched together: each dispatch
+    blocks on a shared gate that only opens once every key has started, so a
+    sequential implementation would deadlock under the `fail_after`."""
+    keys = ["a", "b", "c"]
+    started: set[str] = set()
+    all_started = anyio.Event()
+
+    async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
+        started.add(key)
+        if started == set(keys):
+            all_started.set()
+        await all_started.wait()  # blocks until every sibling is in-flight
+        return ElicitResult(action="accept", content={"name": key})
+
+    received: list[InputResponses | None] = []
+
+    async def retry(responses: InputResponses | None, state: str | None) -> CallToolResult | InputRequiredResult:
+        received.append(responses)
+        return CallToolResult(content=[])
+
+    first = InputRequiredResult(input_requests={k: _elicit() for k in keys})
+    with anyio.fail_after(5):
+        await run_input_required_driver(first, dispatch=dispatch, retry=retry, max_rounds=2)
+
+    assert received[0] is not None
+    assert received[0] == {k: ElicitResult(action="accept", content={"name": k}) for k in keys}
+
+
+def test_default_max_rounds_constant() -> None:
+    """SDK-defined default; matches the typescript-sdk."""
+    assert DEFAULT_INPUT_REQUIRED_MAX_ROUNDS == 10