Add resolver dependency injection for MCPServer tools (#2969)

Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>
This commit is contained in:
Marcelo Trylesinski 2026-06-29 12:51:46 +02:00 • committed by GitHub
commit f664db8952
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1844 additions and 15 deletions

View file

@ -11,6 +11,7 @@ from mcp_types.v2025_11_25 import PrimitiveSchemaDefinition
from pydantic import BaseModel, ValidationError
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from pydantic_core import core_schema
from typing_extensions import TypeAliasType
from mcp.server.session import ServerSession
@ -36,7 +37,11 @@ class CancelledElicitation(BaseModel):
action: Literal["cancel"] = "cancel"
ElicitationResult = AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation
ElicitationResult = TypeAliasType(
"ElicitationResult",
AcceptedElicitation[ElicitSchemaModelT] | DeclinedElicitation | CancelledElicitation,
type_params=(ElicitSchemaModelT,),
)
class AcceptedUrlElicitation(BaseModel):

View file

@ -5,6 +5,14 @@ from mcp_types import Icon
from mcp.server.extension import Extension, MethodBinding, ResourceBinding, ToolBinding
from .context import Context
from .resolve import (
AcceptedElicitation,
CancelledElicitation,
DeclinedElicitation,
Elicit,
ElicitationResult,
Resolve,
)
from .resources import DEFAULT_RESOURCE_SECURITY, ResourceSecurity
from .server import MCPServer, require_client_extension
from .utilities.types import Audio, Image
@ -15,6 +23,12 @@ __all__ = [
"Image",
"Audio",
"Icon",
"Resolve",
"Elicit",
"ElicitationResult",
"AcceptedElicitation",
"DeclinedElicitation",
"CancelledElicitation",
"Extension",
"ToolBinding",
"ResourceBinding",

View file

@ -1,7 +1,7 @@
from __future__ import annotations
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, Generic
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Any, Generic, cast
from mcp_types import ClientCapabilities, InputResponseRequestParams, InputResponses, LoggingLevel
from pydantic import AnyUrl, BaseModel
@ -217,6 +217,16 @@ class Context(BaseModel, Generic[LifespanContextT, RequestT]):
"""
return self.request_context.meta.get("client_id") if self.request_context.meta else None # pragma: no cover
@property
def headers(self) -> Mapping[str, str] | None:
"""Request headers carried by this message, when the transport has them.
Populated by HTTP-based transports; `None` on stdio or when the
transport's request object carries no headers. Headers are
client-supplied input - never treat one as an identity assertion.
"""
return cast("Mapping[str, str] | None", getattr(self.request_context.request, "headers", None))
@property
def request_id(self) -> str:
"""Get the unique ID for this request."""

View file

@ -0,0 +1,324 @@
"""Resolver dependency injection for MCPServer tools.
A tool parameter annotated `Annotated[T, Resolve(fn)]` is filled by running the
resolver `fn` before the tool body, instead of from the LLM-supplied arguments.
Resolvers form a DAG: a resolver may declare its own `Resolve(...)` dependencies,
take tool arguments by name, and take the `Context`. A resolver may return
`Elicit[T]` to ask the client; the framework runs the elicitation and injects the
answer.
Whether the consumer receives the unwrapped model or the full
`ElicitationResult` union is decided by the consumer's annotation:
- `Annotated[T, Resolve(fn)]` -> unwrapped `T`; decline/cancel aborts the call.
- `Annotated[ElicitationResult[T], Resolve(fn)]` (or a specific member) -> the
full outcome; the consumer branches on accept/decline/cancel.
Each resolver runs at most once per `tools/call` (memoized by function identity).
"""
from __future__ import annotations
import inspect
import typing
from collections.abc import Callable, Hashable, Mapping
from typing import Annotated, Any, Generic, cast, get_args, get_origin
import anyio.to_thread
from pydantic import BaseModel
from typing_extensions import TypeVar
from mcp.server.elicitation import (
AcceptedElicitation,
CancelledElicitation,
DeclinedElicitation,
ElicitationResult,
)
from mcp.server.mcpserver.context import Context
from mcp.server.mcpserver.exceptions import InvalidSignature, ToolError
from mcp.shared._callable_inspection import is_async_callable
T = TypeVar("T", bound=BaseModel)
# The union members the framework injects when a consumer opts into the outcome.
_ELICITATION_RESULT_MEMBERS = (AcceptedElicitation, DeclinedElicitation, CancelledElicitation)
class Resolve:
"""Marker for `Annotated[T, Resolve(fn)]`: fill the parameter by running `fn`."""
def __init__(self, fn: Callable[..., Any]) -> None:
self.fn = fn
class Elicit(Generic[T]):
"""A resolver's request to ask the client.
Returned from a resolver to signal that the value must be elicited. The
framework runs `ctx.elicit(message, schema)` and injects the outcome.
"""
def __init__(self, message: str, schema: type[T]) -> None:
self.message = message
self.schema = schema
class _ParamPlan:
"""How to fill one resolver parameter, decided once at registration."""
kind: str # "context" | "resolve" | "by_name"
resolve: Resolve | None
wants_union: bool
def __init__(self, kind: str, resolve: Resolve | None = None, wants_union: bool = False) -> None:
self.kind = kind
self.resolve = resolve
self.wants_union = wants_union
class _ResolverPlan:
"""A resolver's parameters and whether it is async, analyzed once."""
def __init__(self, fn: Callable[..., Any], params: dict[str, _ParamPlan], is_async: bool) -> None:
self.fn = fn
self.params = params
self.is_async = is_async
def _type_hints(fn: Callable[..., Any]) -> dict[str, Any]:
"""Resolve type hints for a function or a callable object.
`typing.get_type_hints` raises on a callable *instance*; fall back to its
`__call__`. Returns an empty mapping when hints cannot be resolved, matching
`find_context_parameter`'s tolerance so callables without annotations (or with
unresolvable ones) simply have no resolved parameters.
"""
target = fn if inspect.isroutine(fn) else getattr(type(fn), "__call__", fn)
try:
return typing.get_type_hints(target, include_extras=True)
except Exception:
return {}
def _resolver_name(fn: Callable[..., Any]) -> str:
"""Best-effort display name for error messages (callable objects lack `__name__`)."""
return getattr(fn, "__name__", None) or type(fn).__name__
def find_resolved_parameters(fn: Callable[..., Any]) -> dict[str, tuple[Resolve, bool]]:
"""Find parameters of `fn` annotated `Annotated[_, Resolve(...)]`.
Returns a mapping of parameter name to `(Resolve, wants_union)`, where
`wants_union` is True when the annotated type is an `ElicitationResult` member
(the consumer wants the full outcome rather than the unwrapped model).
"""
hints = _type_hints(fn)
resolved: dict[str, tuple[Resolve, bool]] = {}
for name in inspect.signature(fn).parameters:
annotation = hints.get(name)
if get_origin(annotation) is not Annotated:
# A `Resolve` marker is only honored at the top level; flag (rather than
# silently drop) one buried in a union, e.g. `Annotated[T, Resolve(f)] | None`.
if _contains_resolve(annotation):
raise InvalidSignature(
f"Parameter {name!r} of {_resolver_name(fn)!r} wraps `Resolve(...)` in a "
"union; annotate the parameter directly as `Annotated[T, Resolve(...)]`"
)
continue
type_arg, *metadata = get_args(annotation)
marker = next((m for m in metadata if isinstance(m, Resolve)), None)
if marker is not None:
resolved[name] = (marker, _wants_union(type_arg))
return resolved
def _contains_resolve(annotation: Any) -> bool:
"""True when a `Resolve` marker is nested inside `annotation` (e.g. a union member)."""
if get_origin(annotation) is Annotated:
return any(isinstance(m, Resolve) for m in get_args(annotation)[1:])
return any(_contains_resolve(arg) for arg in get_args(annotation))
def _wants_union(type_arg: Any) -> bool:
"""True when `type_arg` is an `ElicitationResult` member (or a union of them).
Handles the subscripted `ElicitationResult[T]` alias (a `TypeAliasType` whose
union is on the origin's `__value__`), the bare `ElicitationResult` alias (the
`__value__` is on `type_arg` itself), an explicit `AcceptedElicitation[T] | ...`
union, and a single member.
"""
# Unwrap the `ElicitationResult` alias whether it is bare or subscripted.
value = getattr(type_arg, "__value__", None) or getattr(get_origin(type_arg), "__value__", None)
if value is not None:
type_arg = value
members = get_args(type_arg) if get_origin(type_arg) is not None else (type_arg,)
return any(isinstance(m, type) and issubclass(m, _ELICITATION_RESULT_MEMBERS) for m in members)
def _resolver_key(fn: Callable[..., Any]) -> Hashable:
"""Identity key for memoizing a resolver.
A bound method - pure-python (`inspect.ismethod`) or built-in (e.g. `obj.meth`
on a C-extension type) - is recreated on each attribute access, so `id(fn)`
differs every time. Key it by its underlying function (or name) plus its
`__self__` identity so `auth.login` referenced in two places memoizes to one
call. Everything else keys by `id`, so two distinct callables never collide
even if they compare equal.
"""
bound_self = getattr(fn, "__self__", None)
if bound_self is not None:
# `__func__` (pure-python) has a stable identity; built-ins expose only a
# stable `__name__`. Use the function's id or the name's value accordingly.
func = getattr(fn, "__func__", None)
underlying: Hashable = id(func) if func is not None else getattr(fn, "__name__", id(fn))
return (underlying, id(bound_self))
return id(fn)
def build_resolver_plans(
resolved_params: Mapping[str, tuple[Resolve, bool]],
tool_arg_names: set[str],
) -> dict[Hashable, _ResolverPlan]:
"""Statically analyze the resolver DAG rooted at a tool's resolved parameters.
Raises:
InvalidSignature: If a resolver has a cyclic dependency, or a resolver
parameter cannot be classified (not a `Context`, a nested `Resolve`,
or a tool argument by name).
"""
plans: dict[Hashable, _ResolverPlan] = {}
def analyze(fn: Callable[..., Any], stack: tuple[Hashable, ...]) -> None:
key = _resolver_key(fn)
if key in stack:
raise InvalidSignature(f"Resolver {_resolver_name(fn)!r} has a cyclic dependency")
if key in plans:
return
hints = _type_hints(fn)
sig = inspect.signature(fn)
params: dict[str, _ParamPlan] = {}
nested: list[Callable[..., Any]] = []
for param_name in sig.parameters:
annotation = hints.get(param_name)
if annotation is not None and _is_context_annotation(annotation):
params[param_name] = _ParamPlan("context")
continue
marker, wants_union = _resolve_marker(annotation)
if marker is not None:
params[param_name] = _ParamPlan("resolve", marker, wants_union)
nested.append(marker.fn)
continue
if param_name in tool_arg_names:
params[param_name] = _ParamPlan("by_name")
continue
raise InvalidSignature(
f"Resolver {_resolver_name(fn)!r} parameter {param_name!r} cannot be resolved: "
"expected a Context, an Annotated[_, Resolve(...)], or a tool argument by name"
)
plans[key] = _ResolverPlan(fn, params, is_async_callable(fn))
for dep in nested:
analyze(dep, stack + (key,))
for marker, _ in resolved_params.values():
analyze(marker.fn, ())
return plans
def _resolve_marker(annotation: Any) -> tuple[Resolve | None, bool]:
if get_origin(annotation) is not Annotated:
return None, False
type_arg, *metadata = get_args(annotation)
marker = next((m for m in metadata if isinstance(m, Resolve)), None)
return marker, (_wants_union(type_arg) if marker is not None else False)
def _is_context_annotation(annotation: Any) -> bool:
if get_origin(annotation) is Annotated:
annotation = get_args(annotation)[0]
candidates = get_args(annotation) if get_origin(annotation) is not None else (annotation,)
return any(isinstance(c, type) and issubclass(c, Context) for c in candidates)
async def resolve_arguments(
resolved_params: Mapping[str, tuple[Resolve, bool]],
plans: Mapping[Hashable, _ResolverPlan],
tool_args: Mapping[str, Any],
context: Context[Any, Any],
) -> dict[str, Any]:
"""Resolve every `Resolve`-marked tool parameter into a concrete value.
Each resolver runs at most once (memoized by function identity). Returns a
mapping of tool parameter name to the value to inject.
Raises:
ToolError: If an elicited value is declined or cancelled and the consumer
asked for the unwrapped model (rather than the result union).
"""
cache: dict[Hashable, ElicitationResult[Any]] = {}
injected: dict[str, Any] = {}
for name, (marker, wants_union) in resolved_params.items():
outcome = await _resolve(marker.fn, plans, tool_args, context, cache)
injected[name] = outcome if wants_union else _unwrap(outcome, name)
return injected
async def _resolve(
fn: Callable[..., Any],
plans: Mapping[Hashable, _ResolverPlan],
tool_args: Mapping[str, Any],
context: Context[Any, Any],
cache: dict[Hashable, ElicitationResult[Any]],
) -> ElicitationResult[Any]:
key = _resolver_key(fn)
if key in cache:
return cache[key]
plan = plans[key]
kwargs: dict[str, Any] = {}
for param_name, param_plan in plan.params.items():
if param_plan.kind == "context":
kwargs[param_name] = context
elif param_plan.kind == "by_name":
kwargs[param_name] = tool_args[param_name]
else:
assert param_plan.resolve is not None
dep_outcome = await _resolve(param_plan.resolve.fn, plans, tool_args, context, cache)
kwargs[param_name] = dep_outcome if param_plan.wants_union else _unwrap(dep_outcome, param_name)
if plan.is_async:
result = await fn(**kwargs)
else:
result = await anyio.to_thread.run_sync(lambda: fn(**kwargs))
outcome: ElicitationResult[Any]
if isinstance(result, Elicit):
elicit = cast("Elicit[BaseModel]", result)
outcome = await context.elicit(elicit.message, elicit.schema)
else:
# A resolver may return any type (not just `BaseModel`); `model_construct`
# wraps it as an accepted result without validating against the schema bound.
outcome = cast("AcceptedElicitation[Any]", AcceptedElicitation.model_construct(data=result))
cache[key] = outcome
return outcome
def _unwrap(outcome: ElicitationResult[Any], name: str) -> Any:
if isinstance(outcome, AcceptedElicitation):
return outcome.data
raise ToolError(f"Resolver for parameter {name!r} could not resolve: elicitation was {outcome.action}")
__all__ = [
"Resolve",
"Elicit",
"ElicitationResult",
"AcceptedElicitation",
"DeclinedElicitation",
"CancelledElicitation",
"find_resolved_parameters",
"build_resolver_plans",
"resolve_arguments",
]

View file

@ -1,6 +1,6 @@
from __future__ import annotations
from collections.abc import Callable
from collections.abc import Callable, Hashable
from functools import cached_property
from typing import TYPE_CHECKING, Any
@ -8,6 +8,11 @@ from mcp_types import Icon, ToolAnnotations
from pydantic import BaseModel, Field
from mcp.server.mcpserver.exceptions import ToolError
from mcp.server.mcpserver.resolve import (
build_resolver_plans,
find_resolved_parameters,
resolve_arguments,
)
from mcp.server.mcpserver.utilities.context_injection import find_context_parameter
from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, func_metadata
from mcp.shared._callable_inspection import is_async_callable
@ -32,6 +37,14 @@ class Tool(BaseModel):
)
is_async: bool = Field(description="Whether the tool is async")
context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
resolved_params: dict[str, Any] = Field(
default_factory=lambda: {},
exclude=True,
description="Parameters filled by resolvers, mapped to (Resolve, wants_union)",
)
resolver_plans: dict[Hashable, Any] = Field(
default_factory=lambda: {}, exclude=True, description="Static per-resolver parameter plans"
)
annotations: ToolAnnotations | None = Field(None, description="Optional annotations for the tool")
icons: list[Icon] | None = Field(default=None, description="Optional list of icons for this tool")
meta: dict[str, Any] | None = Field(default=None, description="Optional metadata for this tool")
@ -67,13 +80,23 @@ class Tool(BaseModel):
if context_kwarg is None: # pragma: no branch
context_kwarg = find_context_parameter(fn)
resolved_params = find_resolved_parameters(fn)
skip_names = [context_kwarg] if context_kwarg is not None else []
skip_names.extend(resolved_params)
func_arg_metadata = func_metadata(
fn,
skip_names=[context_kwarg] if context_kwarg is not None else [],
skip_names=skip_names,
structured_output=structured_output,
)
parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True)
# Match `model_dump_one_level`'s kwarg keys (alias when present, else field name)
# so a by-name resolver param resolves to a key that exists at call time.
tool_arg_names = {field.alias or name for name, field in func_arg_metadata.arg_model.model_fields.items()}
resolver_plans = build_resolver_plans(resolved_params, tool_arg_names)
return cls(
fn=fn,
name=func_name,
@ -83,6 +106,8 @@ class Tool(BaseModel):
fn_metadata=func_arg_metadata,
is_async=is_async,
context_kwarg=context_kwarg,
resolved_params=dict(resolved_params),
resolver_plans=resolver_plans,
annotations=annotations,
icons=icons,
meta=meta,
@ -100,11 +125,26 @@ class Tool(BaseModel):
ToolError: If the tool function raises during execution.
"""
try:
pass_directly: dict[str, Any] = {}
if self.context_kwarg is not None:
pass_directly[self.context_kwarg] = context
# Resolvers see the same validated arguments the tool body receives:
# validate once and reuse it, so a `default_factory`/stateful validator
# can't hand a by-name resolver a different value than the body.
pre_validated: dict[str, Any] | None = None
if self.resolved_params:
pre_validated = self.fn_metadata.validate_arguments(arguments)
pass_directly |= await resolve_arguments(
self.resolved_params, self.resolver_plans, pre_validated, context
)
result = await self.fn_metadata.call_fn_with_arg_validation(
self.fn,
self.is_async,
arguments,
{self.context_kwarg: context} if self.context_kwarg is not None else None,
pass_directly or None,
pre_validated=pre_validated,
)
if convert_result:

View file

@ -69,21 +69,36 @@ class FuncMetadata(BaseModel):
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
wrap_output: bool = False
def validate_arguments(self, arguments_to_validate: dict[str, Any]) -> dict[str, Any]:
"""Validate raw arguments into a one-level kwargs dict (no function call).
Used to feed resolver dependency injection the validated tool arguments
before the tool function itself runs.
"""
arguments_pre_parsed = self.pre_parse_json(arguments_to_validate)
arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed)
return arguments_parsed_model.model_dump_one_level()
async def call_fn_with_arg_validation(
self,
fn: Callable[..., Any | Awaitable[Any]],
fn_is_async: bool,
arguments_to_validate: dict[str, Any],
arguments_to_pass_directly: dict[str, Any] | None,
pre_validated: dict[str, Any] | None = None,
) -> Any:
"""Call the given function with arguments validated and injected.
Arguments are first attempted to be parsed from JSON, then validated against
the argument model, before being passed to the function.
the argument model, before being passed to the function. Pass `pre_validated`
(the output of `validate_arguments`) to reuse an earlier validation pass -
validating twice can re-run `default_factory`/stateful validators and hand the
function different values than a caller already observed.
"""
arguments_pre_parsed = self.pre_parse_json(arguments_to_validate)
arguments_parsed_model = self.arg_model.model_validate(arguments_pre_parsed)
arguments_parsed_dict = arguments_parsed_model.model_dump_one_level()
# Copy so a caller-provided `pre_validated` dict is never mutated in place.
arguments_parsed_dict = dict(
pre_validated if pre_validated is not None else self.validate_arguments(arguments_to_validate)
)
arguments_parsed_dict |= arguments_to_pass_directly or {}
@ -150,7 +165,7 @@ class FuncMetadata(BaseModel):
key_to_field_info[field_info.alias] = field_info
for data_key, data_value in data.items():
if data_key not in key_to_field_info: # pragma: no cover
if data_key not in key_to_field_info:
continue
field_info = key_to_field_info[data_key]