[Z2O-1694] fix(mcp): resolve ${context:} headers caller-side, not on the MCP loop

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

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,13 +1195,41 @@ 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``,
         # ``.capabilities.prompts``) instead of assuming every ``ClientSession``
         # 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,37 +1666,35 @@ 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 = (
                         target.scheme, target.host, target.port,
                     ) == (
                         _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 = ""