Add resolver dependency injection for MCPServer tools (#2969)

Co-authored-by: Max Isbey <224885523+maxisbey@users.noreply.github.com>

tests/server/mcpserver/test_func_metadata.py

@@ -155,6 +155,28 @@ async def test_complex_function_runtime_arg_validation_with_json():
     assert result == "ok!"
 
 
+@pytest.mark.anyio
+async def test_call_fn_does_not_mutate_pre_validated():
+    """A caller-provided `pre_validated` dict must not be mutated by the call."""
+
+    def fn(x: int, ctx: str) -> str:
+        return f"{x}:{ctx}"
+
+    meta = func_metadata(fn, skip_names=["ctx"])
+    pre_validated = meta.validate_arguments({"x": 1})
+    snapshot = dict(pre_validated)
+
+    result = await meta.call_fn_with_arg_validation(
+        fn,
+        fn_is_async=False,
+        arguments_to_validate={"x": 1},
+        arguments_to_pass_directly={"ctx": "injected"},
+        pre_validated=pre_validated,
+    )
+    assert result == "1:injected"
+    assert pre_validated == snapshot  # `ctx` was not leaked into the caller's dict
+
+
 def test_str_vs_list_str():
     """Test handling of string vs list[str] type annotations.
 

tests/server/mcpserver/test_resolve.py

@@ -0,0 +1,571 @@
+"""Tests for resolver dependency injection (MRTR) on MCPServer tools."""
+
+from typing import Annotated, Any, Literal
+
+import pytest
+from mcp_types import ElicitRequestParams, ElicitResult, TextContent
+from pydantic import BaseModel, Field
+
+from mcp import Client
+from mcp.client import ClientRequestContext
+from mcp.server.mcpserver import (
+    AcceptedElicitation,
+    CancelledElicitation,
+    Context,
+    DeclinedElicitation,
+    Elicit,
+    ElicitationResult,
+    MCPServer,
+    Resolve,
+)
+from mcp.server.mcpserver.exceptions import InvalidSignature
+from mcp.server.mcpserver.resolve import _resolver_key, find_resolved_parameters
+from mcp.server.mcpserver.tools.base import Tool
+
+
+class Login(BaseModel):
+    username: str
+
+
+class Confirm(BaseModel):
+    ok: bool
+
+
+async def _alias_login(ctx: Context) -> Login:
+    return Login(username="x")  # pragma: no cover - only the signature is inspected
+
+
+def _accept(content: dict[str, str | int | float | bool | list[str] | None]):
+    async def callback(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
+        return ElicitResult(action="accept", content=content)
+
+    return callback
+
+
+async def _decline(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
+    return ElicitResult(action="decline")
+
+
+async def _text(client: Client, tool: str, args: dict[str, object]) -> str:
+    result = await client.call_tool(tool, args)
+    assert len(result.content) == 1
+    assert isinstance(result.content[0], TextContent)
+    return result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_resolver_returns_value_directly_without_eliciting():
+    mcp = MCPServer(name="Direct")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        username = (ctx.headers or {}).get("x-github-user")
+        if username:  # pragma: no cover - no headers on in-memory transport
+            return Login(username=username)
+        return Login(username="from-resolver")
+
+    @mcp.tool()
+    async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
+        return login.username
+
+    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, "whoami", {}) == "from-resolver"
+
+
+@pytest.mark.anyio
+async def test_resolver_elicits_and_injects_unwrapped_model_on_accept():
+    mcp = MCPServer(name="Accept")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    @mcp.tool()
+    async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
+        return login.username
+
+    async with Client(mcp, mode="legacy", elicitation_callback=_accept({"username": "octocat"})) as client:
+        assert await _text(client, "whoami", {}) == "octocat"
+
+
+@pytest.mark.anyio
+async def test_consumer_receives_result_union_and_branches():
+    mcp = MCPServer(name="Union")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    @mcp.tool()
+    async def whoami(login: Annotated[ElicitationResult[Login], Resolve(login)]) -> str:
+        match login:
+            case AcceptedElicitation(data=data):
+                return f"hi {data.username}"
+            case _:  # pragma: no cover - accepted in this test
+                return "no username"
+
+    async with Client(mcp, mode="legacy", elicitation_callback=_accept({"username": "octocat"})) as client:
+        assert await _text(client, "whoami", {}) == "hi octocat"
+
+
+@pytest.mark.anyio
+async def test_decline_reaches_union_consumer_without_aborting():
+    mcp = MCPServer(name="UnionDecline")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    @mcp.tool()
+    async def whoami(
+        login: Annotated[AcceptedElicitation[Login] | DeclinedElicitation | CancelledElicitation, Resolve(login)],
+    ) -> str:
+        if isinstance(login, DeclinedElicitation):
+            return "declined gracefully"
+        raise NotImplementedError
+
+    async with Client(mcp, mode="legacy", elicitation_callback=_decline) as client:
+        assert await _text(client, "whoami", {}) == "declined gracefully"
+
+
+@pytest.mark.anyio
+async def test_decline_aborts_when_consumer_wants_unwrapped():
+    mcp = MCPServer(name="UnwrappedDecline")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    @mcp.tool()
+    async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
+        raise NotImplementedError  # pragma: no cover - never reached
+
+    async with Client(mcp, mode="legacy", elicitation_callback=_decline) as client:
+        result = await client.call_tool("whoami", {})
+        assert result.is_error
+        assert isinstance(result.content[0], TextContent)
+        assert "decline" in result.content[0].text
+
+
+@pytest.mark.anyio
+async def test_nested_resolver_sees_dependency_and_tool_args():
+    mcp = MCPServer(name="Nested")
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    async def confirm(repo: str, login: Annotated[Login, Resolve(login)]) -> Elicit[Confirm]:
+        return Elicit(f"Star {repo} as {login.username}?", Confirm)
+
+    @mcp.tool()
+    async def star_repo(
+        repo: str,
+        login: Annotated[Login, Resolve(login)],
+        confirm: Annotated[Confirm, Resolve(confirm)],
+    ) -> str:
+        if confirm.ok:
+            return f"starred {repo} as {login.username}"
+        raise NotImplementedError
+
+    async def callback(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
+        if "username" in params.message:
+            return ElicitResult(action="accept", content={"username": "octocat"})
+        assert "Star modelcontextprotocol/python-sdk as octocat?" in params.message
+        return ElicitResult(action="accept", content={"ok": True})
+
+    async with Client(mcp, mode="legacy", elicitation_callback=callback) as client:
+        text = await _text(client, "star_repo", {"repo": "modelcontextprotocol/python-sdk"})
+        assert text == "starred modelcontextprotocol/python-sdk as octocat"
+
+
+@pytest.mark.anyio
+async def test_resolver_runs_once_for_two_consumers():
+    mcp = MCPServer(name="ExactlyOnce")
+    elicit_count = 0
+
+    async def login(ctx: Context) -> Login | Elicit[Login]:
+        return Elicit("GitHub username?", Login)
+
+    async def confirm(login: Annotated[Login, Resolve(login)]) -> Elicit[Confirm]:
+        return Elicit(f"As {login.username}?", Confirm)
+
+    @mcp.tool()
+    async def star_repo(
+        login: Annotated[Login, Resolve(login)],
+        confirm: Annotated[Confirm, Resolve(confirm)],
+    ) -> str:
+        return f"{login.username}:{confirm.ok}"
+
+    async def callback(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
+        nonlocal elicit_count
+        if "username" in params.message:
+            elicit_count += 1
+            return ElicitResult(action="accept", content={"username": "octocat"})
+        return ElicitResult(action="accept", content={"ok": True})
+
+    async with Client(mcp, mode="legacy", elicitation_callback=callback) as client:
+        assert await _text(client, "star_repo", {}) == "octocat:True"
+    assert elicit_count == 1
+
+
+@pytest.mark.anyio
+async def test_sync_resolver():
+    mcp = MCPServer(name="Sync")
+
+    def login(ctx: Context) -> Login:
+        return Login(username="sync-user")
+
+    @mcp.tool()
+    async def whoami(login: Annotated[Login, Resolve(login)]) -> str:
+        return login.username
+
+    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, "whoami", {}) == "sync-user"
+
+
+def test_resolved_params_absent_from_input_schema():
+    async def login(ctx: Context) -> Login:
+        return Login(username="x")  # pragma: no cover - only the schema is inspected
+
+    async def tool(
+        repo: Annotated[str, Field(description="repo name")],
+        login: Annotated[Login, Resolve(login)],
+    ) -> str:
+        return repo  # pragma: no cover - only the schema is inspected
+
+    built = Tool.from_function(tool)
+    properties = built.parameters["properties"]
+    assert "repo" in properties
+    assert "login" not in properties
+
+
+def test_cycle_detection_raises_at_registration():
+    async def a(dep: Login) -> Login:
+        return dep  # pragma: no cover
+
+    async def b(dep: Login) -> Login:
+        return dep  # pragma: no cover
+
+    # Close the loop after both exist: a depends on b, b depends on a.
+    a.__annotations__["dep"] = Annotated[Login, Resolve(b)]
+    b.__annotations__["dep"] = Annotated[Login, Resolve(a)]
+
+    async def tool(value: Annotated[Login, Resolve(a)]) -> str:
+        return value.username  # pragma: no cover
+
+    with pytest.raises(InvalidSignature, match="cyclic"):
+        Tool.from_function(tool)
+
+
+def test_find_resolved_parameters_tolerates_unresolvable_hints():
+    def fn(x: int) -> int:
+        return x  # pragma: no cover
+
+    fn.__annotations__["x"] = "DoesNotExist"
+    assert find_resolved_parameters(fn) == {}
+
+
+def test_elicitation_result_alias_resolves_under_postponed_annotations():
+    # Reproduces the case where `from __future__ import annotations` stringifies
+    # `Annotated[ElicitationResult[Login], Resolve(_alias_login)]`: the alias must be
+    # subscriptable so the resolver is detected (not silently dropped) and the
+    # consumer is recognized as wanting the result union.
+    def tool(login: str) -> str:
+        return login  # pragma: no cover
+
+    tool.__annotations__["login"] = "Annotated[ElicitationResult[Login], Resolve(_alias_login)]"
+    resolved = find_resolved_parameters(tool)
+    assert "login" in resolved
+    assert resolved["login"][1] is True  # wants_union
+
+
+def test_unresolvable_resolver_param_raises_at_registration():
+    async def login(mystery: int) -> Login:
+        return Login(username="x")  # pragma: no cover
+
+    async def tool(login: Annotated[Login, Resolve(login)]) -> str:
+        return login.username  # pragma: no cover
+
+    with pytest.raises(InvalidSignature, match="cannot be resolved"):
+        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
+
+    async def tool(repo: str) -> Annotated[str, Resolve(login)]:
+        return repo  # pragma: no cover
+
+    assert find_resolved_parameters(tool) == {}
+
+
+def test_callable_object_resolver_error_uses_type_name():
+    class BadResolver:
+        async def __call__(self, mystery: int) -> Login:
+            return Login(username="x")  # pragma: no cover
+
+    async def tool(login: Annotated[Login, Resolve(BadResolver())]) -> str:
+        return login.username  # pragma: no cover
+
+    with pytest.raises(InvalidSignature, match="'BadResolver'"):
+        Tool.from_function(tool)
+
+
+@pytest.mark.anyio
+async def test_by_name_resolver_param_uses_aliased_tool_arg():
+    mcp = MCPServer(name="Aliased")
+
+    # `schema` collides with a BaseModel attribute, so func_metadata aliases the field;
+    # the runtime kwarg key is the alias, which is what a by-name resolver must match.
+    async def upper(schema: str) -> Login:
+        return Login(username=schema.upper())
+
+    @mcp.tool()
+    async def run(schema: str, shouted: Annotated[Login, Resolve(upper)]) -> str:
+        return shouted.username
+
+    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", {"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)
+
+
+@pytest.mark.anyio
+async def test_resolver_and_body_see_the_same_validated_default():
+    mcp = MCPServer(name="DefaultFactory")
+    counter = {"n": 0}
+
+    def next_id() -> int:
+        counter["n"] += 1
+        return counter["n"]
+
+    # A by-name resolver and the tool body must observe one validation pass, so the
+    # `default_factory` runs once and both see the same generated value.
+    async def echo_id(request_id: int) -> int:
+        return request_id
+
+    @mcp.tool()
+    async def run(
+        request_id: Annotated[int, Field(default_factory=next_id)],
+        resolved_id: Annotated[int, Resolve(echo_id)],
+    ) -> str:
+        return f"{request_id}:{resolved_id}"
+
+    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", {}) == "1:1"
+    assert counter["n"] == 1
+
+
+def test_resolver_key_is_stable_for_methods_and_distinct_callables():
+    class Service:
+        def handler(self) -> None: ...  # pragma: no cover
+
+    a, b = Service(), Service()
+
+    # Pure-python bound methods: stable across accesses, distinct per instance.
+    assert _resolver_key(a.handler) == _resolver_key(a.handler)
+    assert _resolver_key(a.handler) != _resolver_key(b.handler)
+
+    # Built-in bound methods (no `__func__`): fresh object each access, but the key
+    # is stable and keyed to `__self__`.
+    items: list[int] = []
+    others: list[int] = []
+    assert _resolver_key(items.append) == _resolver_key(items.append)
+    assert _resolver_key(items.append) != _resolver_key(others.append)
+    assert _resolver_key(items.append) != _resolver_key(items.pop)
+
+    # Plain functions key by identity.
+    def fn() -> None: ...  # pragma: no cover
+
+    assert _resolver_key(fn) == _resolver_key(fn)
+
+
+def _delete_folder_server() -> tuple[MCPServer, dict[str, list[str]]]:
+    """The `delete_folder` example from docs/migration.md, wired to an in-memory fs."""
+    mcp = MCPServer(name="files")
+    fs: dict[str, list[str]] = {}
+
+    async def confirm_delete(path: str) -> Confirm | Elicit[Confirm]:
+        file_count = len(fs.get(path, []))
+        if file_count == 0:
+            return Confirm(ok=True)
+        return Elicit(f"{path} has {file_count} file(s). Delete anyway?", Confirm)
+
+    @mcp.tool()
+    async def delete_folder(
+        path: str,
+        confirm: Annotated[ElicitationResult[Confirm], Resolve(confirm_delete)],
+    ) -> str:
+        match confirm:
+            case AcceptedElicitation(data=Confirm(ok=True)):
+                fs.pop(path, None)
+                return f"deleted {path}"
+            case AcceptedElicitation():
+                return "kept the folder"
+            case DeclinedElicitation():
+                return "declined: folder not deleted"
+            case CancelledElicitation():  # pragma: no branch
+                return "cancelled: folder not deleted"
+
+    return mcp, fs
+
+
+@pytest.mark.anyio
+async def test_delete_empty_folder_does_not_elicit():
+    mcp, fs = _delete_folder_server()
+    fs["/empty"] = []
+
+    async def never(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:  # pragma: no cover
+        raise AssertionError("should not elicit for an empty folder")
+
+    async with Client(mcp, mode="legacy", elicitation_callback=never) as client:
+        assert await _text(client, "delete_folder", {"path": "/empty"}) == "deleted /empty"
+    assert "/empty" not in fs
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize(
+    ("action", "content", "expected"),
+    [
+        ("accept", {"ok": True}, "deleted /docs"),
+        ("accept", {"ok": False}, "kept the folder"),
+        ("decline", None, "declined: folder not deleted"),
+        ("cancel", None, "cancelled: folder not deleted"),
+    ],
+)
+async def test_delete_non_empty_folder_handles_every_outcome(
+    action: Literal["accept", "decline", "cancel"],
+    content: dict[str, str | int | float | bool | list[str] | None] | None,
+    expected: str,
+):
+    mcp, fs = _delete_folder_server()
+    fs["/docs"] = ["a.txt", "b.txt"]
+
+    async def callback(context: ClientRequestContext, params: ElicitRequestParams) -> ElicitResult:
+        assert "/docs has 2 file(s)" in params.message
+        return ElicitResult(action=action, content=content)
+
+    async with Client(mcp, mode="legacy", elicitation_callback=callback) as client:
+        assert await _text(client, "delete_folder", {"path": "/docs"}) == expected
+    assert ("/docs" in fs) is (expected != "deleted /docs")

tests/server/mcpserver/test_server.py

@@ -1,5 +1,6 @@
 import base64
 from pathlib import Path
+from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -1801,6 +1802,33 @@ async def test_report_progress_delegates_to_session_report_progress():
     mock_session.report_progress.assert_awaited_once_with(50, 100, "halfway")
 
 
+def _request_context(request: object | None) -> ServerRequestContext[None, object]:
+    return ServerRequestContext(
+        session=AsyncMock(),
+        method="tools/call",
+        lifespan_context=None,
+        protocol_version="2025-11-25",
+        request=request,
+    )
+
+
+def test_context_headers_returns_request_headers():
+    request = SimpleNamespace(headers={"x-github-user": "octocat"})
+    ctx = Context(request_context=_request_context(request), mcp_server=MagicMock())
+    assert ctx.headers == {"x-github-user": "octocat"}
+
+
+def test_context_headers_is_none_without_request():
+    ctx = Context(request_context=_request_context(None), mcp_server=MagicMock())
+    assert ctx.headers is None
+
+
+def test_context_headers_is_none_when_request_carries_no_headers():
+    """A transport may attach a custom request object that has no headers attribute."""
+    ctx = Context(request_context=_request_context(object()), mcp_server=MagicMock())
+    assert ctx.headers is None
+
+
 async def test_read_resource_template_error():
     """Template-creation failure must surface as INTERNAL_ERROR, not INVALID_PARAMS (not-found)."""
     mcp = MCPServer()