Fix resolver edge cases: non-BaseModel returns, optional Context, bound-method memoization

bughunter findings on #2969:
- Resolvers may return any type, not just BaseModel. Wrapping the return in
  AcceptedElicitation(data=...) validated it against the schema bound, so e.g.
  Annotated[str, Resolve(get_token)] failed every call with a cryptic
  ValidationError. Use model_construct to wrap the value without validation
  (the Elicit[T] path still validates via ctx.elicit).
- _is_context_annotation now unwraps unions, so a resolver param typed
  Context | None is accepted, matching find_context_parameter on tools.
- Memoize resolvers by the callable itself (hash/eq) instead of id(fn), so a
  bound-method resolver referenced as auth.login in two places runs at most
  once and participates in cycle detection. Fresh bound-method objects share
  identity by (__func__, __self__).

Add regression tests for each.

Author: Kludex

src/mcp/server/mcpserver/resolve.py

@@ -21,7 +21,7 @@
 
 import inspect
 import typing
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Hashable, Mapping
 from typing import Annotated, Any, Generic, cast, get_args, get_origin
 
 import anyio.to_thread
@@ -131,21 +131,35 @@ def _wants_union(type_arg: Any) -> bool:
     return any(isinstance(m, type) and issubclass(m, _ELICITATION_RESULT_MEMBERS) for m in members)
 
 
+def _resolver_key(fn: Callable[..., Any]) -> Hashable:
+    """Stable, equality-based key for memoizing a resolver.
+
+    Bound methods are recreated on each attribute access (`id(auth.login)` differs
+    every time) but hash/compare by `(__func__, __self__)`, so the callable itself
+    is the right key. Falls back to `id` only for the rare unhashable callable.
+    """
+    try:
+        hash(fn)
+    except TypeError:  # pragma: no cover - unhashable callables are pathological
+        return id(fn)
+    return fn
+
+
 def build_resolver_plans(
     resolved_params: Mapping[str, tuple[Resolve, bool]],
     tool_arg_names: set[str],
-) -> dict[int, _ResolverPlan]:
+) -> 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[int, _ResolverPlan] = {}
+    plans: dict[Hashable, _ResolverPlan] = {}
 
-    def analyze(fn: Callable[..., Any], stack: tuple[int, ...]) -> None:
-        key = id(fn)
+    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:
@@ -193,12 +207,13 @@ def _resolve_marker(annotation: Any) -> tuple[Resolve | None, bool]:
 def _is_context_annotation(annotation: Any) -> bool:
     if get_origin(annotation) is Annotated:
         annotation = get_args(annotation)[0]
-    return isinstance(annotation, type) and issubclass(annotation, Context)
+    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[int, _ResolverPlan],
+    plans: Mapping[Hashable, _ResolverPlan],
     tool_args: Mapping[str, Any],
     context: Context[Any, Any],
 ) -> dict[str, Any]:
@@ -211,7 +226,7 @@ async def resolve_arguments(
         ToolError: If an elicited value is declined or cancelled and the consumer
             asked for the unwrapped model (rather than the result union).
     """
-    cache: dict[int, ElicitationResult[BaseModel]] = {}
+    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)
@@ -221,12 +236,12 @@ async def resolve_arguments(
 
 async def _resolve(
     fn: Callable[..., Any],
-    plans: Mapping[int, _ResolverPlan],
+    plans: Mapping[Hashable, _ResolverPlan],
     tool_args: Mapping[str, Any],
     context: Context[Any, Any],
-    cache: dict[int, ElicitationResult[BaseModel]],
-) -> ElicitationResult[BaseModel]:
-    key = id(fn)
+    cache: dict[Hashable, ElicitationResult[Any]],
+) -> ElicitationResult[Any]:
+    key = _resolver_key(fn)
     if key in cache:
         return cache[key]
 
@@ -247,18 +262,20 @@ async def _resolve(
     else:
         result = await anyio.to_thread.run_sync(lambda: fn(**kwargs))
 
-    outcome: ElicitationResult[BaseModel]
+    outcome: ElicitationResult[Any]
     if isinstance(result, Elicit):
         elicit = cast("Elicit[BaseModel]", result)
         outcome = await context.elicit(elicit.message, elicit.schema)
     else:
-        outcome = AcceptedElicitation(data=result)
+        # 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[BaseModel], name: str) -> BaseModel:
+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}")

src/mcp/server/mcpserver/tools/base.py

@@ -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
 
@@ -42,7 +42,7 @@ class Tool(BaseModel):
         exclude=True,
         description="Parameters filled by resolvers, mapped to (Resolve, wants_union)",
     )
-    resolver_plans: dict[int, Any] = Field(
+    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")

tests/server/mcpserver/test_resolve.py

@@ -314,3 +314,92 @@ async def never(context: ClientRequestContext, params: ElicitRequestParams) -> E
 
     async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
         assert await _text(client, "run", {"schema": "gpt"}) == "GPT"
+
+
+@pytest.mark.anyio
+async def test_resolver_may_return_non_basemodel_value():
+    mcp = MCPServer(name="NonModel")
+
+    async def get_token(ctx: Context) -> str:
+        return "secret-token"
+
+    @mcp.tool()
+    async def use_token(token: Annotated[str, Resolve(get_token)]) -> str:
+        return token
+
+    async def never(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:  # pragma: no cover
+        raise AssertionError("should not elicit")
+
+    async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
+        assert await _text(client, "use_token", {}) == "secret-token"
+
+
+@pytest.mark.anyio
+async def test_resolver_accepts_optional_context_annotation():
+    mcp = MCPServer(name="OptionalContext")
+
+    async def whoami(ctx: Context | None) -> str:
+        assert ctx is not None
+        return "has-context"
+
+    @mcp.tool()
+    async def run(who: Annotated[str, Resolve(whoami)]) -> str:
+        return who
+
+    async def never(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:  # pragma: no cover
+        raise AssertionError("should not elicit")
+
+    async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
+        assert await _text(client, "run", {}) == "has-context"
+
+
+@pytest.mark.anyio
+async def test_bound_method_resolver_runs_once_across_references():
+    mcp = MCPServer(name="BoundMethod")
+    calls = 0
+
+    class Service:
+        async def token(self, ctx: Context) -> str:
+            nonlocal calls
+            calls += 1
+            return "tok"
+
+    service = Service()
+
+    # Each `service.token` access is a fresh bound-method object; keying by the
+    # callable (not id) keeps the resolver memoized to a single call.
+    async def downstream(token: Annotated[str, Resolve(service.token)]) -> str:
+        return token.upper()
+
+    @mcp.tool()
+    async def run(
+        token: Annotated[str, Resolve(service.token)],
+        shouted: Annotated[str, Resolve(downstream)],
+    ) -> str:
+        return f"{token}:{shouted}"
+
+    async def never(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:  # pragma: no cover
+        raise AssertionError("should not elicit")
+
+    async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
+        assert await _text(client, "run", {}) == "tok:TOK"
+    assert calls == 1
+
+
+def test_bound_method_cycle_is_detected():
+    class Service:
+        async def a(self, dep: Login) -> Login:
+            return dep  # pragma: no cover
+
+        async def b(self, dep: Login) -> Login:
+            return dep  # pragma: no cover
+
+    service = Service()
+    service.a.__func__.__annotations__["dep"] = Annotated[Login, Resolve(service.b)]
+    service.b.__func__.__annotations__["dep"] = Annotated[Login, Resolve(service.a)]
+
+    async def tool(value: Annotated[Login, Resolve(service.a)]) -> str:
+        return value.username  # pragma: no cover
+
+    with pytest.raises(InvalidSignature, match="cyclic"):
+        Tool.from_function(tool)