Worker-stable wire keys; restore typed model for unannotated eliciting resolvers

Review follow-ups on #2986:
- _state_key carries no id(...): it is module:qualname (a callable object uses its
  type's), so request_state round-trips and resumes on any worker (stateless HTTP).
  Two resolvers sharing that base (method instances, factory closures) are already
  disambiguated deterministically at registration (#N), so dropping the id is safe.
- An eliciting resolver whose annotation lacks an Elicit[T] arm has elicit_schema
  None; its answer restored from request_state is now re-validated against the live
  Elicit.schema (via _elicit consulting res.state) instead of injecting a raw dict.
- Move the datetime import to module top (AGENTS.md).

Add regression tests: an unannotated eliciting resolver in a multi-round flow, and
worker-stable wire keys for method instances and callable objects.

Author: Kludex

src/mcp/server/mcpserver/resolve.py

@@ -155,6 +155,13 @@ def find_resolved_parameters(fn: Callable[..., Any]) -> dict[str, tuple[Resolve,
     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)
@@ -163,6 +170,13 @@ def find_resolved_parameters(fn: Callable[..., Any]) -> dict[str, tuple[Resolve,
     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 _elicit_return_schema(return_annotation: Any) -> type[BaseModel] | None:
     """Extract `T` from a resolver return type's `Elicit[T]` arm, if present.
 
@@ -187,12 +201,13 @@ def _is_union(annotation: Any) -> bool:
 def _wants_union(type_arg: Any) -> bool:
     """True when `type_arg` is an `ElicitationResult` member (or a union of them).
 
-    Handles the bare `ElicitationResult[T]` alias (a `TypeAliasType` carrying the
-    union on `__value__`), an explicit `AcceptedElicitation[T] | ... ` union, and a
-    single member.
+    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.
     """
-    origin = get_origin(type_arg)
-    value = getattr(origin, "__value__", None)
+    # 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,)
@@ -330,16 +345,17 @@ def __init__(
 
 
 def _state_key(fn: Callable[..., Any]) -> str:
-    """Process-stable wire key for a resolver's elicitation.
+    """Worker-stable base wire key for a resolver, derived only from registration data.
 
-    `id(fn)` isn't stable across `input_required` rounds, so key `input_requests` /
-    `request_state` by `module:qualname`. Bound methods add their `__self__` id so
-    two instances of the same method get distinct questions and stored outcomes
-    (the registered `Resolve(...)` holds the instance for the call's lifetime).
+    `input_requests`/`request_state` must round-trip through the client and resume on
+    any worker (stateless HTTP), so the key carries no `id(...)`: it is the resolver's
+    `module:qualname` (a callable object uses its type's). Distinct resolvers that
+    share this base - two instances of one method, two closures from one factory - are
+    disambiguated deterministically by `build_resolver_plans` (`base`, `base#1`, ...).
     """
-    base = f"{getattr(fn, '__module__', '')}:{getattr(fn, '__qualname__', fn)!s}"
-    bound_self = getattr(fn, "__self__", None)
-    return f"{base}#{id(bound_self)}" if bound_self is not None else base
+    qualname = getattr(fn, "__qualname__", None) or type(fn).__qualname__
+    module = getattr(fn, "__module__", None) or type(fn).__module__
+    return f"{module}:{qualname}"
 
 
 async def resolve_arguments(
@@ -396,7 +412,11 @@ async def _resolve(fn: Callable[..., Any], res: _Resolution) -> ElicitationResul
     if wire_key in res.pending:
         # Already asked this round by another consumer; don't run the resolver again.
         raise _Pending
-    if wire_key in res.state:
+    # Restore a prior round's outcome directly only when its model is known from the
+    # `Elicit[T]` return arm. Without that (a resolver that elicits but isn't annotated
+    # `-> ... Elicit[T]`), fall through and re-run the resolver so `_elicit` can
+    # re-validate the stored answer against the live `Elicit.schema`.
+    if wire_key in res.state and (plan.elicit_schema is not None or res.state[wire_key].action != "accept"):
         outcome = _outcome_from_state(res.state[wire_key], plan.elicit_schema)
         res.cache[cache_key] = outcome
         # Carry the restored answer forward: if a later resolver is still pending,
@@ -449,6 +469,13 @@ async def _elicit(elicit: Elicit[Any], key: str, res: _Resolution) -> Elicitatio
     if not res.input_required:
         return await res.context.elicit(elicit.message, elicit.schema)
 
+    # Answered in a prior round (restored without a known schema, e.g. an unannotated
+    # resolver): re-validate the stored entry against the live `Elicit.schema`.
+    if key in res.state and key not in res.answers:
+        outcome = _outcome_from_state(res.state[key], elicit.schema)
+        res.elicited[key] = outcome
+        return outcome
+
     answer = res.answers.get(key)
     if answer is None:
         res.pending[key] = _elicit_request(elicit)

tests/server/mcpserver/test_resolve.py

@@ -1,7 +1,8 @@
 """Tests for resolver dependency injection (MRTR) on MCPServer tools."""
 
 from collections.abc import Callable
-from typing import Annotated, Literal
+from datetime import datetime
+from typing import Annotated, Any, Literal
 
 import pytest
 from mcp_types import (
@@ -340,6 +341,32 @@ async def tool(login: Annotated[Login, Resolve(login)]) -> str:
         Tool.from_function(tool)
 
 
+def test_resolve_marker_inside_a_union_raises_at_registration():
+    async def login(ctx: Context) -> Login:
+        return Login(username="x")  # pragma: no cover
+
+    async def tool(login: Annotated[Login, Resolve(login)] | None = None) -> str:
+        return login.username if login else ""  # pragma: no cover
+
+    with pytest.raises(InvalidSignature, match="wraps `Resolve"):
+        Tool.from_function(tool)
+
+
+def test_bare_elicitation_result_alias_wants_the_outcome_union():
+    # The bare `ElicitationResult` alias (no `[T]` subscription) must still opt into
+    # the result union, not be treated as wanting the unwrapped model.
+    async def login(ctx: Context) -> Login:
+        return Login(username="x")  # pragma: no cover
+
+    async def tool(login: object) -> str:
+        return "x"  # pragma: no cover
+
+    bare_alias: Any = ElicitationResult
+    tool.__annotations__["login"] = Annotated[bare_alias, Resolve(login)]
+    (_, wants_union) = find_resolved_parameters(tool)["login"]
+    assert wants_union is True
+
+
 def test_resolve_marker_on_return_annotation_is_ignored():
     async def login(ctx: Context) -> Login:
         return Login(username="x")  # pragma: no cover
@@ -853,8 +880,6 @@ async def both(
 
 @pytest.mark.anyio
 async def test_non_serializable_sibling_resolver_does_not_break_rounds():
-    from datetime import datetime
-
     mcp = MCPServer(name="NonSerializable")
 
     async def clock(ctx: Context) -> datetime:
@@ -1052,3 +1077,58 @@ def answer(key: str, params: ElicitRequestFormParams) -> ElicitResult:
         result = await _drive_mrtr(client, "tool", {}, answer)
         assert isinstance(result.content[0], TextContent)
         assert result.content[0].text == "A,B"
+
+
+@pytest.mark.anyio
+async def test_eliciting_resolver_without_elicit_arm_restores_a_typed_model():
+    # A resolver annotated `-> Login` that actually returns `Elicit(...)` has no
+    # `Elicit[T]` return arm, so `elicit_schema` is None. Its answer, restored from
+    # request_state in a 3+ round flow, must still come back as a Login model (not a
+    # raw dict) so a dependent resolver/tool can use its attributes.
+    mcp = MCPServer(name="LyingAnnotation")
+
+    # Annotated without an `Elicit[T]` return arm, so `elicit_schema` is None.
+    async def login(ctx: Context) -> object:
+        return Elicit("user?", Login)
+
+    async def confirm(login: Annotated[Login, Resolve(login)]) -> Elicit[Confirm]:
+        return Elicit(f"as {login.username}?", Confirm)
+
+    @mcp.tool()
+    async def act(
+        login: Annotated[Login, Resolve(login)],
+        confirm: Annotated[Confirm, Resolve(confirm)],
+    ) -> str:
+        return f"{login.username}:{confirm.ok}"
+
+    def answer(key: str, params: ElicitRequestFormParams) -> ElicitResult:
+        if "user" in params.message:
+            return ElicitResult(action="accept", content={"username": "octocat"})
+        assert "as octocat?" in params.message  # login restored as a real model
+        return ElicitResult(action="accept", content={"ok": True})
+
+    async with Client(mcp) as client:
+        result = await _drive_mrtr(client, "act", {}, answer)
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == "octocat:True"
+
+
+def test_wire_key_is_worker_stable_for_methods_and_callable_objects():
+    from mcp.server.mcpserver.resolve import _state_key
+
+    class Service:
+        async def token(self, ctx: Context) -> Login:
+            return Login(username="x")  # pragma: no cover
+
+    class CallableResolver:
+        async def __call__(self, ctx: Context) -> Login:
+            return Login(username="x")  # pragma: no cover
+
+    a, b = Service(), Service()
+    # No id(...) in the key: two instances of one method get the same base (they are
+    # disambiguated at registration, not here), and the key carries no memory address.
+    assert _state_key(a.token) == _state_key(b.token)
+    assert "#" not in _state_key(a.token)
+    assert _state_key(a.token).endswith("Service.token")
+    # Callable objects key by their type's qualname (they have no `__qualname__`).
+    assert _state_key(CallableResolver()).endswith("CallableResolver")