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

This commit is contained in:
Max 2026-06-26 17:35:23 +02:00 • committed by GitHub
commit 08b62308d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1258 additions and 202 deletions

View file

@ -2061,7 +2061,7 @@ class InputRequiredResult(Result):
@model_validator(mode="after")
def _require_one_field(self) -> Self:
if self.input_requests is None and self.request_state is None:
if not self.input_requests and self.request_state is None:
raise ValueError("InputRequiredResult requires at least one of input_requests or request_state")
return self

View file

@ -58,6 +58,7 @@ from mcp_types import (
)
from mcp_types import Role as SamplingRole
from .client._input_required import InputRequiredRoundsExceededError
from .client.client import Client
from .client.session import ClientSession
from .client.session_group import ClientSessionGroup
@ -87,6 +88,7 @@ __all__ = [
"InitializeRequest",
"InitializeResult",
"InitializedNotification",
"InputRequiredRoundsExceededError",
"JSONRPCError",
"JSONRPCRequest",
"JSONRPCResponse",

View file

@ -1,8 +1,9 @@
"""MCP Client module."""
from mcp.client._input_required import InputRequiredRoundsExceededError
from mcp.client._transport import Transport
from mcp.client.client import Client
from mcp.client.context import ClientRequestContext
from mcp.client.session import ClientSession
__all__ = ["Client", "ClientRequestContext", "ClientSession", "Transport"]
__all__ = ["Client", "ClientRequestContext", "ClientSession", "InputRequiredRoundsExceededError", "Transport"]

View file

@ -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

View file

@ -5,7 +5,7 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping
from contextlib import AsyncExitStack
from dataclasses import KW_ONLY, dataclass, field
from typing import Any, Literal, TypeVar, overload
from typing import Any, Literal, TypeVar
import anyio
import mcp_types as types
@ -13,9 +13,12 @@ from mcp_types import (
CallToolResult,
CompleteResult,
EmptyResult,
ErrorData,
GetPromptResult,
Implementation,
InputRequest,
InputRequiredResult,
InputResponse,
InputResponses,
ListPromptsResult,
ListResourcesResult,
@ -32,10 +35,19 @@ from mcp_types import (
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS, MODERN_PROTOCOL_VERSIONS
from typing_extensions import deprecated
from mcp.client._input_required import DEFAULT_INPUT_REQUIRED_MAX_ROUNDS, run_input_required_driver
from mcp.client._memory import InMemoryTransport
from mcp.client._probe import negotiate_auto
from mcp.client._transport import Transport
from mcp.client.session import ClientSession, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
from mcp.client.session import (
ClientRequestContext,
ClientSession,
ElicitationFnT,
ListRootsFnT,
LoggingFnT,
MessageHandlerFnT,
SamplingFnT,
)
from mcp.client.streamable_http import streamable_http_client
from mcp.server import Server
from mcp.server.mcpserver import MCPServer
@ -51,6 +63,7 @@ initialize), or a modern protocol-version string (adopt directly). The ``str`` a
forward-compat; ``Client.__post_init__`` rejects anything outside that set at construction."""
_T = TypeVar("_T")
_ResultT = TypeVar("_ResultT")
_Connector = Callable[[AsyncExitStack, ConnectMode, bool], Awaitable["Dispatcher[Any]"]]
"""Resolved at ``__post_init__`` from the shape of ``server`` alone: enter whatever resources
@ -199,6 +212,11 @@ class Client:
elicitation_callback: ElicitationFnT | None = None
"""Callback for handling elicitation requests."""
input_required_max_rounds: int = DEFAULT_INPUT_REQUIRED_MAX_ROUNDS
"""Cap on `InputRequiredResult` retry rounds before `call_tool` / `get_prompt` /
`read_resource` give up. Use `client.session.<method>(..., allow_input_required=True)`
to drive the loop manually instead."""
_entered: bool = field(init=False, default=False)
_session: ClientSession | None = field(init=False, default=None)
_exit_stack: AsyncExitStack | None = field(init=False, default=None)
@ -356,17 +374,42 @@ class Client:
"""List available resource templates from the server."""
return await self.session.list_resource_templates(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
async def read_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> ReadResourceResult:
async def read_resource(
self,
uri: str,
*,
input_responses: InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
) -> ReadResourceResult:
"""Read a resource from the server.
If the server returns an `InputRequiredResult`, the embedded input
requests are dispatched to this client's sampling / elicitation / roots
callbacks and the read is retried automatically (up to
`input_required_max_rounds`).
Args:
uri: The URI of the resource to read.
input_responses: Responses to seed the first call with (e.g. when
resuming from a persisted `InputRequiredResult`).
request_state: Opaque state to seed the first call with.
meta: Additional metadata for the request.
Returns:
The resource content.
Raises:
InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted.
MCPError: A callback returned `ErrorData` for an embedded input request.
"""
return await self.session.read_resource(uri, meta=meta)
async def retry(r: InputResponses | None, s: str | None) -> ReadResourceResult | InputRequiredResult:
return await self.session.read_resource(
uri, input_responses=r, request_state=s, meta=meta, allow_input_required=True
)
return await self._drive_input_required(await retry(input_responses, request_state), retry)
async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> EmptyResult:
"""Subscribe to resource updates."""
@ -376,7 +419,6 @@ class Client:
"""Unsubscribe from resource updates."""
return await self.session.unsubscribe_resource(uri, meta=meta)
@overload
async def call_tool(
self,
name: str,
@ -387,69 +429,47 @@ class Client:
input_responses: InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: Literal[False] = False,
) -> CallToolResult: ...
@overload
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: float | None = None,
progress_callback: ProgressFnT | None = None,
*,
input_responses: InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool,
) -> CallToolResult | InputRequiredResult: ...
async def call_tool(
self,
name: str,
arguments: dict[str, Any] | None = None,
read_timeout_seconds: float | None = None,
progress_callback: ProgressFnT | None = None,
*,
input_responses: InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool = False,
) -> CallToolResult | InputRequiredResult:
) -> CallToolResult:
"""Call a tool on the server.
If the server returns an `InputRequiredResult`, the embedded input
requests are dispatched to this client's sampling / elicitation / roots
callbacks and the call is retried automatically (up to
`input_required_max_rounds`). To drive the loop yourself — e.g. to
persist `request_state` across process restarts — use
`client.session.call_tool(..., allow_input_required=True)`.
Args:
name: The name of the tool to call
arguments: Arguments to pass to the tool
read_timeout_seconds: Timeout for the tool call
progress_callback: Callback for progress updates
input_responses: Responses to a prior `InputRequiredResult.input_requests`
request_state: Opaque state echoed from a prior `InputRequiredResult`
meta: Additional metadata for the request
allow_input_required: When ``False`` (default), an `InputRequiredResult`
from the server raises `RuntimeError`; when ``True``, it is returned
so the caller can resolve the requests and retry.
name: The name of the tool to call.
arguments: Arguments to pass to the tool.
read_timeout_seconds: Timeout for each underlying `tools/call` round.
progress_callback: Callback for progress updates.
input_responses: Responses to seed the first call with (e.g. when
resuming from a persisted `InputRequiredResult`).
request_state: Opaque state to seed the first call with.
meta: Additional metadata for the request.
Returns:
The tool result. When ``allow_input_required=True``, may instead be an
`InputRequiredResult` carrying the server's input requests and opaque
``request_state`` for the retry.
The tool result.
Raises:
RuntimeError: If the server returns an `InputRequiredResult` and
``allow_input_required`` is ``False``.
InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted.
MCPError: A callback returned `ErrorData` for an embedded input request.
"""
# TODO(L84): stop forwarding allow_input_required; run the MRTR auto-loop driver here (S6).
return await self.session.call_tool(
name=name,
arguments=arguments,
read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
input_responses=input_responses,
request_state=request_state,
meta=meta,
allow_input_required=allow_input_required,
)
async def retry(r: InputResponses | None, s: str | None) -> CallToolResult | InputRequiredResult:
return await self.session.call_tool(
name,
arguments,
read_timeout_seconds=read_timeout_seconds,
progress_callback=progress_callback,
input_responses=r,
request_state=s,
meta=meta,
allow_input_required=True,
)
return await self._drive_input_required(await retry(input_responses, request_state), retry)
async def list_prompts(
self,
@ -461,19 +481,66 @@ class Client:
return await self.session.list_prompts(params=PaginatedRequestParams(cursor=cursor, _meta=meta))
async def get_prompt(
self, name: str, arguments: dict[str, str] | None = None, *, meta: RequestParamsMeta | None = None
self,
name: str,
arguments: dict[str, str] | None = None,
*,
input_responses: InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
) -> GetPromptResult:
"""Get a prompt from the server.
If the server returns an `InputRequiredResult`, the embedded input
requests are dispatched to this client's sampling / elicitation / roots
callbacks and the get is retried automatically (up to
`input_required_max_rounds`).
Args:
name: The name of the prompt
arguments: Arguments to pass to the prompt
meta: Additional metadata for the request
name: The name of the prompt.
arguments: Arguments to pass to the prompt.
input_responses: Responses to seed the first call with (e.g. when
resuming from a persisted `InputRequiredResult`).
request_state: Opaque state to seed the first call with.
meta: Additional metadata for the request.
Returns:
The prompt content.
Raises:
InputRequiredRoundsExceededError: `input_required_max_rounds` exhausted.
MCPError: A callback returned `ErrorData` for an embedded input request.
"""
return await self.session.get_prompt(name=name, arguments=arguments, meta=meta)
async def retry(r: InputResponses | None, s: str | None) -> GetPromptResult | InputRequiredResult:
return await self.session.get_prompt(
name, arguments, input_responses=r, request_state=s, meta=meta, allow_input_required=True
)
return await self._drive_input_required(await retry(input_responses, request_state), retry)
async def _drive_input_required(
self,
first: _ResultT | InputRequiredResult,
retry: Callable[[InputResponses | None, str | None], Awaitable[_ResultT | InputRequiredResult]],
) -> _ResultT:
"""Hand an `InputRequiredResult` to the SEP-2322 driver, or pass a terminal result through.
`dispatch` routes each embedded request through the same callback table
that serves legacy server→client RPCs, so the two paths stay
behaviourally identical by construction.
"""
if not isinstance(first, InputRequiredResult):
return first
session = self.session
async def dispatch(key: str, req: InputRequest) -> InputResponse | ErrorData:
ctx = ClientRequestContext(session=session, request_id=key, meta=req.params.meta if req.params else None)
return await session._dispatch_input_request(ctx, req) # pyright: ignore[reportPrivateUsage]
return await run_input_required_driver(
first, dispatch=dispatch, retry=retry, max_rounds=self.input_required_max_rounds
)
async def complete(
self,

View file

@ -185,6 +185,19 @@ ClientResponse: TypeAdapter[types.ClientResult | types.ErrorData] = TypeAdapter(
_CallToolResultAdapter: TypeAdapter[types.CallToolResult | types.InputRequiredResult] = TypeAdapter(
types.CallToolResult | types.InputRequiredResult
)
_GetPromptResultAdapter: TypeAdapter[types.GetPromptResult | types.InputRequiredResult] = TypeAdapter(
types.GetPromptResult | types.InputRequiredResult
)
_ReadResourceResultAdapter: TypeAdapter[types.ReadResourceResult | types.InputRequiredResult] = TypeAdapter(
types.ReadResourceResult | types.InputRequiredResult
)
def _input_required_unexpected(method: str) -> RuntimeError:
return RuntimeError(
"Server returned InputRequiredResult; pass allow_input_required=True to receive it "
f"and retry {method}(..., input_responses=..., request_state=result.request_state)."
)
class ClientSession:
@ -591,12 +604,64 @@ class ClientSession:
types.ListResourceTemplatesResult,
)
async def read_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.ReadResourceResult:
"""Send a resources/read request."""
return await self.send_request(
types.ReadResourceRequest(params=types.ReadResourceRequestParams(uri=uri, _meta=meta)),
types.ReadResourceResult,
@overload
async def read_resource(
self,
uri: str,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: Literal[False] = False,
) -> types.ReadResourceResult: ...
@overload
async def read_resource(
self,
uri: str,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool,
) -> types.ReadResourceResult | types.InputRequiredResult: ...
async def read_resource(
self,
uri: str,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool = False,
) -> types.ReadResourceResult | types.InputRequiredResult:
"""Send a resources/read request.
Args:
input_responses: Responses to a prior `InputRequiredResult.input_requests`.
request_state: Opaque state echoed from a prior `InputRequiredResult`.
allow_input_required: When `False` (default), an `InputRequiredResult`
from the server raises `RuntimeError`; when `True`, it is returned
so the caller can resolve the requests and retry.
Raises:
RuntimeError: If the server returns an `InputRequiredResult` and
`allow_input_required` is `False`.
"""
result = await self.send_request(
types.ReadResourceRequest(
params=types.ReadResourceRequestParams(
uri=uri,
input_responses=input_responses,
request_state=request_state,
_meta=meta,
),
),
_ReadResourceResultAdapter,
)
if isinstance(result, types.InputRequiredResult) and not allow_input_required:
raise _input_required_unexpected("read_resource")
return result
async def subscribe_resource(self, uri: str, *, meta: RequestParamsMeta | None = None) -> types.EmptyResult:
"""Send a resources/subscribe request."""
@ -689,10 +754,7 @@ class ClientSession:
await self._validate_tool_result(name, result)
if isinstance(result, types.InputRequiredResult) and not allow_input_required:
raise RuntimeError(
"Server returned InputRequiredResult; pass allow_input_required=True to receive it "
"and retry call_tool(..., input_responses=..., request_state=result.request_state)."
)
raise _input_required_unexpected("call_tool")
return result
def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]:
@ -734,18 +796,68 @@ class ClientSession:
"""
return await self.send_request(types.ListPromptsRequest(params=params), types.ListPromptsResult)
@overload
async def get_prompt(
self,
name: str,
arguments: dict[str, str] | None = None,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
) -> types.GetPromptResult:
"""Send a prompts/get request."""
return await self.send_request(
types.GetPromptRequest(params=types.GetPromptRequestParams(name=name, arguments=arguments, _meta=meta)),
types.GetPromptResult,
allow_input_required: Literal[False] = False,
) -> types.GetPromptResult: ...
@overload
async def get_prompt(
self,
name: str,
arguments: dict[str, str] | None = None,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool,
) -> types.GetPromptResult | types.InputRequiredResult: ...
async def get_prompt(
self,
name: str,
arguments: dict[str, str] | None = None,
*,
input_responses: types.InputResponses | None = None,
request_state: str | None = None,
meta: RequestParamsMeta | None = None,
allow_input_required: bool = False,
) -> types.GetPromptResult | types.InputRequiredResult:
"""Send a prompts/get request.
Args:
input_responses: Responses to a prior `InputRequiredResult.input_requests`.
request_state: Opaque state echoed from a prior `InputRequiredResult`.
allow_input_required: When `False` (default), an `InputRequiredResult`
from the server raises `RuntimeError`; when `True`, it is returned
so the caller can resolve the requests and retry.
Raises:
RuntimeError: If the server returns an `InputRequiredResult` and
`allow_input_required` is `False`.
"""
result = await self.send_request(
types.GetPromptRequest(
params=types.GetPromptRequestParams(
name=name,
arguments=arguments,
input_responses=input_responses,
request_state=request_state,
_meta=meta,
),
),
_GetPromptResultAdapter,
)
if isinstance(result, types.InputRequiredResult) and not allow_input_required:
raise _input_required_unexpected("get_prompt")
return result
async def complete(
self,
@ -829,13 +941,7 @@ class ClientSession:
ctx = ClientRequestContext(
session=self, request_id=dctx.request_id, meta=request.params.meta if request.params else None
)
match request:
case types.CreateMessageRequest(params=sampling_params):
response = await self._sampling_callback(ctx, sampling_params)
case types.ElicitRequest(params=elicit_params):
response = await self._elicitation_callback(ctx, elicit_params)
case types.ListRootsRequest(): # pragma: no branch
response = await self._list_roots_callback(ctx)
response = await self._dispatch_input_request(ctx, request)
client_response = ClientResponse.validate_python(response)
if isinstance(client_response, types.ErrorData):
raise MCPError.from_error_data(client_response)
@ -847,6 +953,23 @@ class ClientSession:
raise MCPError(code=INTERNAL_ERROR, message="Client callback returned an invalid result") from None
return dumped
async def _dispatch_input_request(
self, ctx: ClientRequestContext, req: types.InputRequest
) -> types.InputResponse | types.ErrorData:
"""Route a server-initiated input request to the matching constructor callback.
Shared by the legacy server→client RPC path (`_on_request`) and the
2026-07-28 multi-round-trip driver, which dispatches the embedded
`InputRequiredResult.input_requests` through the same callbacks.
"""
match req:
case types.CreateMessageRequest(params=p):
return await self._sampling_callback(ctx, p)
case types.ElicitRequest(params=p):
return await self._elicitation_callback(ctx, p)
case types.ListRootsRequest(): # pragma: no branch
return await self._list_roots_callback(ctx)
async def _on_notify(
self, dctx: DispatchContext[TransportContext], method: str, params: Mapping[str, Any] | None
) -> None: