49b9bb6724
Deploy Docs / deploy-docs (push) Failing after 1s
Conformance Tests / client-conformance (push) Failing after 3s
Conformance Tests / server-conformance (push) Failing after 1s
GitHub Actions Security Analysis / zizmor (push) Failing after 1s
CI / checks (push) Failing after 59m20s
CI / all-green (push) Has been cancelled
272 lines
12 KiB
Python
272 lines
12 KiB
Python
"""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
|