diff --git a/tests/tools/test_mcp_context_template.py b/tests/tools/test_mcp_context_template.py index cc44ea487b7b..0ba3d0a3a966 100644 --- a/tests/tools/test_mcp_context_template.py +++ b/tests/tools/test_mcp_context_template.py @@ -21,6 +21,8 @@ register_session_context_var, ) from tools.mcp_tool import ( + MCPServerTask, + _apply_resolved_headers, _has_context_template, _resolve_context_templates, _resolve_frozen_headers, @@ -250,7 +252,132 @@ def test_non_string_values_treated_as_static(self): # --------------------------------------------------------------------------- -# Integration test: the actual event-hook injection pattern using httpx +# Z2O-1694: caller-side resolution + instance-attr bridge to the MCP loop +# --------------------------------------------------------------------------- +# +# The real MCP request runs on a dedicated background event loop (and a task +# spawned by the SDK's post_writer), so resolving ``${context:}`` inside the +# httpx request hook *there* always saw an empty context — the per-turn +# session vars are set on the caller's thread, never the MCP loop's. The fix +# resolves on the caller side and bridges the values across via an instance +# attribute (``MCPServerTask._outbound_resolved_headers``), applied under +# ``_rpc_lock`` so per-server calls can't race. These tests pin that contract. + +class TestResolveTemplatedHeadersMethod: + """Caller-side resolution: ``MCPServerTask._resolve_templated_headers`` + snapshots the current task's session context into concrete values.""" + + def test_resolves_configured_templates(self, custom_var): + server = MCPServerTask("meridian") + server._templated_headers = { + "X-Meridian-Delegated-Principal": "${context:TEST_PRINCIPAL}", + } + custom_var.set("thomas@verdigris.co") + assert server._resolve_templated_headers() == { + "X-Meridian-Delegated-Principal": "thomas@verdigris.co", + } + + def test_unset_value_omitted_from_dict(self, custom_var): + server = MCPServerTask("meridian") + server._templated_headers = { + "X-Meridian-Delegated-Principal": "${context:TEST_PRINCIPAL}", + } + # custom_var unset → "" → dropped (not present with empty value). + assert server._resolve_templated_headers() == {} + + def test_no_templates_returns_empty(self): + server = MCPServerTask("meridian") + assert server._resolve_templated_headers() == {} + + +class TestApplyResolvedHeaders: + """Loop-side application: ``_apply_resolved_headers`` stamps the already- + resolved values onto the outbound request (no context lookup here).""" + + def test_injects_when_same_origin(self): + req = httpx.Request("POST", "http://meridian.invalid/mcp") + _apply_resolved_headers( + req, same_origin=True, + names=["X-Meridian-Delegated-Principal"], + resolved={"X-Meridian-Delegated-Principal": "thomas@verdigris.co"}, + ) + assert req.headers["X-Meridian-Delegated-Principal"] == "thomas@verdigris.co" + + def test_drops_when_value_missing(self): + req = httpx.Request("POST", "http://meridian.invalid/mcp") + req.headers["X-Meridian-Delegated-Principal"] = "stale@example.com" + _apply_resolved_headers( + req, same_origin=True, + names=["X-Meridian-Delegated-Principal"], + resolved={}, # nothing resolved this call → header must be removed + ) + assert "X-Meridian-Delegated-Principal" not in req.headers + + def test_strips_on_cross_origin(self): + req = httpx.Request("POST", "http://attacker.invalid/harvest") + req.headers["X-Meridian-Delegated-Principal"] = "thomas@verdigris.co" + _apply_resolved_headers( + req, same_origin=False, + names=["X-Meridian-Delegated-Principal"], + resolved={"X-Meridian-Delegated-Principal": "thomas@verdigris.co"}, + ) + assert "X-Meridian-Delegated-Principal" not in req.headers + + +def test_resolved_headers_survive_lost_caller_context(custom_var): + """The crux of Z2O-1694: resolve while the per-turn var is set, then + apply *after* that context is gone (simulating the MCP background loop, + whose task never saw the var). The header must still be injected — + proving the value travels via the instance attr, not a live context + lookup at request time.""" + server = MCPServerTask("meridian") + server._templated_headers = { + "X-Meridian-Delegated-Principal": "${context:TEST_PRINCIPAL}", + } + custom_var.set("thomas@verdigris.co") + resolved = server._resolve_templated_headers() # caller thread/context + server._outbound_resolved_headers = resolved # bridged onto instance + custom_var.set(_UNSET) # caller context gone + + req = httpx.Request("POST", "http://meridian.invalid/mcp") + _apply_resolved_headers( + req, same_origin=True, + names=server._templated_headers, + resolved=server._outbound_resolved_headers, + ) + assert req.headers["X-Meridian-Delegated-Principal"] == "thomas@verdigris.co" + + +@pytest.mark.asyncio +async def test_rpc_scopes_headers_to_call_and_clears(): + """``_rpc`` exposes the resolved headers only for the duration of the call + and clears them after — so a later system RPC (keepalive) passing no + headers can't leak the prior user call's delegated principal (Z2O-1694).""" + server = MCPServerTask("meridian") + assert server._outbound_resolved_headers == {} + + async with server._rpc({"X-Meridian-Delegated-Principal": "thomas@verdigris.co"}): + assert server._outbound_resolved_headers == { + "X-Meridian-Delegated-Principal": "thomas@verdigris.co", + } + assert server._outbound_resolved_headers == {} # cleared after the call + + # A system RPC (keepalive/discovery) passes no headers → slot stays empty. + async with server._rpc(): + assert server._outbound_resolved_headers == {} + assert server._outbound_resolved_headers == {} + + +# --------------------------------------------------------------------------- +# Integration tests: httpx request-event-hook mechanics +# --------------------------------------------------------------------------- +# NOTE: these exercise the httpx event-hook *mechanism* in the test's own +# event loop, where the calling task's context IS visible. The real MCP +# request runs on a dedicated background loop whose task can't see the +# caller's context (Z2O-1694), so the production hook no longer resolves +# ``${context:}`` here — it applies values resolved caller-side (see the +# TestResolveTemplatedHeadersMethod / _apply_resolved_headers tests above). +# These remain valid as httpx-behavior documentation. # --------------------------------------------------------------------------- @pytest.mark.asyncio diff --git a/tests/tools/test_mcp_structured_content.py b/tests/tools/test_mcp_structured_content.py index f4cda00f9f01..165d332bfd52 100644 --- a/tests/tools/test_mcp_structured_content.py +++ b/tests/tools/test_mcp_structured_content.py @@ -2,7 +2,7 @@ import asyncio import json -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -56,7 +56,15 @@ def _patch_mcp_server(): # `_rpc_lock` is acquired by _make_tool_handler's call path (mcp_tool.py # ~L2008) to serialize JSON-RPC against the server — build it inside the # fresh loop that _fake_run_on_mcp_loop spins up, not at fixture import. - fake_server = SimpleNamespace(session=fake_session, _rpc_lock=None) + # The tool handler now resolves ${context:} headers caller-side and applies + # them via server._rpc() under the lock (Z2O-1694). Give the fake those + # members: an empty templated-header resolver and the real _rpc CM (which + # only needs _rpc_lock + _outbound_resolved_headers, both present here). + fake_server = SimpleNamespace( + session=fake_session, _rpc_lock=None, _outbound_resolved_headers={}, + ) + fake_server._resolve_templated_headers = lambda: {} + fake_server._rpc = MethodType(mcp_tool.MCPServerTask._rpc, fake_server) with patch.dict(mcp_tool._servers, {"test-server": fake_server}), \ patch("tools.mcp_tool._run_on_mcp_loop", side_effect=_fake_run_on_mcp_loop): yield fake_session diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index 4da824a86386..50cce1babcea 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -79,6 +79,7 @@ import asyncio import concurrent.futures +import contextlib import inspect import json import logging @@ -710,6 +711,30 @@ def _replace(match: re.Match) -> str: return _CONTEXT_TEMPLATE_RE.sub(_replace, value).strip() +def _apply_resolved_headers(request, *, same_origin: bool, names, resolved: dict) -> None: + """Stamp already-resolved templated headers onto an outbound MCP request. + + Unlike per-request ``${context:}`` resolution (which must happen in the + caller's task — see ``MCPServerTask._resolve_templated_headers``), this is + pure application of pre-resolved values, safe to run on the MCP background + loop. ``resolved`` maps header-name -> value (already non-empty); a name + absent from ``resolved`` means "drop this header" (don't send it blank). + + Cross-origin safety: on a redirect to a different origin, identity-carrying + templated headers are stripped, never injected — same threat model as the + static-``Authorization`` cross-origin strip. + """ + for name in names: + if not same_origin: + request.headers.pop(name, None) + continue + value = resolved.get(name, "") + if value: + request.headers[name] = value + else: + request.headers.pop(name, None) + + def _split_static_and_templated_headers( headers: Dict[str, Any], ) -> tuple[Dict[str, Any], Dict[str, str]]: @@ -1139,6 +1164,7 @@ class MCPServerTask: "_tools", "_error", "_config", "_sampling", "_registered_tool_names", "_auth_type", "_refresh_lock", "_rpc_lock", "_pending_refresh_tasks", + "_templated_headers", "_outbound_resolved_headers", "initialize_result", ) @@ -1170,6 +1196,15 @@ def __init__(self, name: str): # transports for conservative per-server ordering. self._rpc_lock = asyncio.Lock() self._pending_refresh_tasks: set[asyncio.Task] = set() + # ${context:NAME} header templates for this server (HTTP transports), + # captured at connect time so the SYNC tool handler can resolve them + # against the caller's session context. The MCP request runs on a + # dedicated background loop whose task never sees those vars, so + # resolving in the request hook there always yielded empty — Z2O-1694. + # The handler resolves caller-side into _outbound_resolved_headers and + # the request hook (on the MCP loop) reads that instance attr. + self._templated_headers: dict[str, str] = {} + self._outbound_resolved_headers: dict[str, str] = {} # Captures the ``InitializeResult`` returned by # ``await session.initialize()`` so downstream code can inspect the # server's real advertised capabilities (``.capabilities.resources``, @@ -1177,6 +1212,44 @@ def __init__(self, name: str): # method attribute corresponds to a supported server method. See #18051. self.initialize_result: Optional[Any] = None + def _resolve_templated_headers(self) -> dict[str, str]: + """Resolve this server's ``${context:NAME}`` header templates against + the CURRENT task's session context, returning name -> value for every + template that resolved to a non-empty value. + + Must be called on the caller's thread (the sync tool handler), where + the per-turn session vars set by plugin hooks (e.g. delegated + principal) are visible. The MCP request itself runs on a dedicated + background loop whose task never sees those vars, so resolution there + always yielded empty — Z2O-1694. The returned dict is bridged onto + ``self._outbound_resolved_headers`` and applied by the request hook. + """ + out: dict[str, str] = {} + for name, template in self._templated_headers.items(): + value = _resolve_context_templates(template) + if value: + out[name] = value + return out + + @contextlib.asynccontextmanager + async def _rpc(self, resolved_headers: Optional[dict] = None): + """Serialize a single client-initiated RPC on this server and, for the + duration of the call, expose caller-resolved templated headers to the + request hook (which runs on the MCP background loop and cannot resolve + ``${context:}`` itself — Z2O-1694). + + System RPCs (keepalive pings, discovery, refresh) pass no headers, so + they never carry a user's delegated principal. Holding ``_rpc_lock`` + also guarantees a system ping can't overlap — and read — a user call's + stashed headers. + """ + async with self._rpc_lock: + self._outbound_resolved_headers = resolved_headers or {} + try: + yield + finally: + self._outbound_resolved_headers = {} + def _is_http(self) -> bool: """Check if this server uses HTTP transport.""" return "url" in self._config @@ -1341,10 +1414,14 @@ async def _wait_for_lifecycle_event(self) -> str: # to exercise the connection and detect stale sockets. if self.session: try: - await asyncio.wait_for( - self.session.list_tools(), - timeout=30.0, - ) + # Hold _rpc_lock with NO templated headers so this + # system keepalive can't overlap — and pick up — a + # user call's stashed delegated principal (Z2O-1694). + async with self._rpc(): + await asyncio.wait_for( + self.session.list_tools(), + timeout=30.0, + ) except Exception as exc: logger.warning( "MCP server '%s' keepalive failed, " @@ -1470,6 +1547,11 @@ async def _run_http(self, config: dict): # headers are resolved per-request via a request event hook so each # outgoing MCP call sees the calling task's current session state. static_headers, templated_headers = _split_static_and_templated_headers(headers) + # Stash the templates so the SYNC tool handler can resolve them against + # the caller's session context and bridge the values onto this instance + # (Z2O-1694) — the request hook below runs on the MCP background loop + # and can't see the caller's contextvars. + self._templated_headers = dict(templated_headers) connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT) ssl_verify = config.get("ssl_verify", True) @@ -1608,18 +1690,22 @@ async def _strip_identity_on_cross_origin_redirect(response): } if templated_headers: async def _inject_templated_headers(request): - """Resolve ``${context:NAME}`` templates against the - *calling task's* session context on each outbound request. + """Apply the delegated/templated headers that the SYNC tool + handler resolved against the *caller's* session context and + stashed on ``self._outbound_resolved_headers``. + + This hook runs on the MCP background loop (the SDK's + post_writer spawns the request task there), whose context + NEVER sees the per-turn session vars — so resolving + ``${context:}`` here always yielded empty (Z2O-1694). + Resolution happens caller-side; this only applies the + bridged values, serialized per server by ``_rpc_lock``. **Cross-origin safety:** the hook fires on EVERY outbound - request, including redirects. If a redirect target is on - a different origin than the configured MCP URL, we must - NOT re-inject identity-carrying templated headers there - — same threat model as the static-``Authorization`` - cross-origin strip. Without this guard, a malicious 302 - target could harvest delegated-principal claims even - though the response hook stripped them off the - next_request. + request, including redirects. On a redirect to a different + origin we strip identity-carrying templated headers rather + than inject them — same threat model as the static + ``Authorization`` cross-origin strip. """ target = request.url same_origin = ( @@ -1627,18 +1713,12 @@ async def _inject_templated_headers(request): ) == ( _original_url.scheme, _original_url.host, _original_url.port, ) - for name in _templated_for_hook: - if not same_origin: - # Redirected away from the configured origin: - # never inject identity headers, and proactively - # strip any that survived for some reason. - request.headers.pop(name, None) - continue - resolved = _resolve_context_templates(_templated_for_hook[name]) - if resolved: - request.headers[name] = resolved - elif name in request.headers: - del request.headers[name] + _apply_resolved_headers( + request, + same_origin=same_origin, + names=_templated_for_hook, + resolved=self._outbound_resolved_headers, + ) event_hooks["request"] = [_inject_templated_headers] @@ -2553,8 +2633,16 @@ def _handler(args: dict, **kwargs) -> str: "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) + # Resolve ${context:} headers HERE — on the caller's thread, where the + # per-turn session context (delegated principal, set by plugin hooks) + # is visible. The actual call runs on the MCP background loop whose + # task never sees those vars, so the request hook there can only apply + # already-resolved values (Z2O-1694). We bridge via the server + # instance inside _rpc_lock so concurrent per-server calls can't race. + _resolved_headers = server._resolve_templated_headers() + async def _call(): - async with server._rpc_lock: + async with server._rpc(_resolved_headers): result = await server.session.call_tool(tool_name, arguments=args) # MCP CallToolResult has .content (list of content blocks) and .isError if result.isError: @@ -2666,8 +2754,10 @@ def _handler(args: dict, **kwargs) -> str: "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) + _resolved_headers = server._resolve_templated_headers() + async def _call(): - async with server._rpc_lock: + async with server._rpc(_resolved_headers): result = await server.session.list_resources() resources = [] for r in (result.resources if hasattr(result, "resources") else []): @@ -2730,8 +2820,10 @@ def _handler(args: dict, **kwargs) -> str: if not uri: return tool_error("Missing required parameter 'uri'") + _resolved_headers = server._resolve_templated_headers() + async def _call(): - async with server._rpc_lock: + async with server._rpc(_resolved_headers): result = await server.session.read_resource(uri) # read_resource returns ReadResourceResult with .contents list parts: List[str] = [] @@ -2784,8 +2876,10 @@ def _handler(args: dict, **kwargs) -> str: "error": f"MCP server '{server_name}' is not connected" }, ensure_ascii=False) + _resolved_headers = server._resolve_templated_headers() + async def _call(): - async with server._rpc_lock: + async with server._rpc(_resolved_headers): result = await server.session.list_prompts() prompts = [] for p in (result.prompts if hasattr(result, "prompts") else []): @@ -2854,8 +2948,10 @@ def _handler(args: dict, **kwargs) -> str: return tool_error("Missing required parameter 'name'") arguments = args.get("arguments", {}) + _resolved_headers = server._resolve_templated_headers() + async def _call(): - async with server._rpc_lock: + async with server._rpc(_resolved_headers): result = await server.session.get_prompt(name, arguments=arguments) # GetPromptResult has .messages list messages = []