From ca3210be87146668468e1160e672f6fedb65210b Mon Sep 17 00:00:00 2001 From: chungty Date: Mon, 1 Jun 2026 23:23:13 -0700 Subject: [PATCH 1/3] [Z2O-1694] fix(mcp): resolve ${context:} headers caller-side, not on the MCP loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ${context:NAME} templated MCP headers (delegated principal, signed assertions) were resolved inside the httpx request event hook — which runs on the dedicated MCP background loop, in a task spawned by the SDK's post_writer. That task's contextvar context is frozen at server-connect time and never sees the per-turn session vars set on the caller's thread, so resolution always returned "" and the header was dropped. Net effect: delegated-principal propagation silently no-ops for every streamable-HTTP MCP server (the default for remote servers). Diagnosed end-to-end from Mercator: the plugin resolved the principal (logs: status=resolved @verdigris.co) but Meridian saw email:null; a direct curl to the gateway with vs without the header confirmed the gateway honors it and Hermes was sending it empty. Fix: resolve the templates on the CALLER's thread (the sync tool handler, where the session vars are live), bridge the resolved values onto the server instance, and have the request hook apply them. The apply happens inside _rpc_lock so concurrent per-server calls can't race, and is cleared after each call so non-tool requests (pings, reconnect GETs) never carry a stale principal. - _resolve_context_templates stays caller-side via the new MCPServerTask._resolve_templated_headers(). - _apply_resolved_headers (module helper) stamps pre-resolved values, preserving the cross-origin identity-header strip. - The request hook reads self._outbound_resolved_headers instead of resolving against its own (wrong) context. SSE / legacy-HTTP transports already freeze templated headers at connect (documented limitation) and are unchanged. Tests: caller-side resolution, loop-side application, cross-origin strip, and the crux regression — a value resolved while the var is set still injects after that context is gone (simulating the MCP loop). Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/tools/test_mcp_context_template.py | 109 +++++++++++++++++++++- tools/mcp_tool.py | 114 ++++++++++++++++++----- 2 files changed, 199 insertions(+), 24 deletions(-) diff --git a/tests/tools/test_mcp_context_template.py b/tests/tools/test_mcp_context_template.py index cc44ea487b7b..a9b57d9ddfe4 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,112 @@ 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" + + +# --------------------------------------------------------------------------- +# 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/tools/mcp_tool.py b/tools/mcp_tool.py index 4da824a86386..a14ebca4278d 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -710,6 +710,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 + elif name in request.headers: + del request.headers[name] + + def _split_static_and_templated_headers( headers: Dict[str, Any], ) -> tuple[Dict[str, Any], Dict[str, str]]: @@ -1139,6 +1163,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 +1195,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 +1211,25 @@ 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 + def _is_http(self) -> bool: """Check if this server uses HTTP transport.""" return "url" in self._config @@ -1470,6 +1523,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 +1666,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 +1689,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,9 +2609,21 @@ 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: - result = await server.session.call_tool(tool_name, arguments=args) + server._outbound_resolved_headers = _resolved_headers + try: + result = await server.session.call_tool(tool_name, arguments=args) + finally: + server._outbound_resolved_headers = {} # MCP CallToolResult has .content (list of content blocks) and .isError if result.isError: error_text = "" From e9aec1477a272f9fad38d2499bedbdebb6d552d1 Mon Sep 17 00:00:00 2001 From: chungty Date: Mon, 1 Jun 2026 23:30:21 -0700 Subject: [PATCH 2/3] [Z2O-1694] address Gemini review: cover all RPC handlers + close keepalive principal leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gemini review on PR #2 flagged two real issues: 1. (high) Other MCP handlers (list_resources, read_resource, list_prompts, get_prompt) also issue outbound HTTP but didn't resolve/stash templated headers — delegated principal silently dropped for resource/prompt ops. 2. (high) Background keepalive ping (list_tools in _wait_for_lifecycle_event) doesn't take _rpc_lock, so it could overlap a user tool call and pick up the stashed _outbound_resolved_headers — leaking the user's principal on a system keepalive. Fix: - Extract `MCPServerTask._rpc(resolved_headers=None)` async context manager: acquires _rpc_lock, exposes the caller-resolved headers for the call's duration, clears them after. System RPCs pass no headers and still hold the lock, so they can't overlap-and-read a user call's principal. - Use it in all five user-facing handlers (tool + 4 resource/prompt) with caller-side _resolve_templated_headers(), and in the keepalive ping (empty headers) — closing the leak. - (medium) Simplify _apply_resolved_headers drop path to headers.pop(name, None). Header partitioning (static vs templated, Gemini ref #3) was already in place via _split_static_and_templated_headers — only templated names are resolved. Tests: + _rpc scoping/clear contract (system RPC carries no headers). 94 passed across context-template + mcp + session-env suites. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/tools/test_mcp_context_template.py | 20 ++++++++ tools/mcp_tool.py | 60 +++++++++++++++++------- 2 files changed, 64 insertions(+), 16 deletions(-) diff --git a/tests/tools/test_mcp_context_template.py b/tests/tools/test_mcp_context_template.py index a9b57d9ddfe4..0ba3d0a3a966 100644 --- a/tests/tools/test_mcp_context_template.py +++ b/tests/tools/test_mcp_context_template.py @@ -348,6 +348,26 @@ def test_resolved_headers_survive_lost_caller_context(custom_var): 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 # --------------------------------------------------------------------------- diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index a14ebca4278d..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 @@ -730,8 +731,8 @@ def _apply_resolved_headers(request, *, same_origin: bool, names, resolved: dict value = resolved.get(name, "") if value: request.headers[name] = value - elif name in request.headers: - del request.headers[name] + else: + request.headers.pop(name, None) def _split_static_and_templated_headers( @@ -1230,6 +1231,25 @@ def _resolve_templated_headers(self) -> dict[str, str]: 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 @@ -1394,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, " @@ -2618,12 +2642,8 @@ def _handler(args: dict, **kwargs) -> str: _resolved_headers = server._resolve_templated_headers() async def _call(): - async with server._rpc_lock: - server._outbound_resolved_headers = _resolved_headers - try: - result = await server.session.call_tool(tool_name, arguments=args) - finally: - server._outbound_resolved_headers = {} + 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: error_text = "" @@ -2734,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 []): @@ -2798,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] = [] @@ -2852,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 []): @@ -2922,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 = [] From d31a0b10fb767b04c08d25693d7f06a806b2e21e Mon Sep 17 00:00:00 2001 From: chungty Date: Mon, 1 Jun 2026 23:43:23 -0700 Subject: [PATCH 3/3] [Z2O-1694] test: fix structured-content fake server for caller-side header resolution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit test_mcp_structured_content fakes `server` as a SimpleNamespace; the tool handler now calls server._resolve_templated_headers() (caller-side) and server._rpc() under the lock, which the bare fake lacked → AttributeError. Give the fake an empty resolver and bind the real _rpc CM (it only needs _rpc_lock + _outbound_resolved_headers, both present). Pure test fixture fix; no production change. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/tools/test_mcp_structured_content.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) 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