diff --git a/gateway/session_context.py b/gateway/session_context.py index 486949fae3de2..802fa6d791882 100644 --- a/gateway/session_context.py +++ b/gateway/session_context.py @@ -36,6 +36,7 @@ platform = get_session_env("HERMES_SESSION_PLATFORM", "") """ +import threading from contextvars import ContextVar from typing import Any @@ -82,6 +83,13 @@ "HERMES_CRON_AUTO_DELIVER_THREAD_ID": _CRON_AUTO_DELIVER_THREAD_ID, } +# Guards _VAR_MAP mutation + iteration. Plugin ``register(ctx)`` typically +# runs at startup, but ``get_registered_var_names()`` may be called from a +# different thread (diagnostics, test introspection). Without the lock, +# concurrent registration + iteration can raise ``RuntimeError: dictionary +# changed size during iteration``. +_VAR_MAP_LOCK = threading.Lock() + def set_session_vars( platform: str = "", @@ -152,6 +160,9 @@ def get_session_env(name: str, default: str = "") -> str: this context — i.e. CLI, cron scheduler, and test processes that don't use ``set_session_vars`` at all). 3. *default* + + Plugin-registered names (via :func:`register_session_context_var`) are + resolved through the same path as built-in ``HERMES_SESSION_*`` names. """ import os @@ -162,3 +173,69 @@ def get_session_env(name: str, default: str = "") -> str: return value # Fall back to os.environ for CLI, cron, and test compatibility return os.getenv(name, default) + + +def register_session_context_var(name: str, var: ContextVar) -> None: + """Register a plugin-owned ``ContextVar`` resolvable through + :func:`get_session_env` and the ``${context:NAME}`` template syntax in + MCP-server header config. + + Plugins call this during ``register(ctx)`` so their per-task state is + visible to Hermes config resolution. The registered ``ContextVar`` + follows the same resolution semantics as built-in ``HERMES_SESSION_*`` + vars: explicitly set values (even ``""``) are returned as-is; the + sentinel ``_UNSET`` default triggers ``os.environ`` fallback. + + Example:: + + from contextvars import ContextVar + from gateway.session_context import register_session_context_var + + PRINCIPAL_EMAIL: ContextVar = ContextVar( + "PRINCIPAL_EMAIL", default="", + ) + + def register(ctx): + register_session_context_var("PRINCIPAL_EMAIL", PRINCIPAL_EMAIL) + # ... plugin sets PRINCIPAL_EMAIL.set(email) per turn + + The config-side consumer (``tools/mcp_tool.py`` ``${context:NAME}`` + expansion) reads via :func:`get_session_env`, so a registered name + becomes usable in ``config.yaml`` like:: + + mcp_servers: + my_server: + headers: + X-Principal-Email: "${context:PRINCIPAL_EMAIL}" + + Registration is idempotent — repeated calls with the same name replace + the prior binding (last-writer-wins). Built-in names (``HERMES_SESSION_*``, + ``HERMES_CRON_AUTO_DELIVER_*``) can be overridden but typically should + not be. + """ + if not isinstance(name, str) or not name or not name.isascii() or not name.isidentifier(): + # The MCP-headers template regex (``${context:NAME}``) only matches + # ASCII identifier characters ``[A-Za-z_][A-Za-z0-9_]*``. Names + # outside that set (hyphens, spaces, unicode) would register + # successfully here but silently fail to resolve in templates — + # surface the invariant at registration time instead. + raise ValueError( + "register_session_context_var: name must be a valid ASCII " + "identifier matching the ${context:NAME} template regex", + ) + if not isinstance(var, ContextVar): + raise TypeError("register_session_context_var: var must be a ContextVar") + with _VAR_MAP_LOCK: + _VAR_MAP[name] = var + + +def get_registered_var_names() -> tuple[str, ...]: + """Return the set of names currently resolvable via :func:`get_session_env`. + + Useful for diagnostics and tests. Includes both built-in and + plugin-registered names. Snapshot under the registry lock so concurrent + plugin registration cannot raise ``RuntimeError: dictionary changed + size during iteration``. + """ + with _VAR_MAP_LOCK: + return tuple(_VAR_MAP.keys()) diff --git a/scripts/release.py b/scripts/release.py index 27f9fbd8f987e..58b2c65a324eb 100755 --- a/scripts/release.py +++ b/scripts/release.py @@ -60,6 +60,7 @@ "0x.badfriend@gmail.com": "discodirector", "altriatree@gmail.com": "TruaShamu", "contact-me@stark-x.cn": "Stark-X", + "chungty@gmail.com": "chungty", "nat@nthrow.io": "nthrow", "m@mobrienv.dev": "mikeyobrien", "saeed919@pm.me": "falasi", diff --git a/tests/gateway/test_session_env.py b/tests/gateway/test_session_env.py index 2b6c983a769b4..f14f6edbf7d40 100644 --- a/tests/gateway/test_session_env.py +++ b/tests/gateway/test_session_env.py @@ -6,8 +6,12 @@ from gateway.config import Platform from gateway.run import GatewayRunner from gateway.session import SessionContext, SessionSource +from contextvars import ContextVar + from gateway.session_context import ( get_session_env, + get_registered_var_names, + register_session_context_var, set_session_vars, clear_session_vars, _VAR_MAP, @@ -322,3 +326,148 @@ def blow_up(): with pytest.raises(ValueError, match="boom"): await runner._run_in_executor_with_context(blow_up) + + +# --------------------------------------------------------------------------- +# register_session_context_var — plugin-extensible registry +# --------------------------------------------------------------------------- + +def test_register_session_context_var_makes_name_resolvable(monkeypatch): + """A registered ContextVar should resolve via get_session_env using + the registered name, even though it isn't a built-in HERMES_SESSION_* + name.""" + var: ContextVar = ContextVar("CUSTOM_PRINCIPAL", default=_UNSET) + register_session_context_var("CUSTOM_PRINCIPAL", var) + monkeypatch.delenv("CUSTOM_PRINCIPAL", raising=False) + try: + var.set("alice@example.com") + assert get_session_env("CUSTOM_PRINCIPAL") == "alice@example.com" + finally: + var.set(_UNSET) + _VAR_MAP.pop("CUSTOM_PRINCIPAL", None) + + +def test_register_session_context_var_falls_back_to_env(monkeypatch): + """When the ContextVar is _UNSET, the resolver falls back to os.environ + (same semantics as built-in vars).""" + var: ContextVar = ContextVar("CUSTOM_FALLBACK", default=_UNSET) + register_session_context_var("CUSTOM_FALLBACK", var) + monkeypatch.setenv("CUSTOM_FALLBACK", "from-env") + try: + assert get_session_env("CUSTOM_FALLBACK") == "from-env" + finally: + _VAR_MAP.pop("CUSTOM_FALLBACK", None) + + +def test_register_session_context_var_explicit_empty_does_not_fall_back( + monkeypatch, +): + """Setting the ContextVar to "" is an explicit clear — the resolver + must return "" and NOT fall back to os.environ.""" + var: ContextVar = ContextVar("CUSTOM_EMPTY", default=_UNSET) + register_session_context_var("CUSTOM_EMPTY", var) + monkeypatch.setenv("CUSTOM_EMPTY", "from-env-should-not-be-seen") + try: + var.set("") + assert get_session_env("CUSTOM_EMPTY") == "" + finally: + var.set(_UNSET) + _VAR_MAP.pop("CUSTOM_EMPTY", None) + + +@pytest.mark.parametrize( + "bad_name,exc_type", + [ + ("", ValueError), + ("HAS-HYPHEN", ValueError), + ("has space", ValueError), + ("1_leading_digit", ValueError), + ("café", ValueError), # non-ASCII + ], +) +def test_register_session_context_var_rejects_invalid_inputs_and_keeps_registry_clean( + bad_name, exc_type, +): + """Defensive checks — bad inputs raise before corrupting the registry. + + The validity bar matches what the ``${context:NAME}`` template regex + will actually resolve: ASCII identifiers only. Names with hyphens, + spaces, leading digits, or unicode would register but silently fail + to resolve in templates — failing fast at registration is friendlier. + + Also asserts the registry is *not* mutated on rejection (defends + against future refactors that move validation below the assignment). + """ + valid_var: ContextVar = ContextVar("X", default=_UNSET) + with pytest.raises(exc_type): + register_session_context_var(bad_name, valid_var) + assert bad_name not in get_registered_var_names() + + +def test_register_session_context_var_rejects_non_contextvar(): + """Type check rejects non-ContextVar values without mutating the registry.""" + with pytest.raises(TypeError): + register_session_context_var("X", "not a contextvar") # type: ignore[arg-type] + # Name 'X' was the proposed key — it must NOT have been bound. + # (We can't assert "X" not in names if a previous test happened to + # register it; just confirm the registry didn't bind a string.) + from gateway.session_context import _VAR_MAP + binding = _VAR_MAP.get("X") + assert binding is None or isinstance(binding, ContextVar) + + +def test_register_session_context_var_is_idempotent_and_last_writer_wins(): + """Re-registering the same name swaps the binding (last writer wins).""" + a: ContextVar = ContextVar("DUP", default=_UNSET) + b: ContextVar = ContextVar("DUP", default=_UNSET) + register_session_context_var("DUP", a) + register_session_context_var("DUP", b) + try: + b.set("from-b") + a.set("from-a") + # b is the registered one — resolver should see "from-b". + assert get_session_env("DUP") == "from-b" + finally: + a.set(_UNSET) + b.set(_UNSET) + _VAR_MAP.pop("DUP", None) + + +def test_get_registered_var_names_includes_builtin_and_plugin(): + """The introspection helper lists both built-in and plugin-registered names.""" + var: ContextVar = ContextVar("INTROSPECT_TEST", default=_UNSET) + register_session_context_var("INTROSPECT_TEST", var) + try: + names = get_registered_var_names() + assert "HERMES_SESSION_USER_ID" in names # built-in + assert "INTROSPECT_TEST" in names # plugin-registered + finally: + _VAR_MAP.pop("INTROSPECT_TEST", None) + + +@pytest.mark.asyncio +async def test_register_session_context_var_asyncio_task_isolation(): + """Custom ContextVars registered via the public API inherit the same + per-asyncio-task isolation as built-in HERMES_SESSION_* vars.""" + var: ContextVar = ContextVar("ISO_TEST", default=_UNSET) + register_session_context_var("ISO_TEST", var) + + results: dict[str, str] = {} + + async def handler(label: str, value: str, delay: float): + token = var.set(value) + try: + await asyncio.sleep(delay) + results[label] = get_session_env("ISO_TEST") + finally: + var.reset(token) + + try: + await asyncio.gather( + handler("a", "alice", 0.02), + handler("b", "bob", 0.01), + ) + # Each task reads its own value back — no bleed. + assert results == {"a": "alice", "b": "bob"} + finally: + _VAR_MAP.pop("ISO_TEST", None) diff --git a/tests/hermes_cli/test_tui_npm_install.py b/tests/hermes_cli/test_tui_npm_install.py index b11d3b4debb3c..8092f12a578a7 100644 --- a/tests/hermes_cli/test_tui_npm_install.py +++ b/tests/hermes_cli/test_tui_npm_install.py @@ -168,7 +168,10 @@ def fail_run(*_args, **_kwargs): argv, cwd = main_mod._make_tui_argv(tmp_path, tui_dev=False) - assert argv == ["/bin/node", str(tmp_path / "dist" / "entry.js")] + # _make_tui_argv prepends --expose-gc per fix(tui): pass --expose-gc as + # node argv (commit 3d2f14646 — also at hermes_cli/main.py:1276,1282,1359). + # This test was missed when --expose-gc was added. + assert argv == ["/bin/node", "--expose-gc", str(tmp_path / "dist" / "entry.js")] assert cwd == tmp_path diff --git a/tests/tools/test_mcp_context_template.py b/tests/tools/test_mcp_context_template.py new file mode 100644 index 0000000000000..cc44ea487b7bb --- /dev/null +++ b/tests/tools/test_mcp_context_template.py @@ -0,0 +1,452 @@ +"""Tests for ``${context:NAME}`` per-request interpolation in MCP-headers config. + +Exercises both the resolution helpers (``_has_context_template``, +``_resolve_context_templates``, ``_split_static_and_templated_headers``) +and the per-request injection behavior via a real ``httpx.AsyncClient`` +event hook — the latter is the contract that matters for downstream +consumers (Mercator's delegated-principal use case). +""" + +from __future__ import annotations + +import asyncio +from contextvars import ContextVar + +import httpx +import pytest + +from gateway.session_context import ( + _UNSET, + _VAR_MAP, + register_session_context_var, +) +from tools.mcp_tool import ( + _has_context_template, + _resolve_context_templates, + _resolve_frozen_headers, + _split_static_and_templated_headers, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def custom_var(): + """Register a fresh ContextVar each test under a stable name.""" + var: ContextVar = ContextVar("TEST_PRINCIPAL", default=_UNSET) + register_session_context_var("TEST_PRINCIPAL", var) + try: + yield var + finally: + var.set(_UNSET) + _VAR_MAP.pop("TEST_PRINCIPAL", None) + + +# --------------------------------------------------------------------------- +# Unit tests: detection + resolution helpers +# --------------------------------------------------------------------------- + +class TestHasContextTemplate: + """Detection helper — recognizes ``${context:NAME}`` in header values.""" + + def test_detects_simple_template(self): + assert _has_context_template("${context:FOO}") + + def test_detects_template_with_surrounding_text(self): + assert _has_context_template("Bearer ${context:JWT}") + + def test_rejects_env_var_style(self): + # Env vars are config-load-time, not per-request. + assert not _has_context_template("${PLAIN_ENV}") + + def test_rejects_literal_string(self): + assert not _has_context_template("plain-static-value") + + def test_rejects_non_string(self): + assert not _has_context_template(None) + assert not _has_context_template(42) + assert not _has_context_template({"x": "${context:Y}"}) + + +class TestResolveContextTemplates: + """Resolution helper — substitutes ``${context:NAME}`` against session_context.""" + + def test_resolves_set_value(self, custom_var): + custom_var.set("alice@example.com") + assert _resolve_context_templates("${context:TEST_PRINCIPAL}") == "alice@example.com" + + def test_unset_resolves_to_empty(self, custom_var): + # ContextVar default is _UNSET — resolver falls back to os.environ + # (also unset for TEST_PRINCIPAL) and finally returns "". + assert _resolve_context_templates("${context:TEST_PRINCIPAL}") == "" + + def test_mixed_literal_and_template(self, custom_var): + custom_var.set("opaque-token-xyz") + assert ( + _resolve_context_templates("Bearer ${context:TEST_PRINCIPAL}") + == "Bearer opaque-token-xyz" + ) + + def test_multiple_templates_in_one_string(self, custom_var): + # Register a second var inline for this test. + other: ContextVar = ContextVar("TEST_OTHER", default=_UNSET) + register_session_context_var("TEST_OTHER", other) + try: + custom_var.set("a") + other.set("b") + result = _resolve_context_templates( + "${context:TEST_PRINCIPAL}|${context:TEST_OTHER}" + ) + assert result == "a|b" + finally: + other.set(_UNSET) + _VAR_MAP.pop("TEST_OTHER", None) + + def test_non_string_passes_through(self): + assert _resolve_context_templates(42) == 42 # type: ignore[arg-type] + + def test_unknown_name_resolves_to_empty(self): + # No var registered for THIS_NAME — falls back to os.environ, + # which is unset in tests, so empty string. + assert _resolve_context_templates("${context:NEVER_REGISTERED}") == "" + + def test_whitespace_only_resolved_value_is_stripped_to_empty(self, custom_var): + """A context-var holding only whitespace must NOT produce a + header sent as ``X-Foo: `` (HTTP-spec-violating, useless to + the recipient). The resolver strips leading/trailing whitespace + so all three transports treat ``" "`` the same as ``""`` — + i.e. "drop this header" at the call site.""" + custom_var.set(" ") + assert _resolve_context_templates("${context:TEST_PRINCIPAL}") == "" + + def test_leading_trailing_whitespace_stripped(self, custom_var): + custom_var.set(" alice@example.com\t") + assert _resolve_context_templates("${context:TEST_PRINCIPAL}") == "alice@example.com" + + def test_inner_whitespace_preserved(self, custom_var): + """Strip only touches the outer edges — inner whitespace + (e.g. ``"Bearer abc 123"``) is intentional and preserved.""" + custom_var.set("abc 123") + assert ( + _resolve_context_templates("Bearer ${context:TEST_PRINCIPAL}") + == "Bearer abc 123" + ) + + @pytest.mark.parametrize( + "raw_value,expected", + [ + (42, "42"), + (0, "0"), + (None, "None"), + (False, "False"), + (True, "True"), + ([], "[]"), + ({"k": "v"}, "{'k': 'v'}"), + ], + ) + def test_non_string_contextvar_value_coerces_to_str(self, raw_value, expected): + """A plugin-registered ContextVar holding any non-string value must + not crash the regex substitution. ``str()`` is the documented + coercion path; future refactors that change this (e.g. drop-header- + when-None) must update both this test and the docstring deliberately. + """ + from contextvars import ContextVar as _ContextVar + weird: _ContextVar = _ContextVar("WEIRD_VAL", default=_UNSET) + register_session_context_var("WEIRD_VAL", weird) + try: + weird.set(raw_value) + assert _resolve_context_templates("${context:WEIRD_VAL}") == expected + finally: + weird.set(_UNSET) + _VAR_MAP.pop("WEIRD_VAL", None) + + +class TestResolveFrozenHeaders: + """Freeze-time merge for SSE / legacy-HTTP transports — drops empty + templated entries so an at-init resolution of an unset context var + doesn't ship a blank header for the connection's lifetime. + """ + + def test_set_value_resolves_into_merged_dict(self, custom_var): + custom_var.set("thomas@verdigris.co") + merged = _resolve_frozen_headers( + {"Authorization": "Bearer static"}, + {"X-Delegated-Principal": "${context:TEST_PRINCIPAL}"}, + ) + assert merged == { + "Authorization": "Bearer static", + "X-Delegated-Principal": "thomas@verdigris.co", + } + + def test_empty_resolved_value_drops_header(self, custom_var): + """The key fix: a templated header whose context var is unset + at freeze time must be ABSENT from the result, not present with + an empty value. Same drop-on-empty discipline as the HTTP + per-request hook.""" + # custom_var default is _UNSET → resolves to "" + merged = _resolve_frozen_headers( + {"Authorization": "Bearer static"}, + {"X-Delegated-Principal": "${context:TEST_PRINCIPAL}"}, + ) + assert merged == {"Authorization": "Bearer static"} + assert "X-Delegated-Principal" not in merged + + def test_whitespace_only_resolved_value_drops_header(self, custom_var): + """Whitespace-only resolution is equivalent to empty — same drop + behavior (avoids sending ``X-Foo: `` which is HTTP-spec + violating and useless to the recipient).""" + custom_var.set(" \t") + merged = _resolve_frozen_headers( + {}, + {"X-Delegated-Principal": "${context:TEST_PRINCIPAL}"}, + ) + assert merged == {} + + def test_static_only(self): + """Pure static headers pass through unchanged.""" + merged = _resolve_frozen_headers( + {"Authorization": "Bearer x", "X-Trace-Id": "abc"}, + {}, + ) + assert merged == {"Authorization": "Bearer x", "X-Trace-Id": "abc"} + + def test_empty_input(self): + assert _resolve_frozen_headers({}, {}) == {} + + +class TestSplitStaticAndTemplatedHeaders: + """Header partition helper — separates client-default-safe from per-request.""" + + def test_partitions_correctly(self): + headers = { + "Authorization": "Bearer static-token", + "X-Delegated-Principal": "${context:PRINCIPAL_EMAIL}", + "X-Other": "static-other", + "X-Mixed": "prefix-${context:NAME}-suffix", + } + static, templated = _split_static_and_templated_headers(headers) + assert static == { + "Authorization": "Bearer static-token", + "X-Other": "static-other", + } + assert templated == { + "X-Delegated-Principal": "${context:PRINCIPAL_EMAIL}", + "X-Mixed": "prefix-${context:NAME}-suffix", + } + + def test_empty_input(self): + assert _split_static_and_templated_headers({}) == ({}, {}) + + def test_non_string_values_treated_as_static(self): + # Edge case: a header value that's not a string (unusual but + # possible from raw YAML) is treated as static — we don't try + # to substitute on non-strings. + headers = {"X-Number": 42, "X-Bool": True} + static, templated = _split_static_and_templated_headers(headers) + assert static == {"X-Number": 42, "X-Bool": True} + assert templated == {} + + +# --------------------------------------------------------------------------- +# Integration test: the actual event-hook injection pattern using httpx +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_per_request_header_injection_via_event_hook(custom_var): + """The pattern used inside ``_run_http`` — install a request event + hook on httpx.AsyncClient that resolves ``${context:NAME}`` templates + against the *calling task's* session context. Verifies the per-request + contract end-to-end (no MCP-stack mocking required).""" + captured_headers: dict[str, str] = {} + + async def _capture_handler(request: httpx.Request) -> httpx.Response: + captured_headers.clear() + captured_headers.update(dict(request.headers)) + return httpx.Response(200, json={"ok": True}) + + templated = {"X-Delegated-Principal": "${context:TEST_PRINCIPAL}"} + + async def _inject(request: httpx.Request) -> None: + for name, template in templated.items(): + resolved = _resolve_context_templates(template) + if resolved: + request.headers[name] = resolved + elif name in request.headers: + del request.headers[name] + + transport = httpx.MockTransport(_capture_handler) + async with httpx.AsyncClient( + transport=transport, + event_hooks={"request": [_inject]}, + headers={"Authorization": "Bearer static-token"}, + ) as client: + # Request 1: principal "alice" — both headers present. + custom_var.set("alice@example.com") + await client.get("http://example.invalid/mcp") + assert captured_headers["x-delegated-principal"] == "alice@example.com" + assert captured_headers["authorization"] == "Bearer static-token" + + # Request 2: principal "bob" — verifies per-request resolution + # (the same client emits a different header value on the next call). + custom_var.set("bob@example.com") + await client.get("http://example.invalid/mcp") + assert captured_headers["x-delegated-principal"] == "bob@example.com" + + # Request 3: principal cleared — header is dropped (NOT sent empty). + custom_var.set("") + await client.get("http://example.invalid/mcp") + assert "x-delegated-principal" not in captured_headers + + +@pytest.mark.asyncio +async def test_principal_mutated_between_requests_in_same_task(custom_var): + """Within a single task, mutating the ContextVar between requests must + cause each request to carry the *current* value — not the value at + hook-install time. + + Catches the latent bug class where the resolution would be cached at + ``_run_http`` invocation time (e.g. ``resolved = _resolve_context_templates(template)`` + lifted out of the inner hook). The + ``test_asyncio_task_isolation_for_per_request_headers`` test below + would silently miss this because each task only sets the var once. + """ + seen: list[str] = [] + + async def _capture(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("x-delegated-principal", "")) + return httpx.Response(200) + + async def _inject(request: httpx.Request) -> None: + resolved = _resolve_context_templates("${context:TEST_PRINCIPAL}") + if resolved: + request.headers["X-Delegated-Principal"] = resolved + elif "x-delegated-principal" in request.headers: + del request.headers["x-delegated-principal"] + + transport = httpx.MockTransport(_capture) + async with httpx.AsyncClient( + transport=transport, event_hooks={"request": [_inject]}, + ) as client: + for value in ("alice@example.com", "bob@example.com", "carol@example.com"): + custom_var.set(value) + await client.get("http://example.invalid/mcp") + + assert seen == [ + "alice@example.com", + "bob@example.com", + "carol@example.com", + ] + + +@pytest.mark.asyncio +async def test_templated_header_dropped_on_cross_origin_redirect(custom_var): + """Identity-carrying templated headers must NOT follow a cross-origin + redirect. Same threat model as ``Authorization``-stripping: if a + redirect target is attacker-controlled, leaking ``X-Delegated-Principal`` + would let the attacker harvest user attribution claims. + """ + # Two responses: the first redirects (302) to a different origin; the + # second is captured so we can inspect the headers that actually crossed. + captured_on_redirect_target: dict[str, str] = {} + + async def _handler(request: httpx.Request) -> httpx.Response: + host = request.url.host + if host == "original.invalid": + return httpx.Response( + 302, + headers={"location": "http://attacker.invalid/harvest"}, + ) + captured_on_redirect_target.clear() + captured_on_redirect_target.update(dict(request.headers)) + return httpx.Response(200) + + # Replicate the production response hook *and* request hook patterns, + # including the origin guard inside the request hook (the response-side + # strip alone is insufficient — the request hook fires again on the + # redirected request). + _original_url = httpx.URL("http://original.invalid/mcp") + templated = {"X-Delegated-Principal": "${context:TEST_PRINCIPAL}"} + + async def _strip_identity_on_redirect(response: httpx.Response) -> None: + if response.is_redirect and response.next_request: + target = response.next_request.url + if (target.scheme, target.host, target.port) != ( + _original_url.scheme, _original_url.host, _original_url.port, + ): + response.next_request.headers.pop("Authorization", None) + response.next_request.headers.pop("authorization", None) + for name in templated: + response.next_request.headers.pop(name, None) + + async def _inject(request: httpx.Request) -> None: + target = request.url + same_origin = (target.scheme, target.host, target.port) == ( + _original_url.scheme, _original_url.host, _original_url.port, + ) + for name in templated: + if not same_origin: + request.headers.pop(name, None) + continue + resolved = _resolve_context_templates(templated[name]) + if resolved: + request.headers[name] = resolved + elif name in request.headers: + del request.headers[name] + + transport = httpx.MockTransport(_handler) + async with httpx.AsyncClient( + transport=transport, + follow_redirects=True, + event_hooks={ + "request": [_inject], + "response": [_strip_identity_on_redirect], + }, + headers={"Authorization": "Bearer static-token"}, + ) as client: + custom_var.set("thomas@verdigris.co") + await client.get("http://original.invalid/mcp") + + # The redirect target should have neither the static Authorization + # header nor the templated delegated-principal header. + assert "authorization" not in captured_on_redirect_target + assert "x-delegated-principal" not in captured_on_redirect_target + + +@pytest.mark.asyncio +async def test_asyncio_task_isolation_for_per_request_headers(custom_var): + """Two concurrent asyncio tasks setting different principals must each + see their own value on outbound requests — no cross-task leak.""" + seen_per_task: dict[str, str] = {} + + async def _record_handler(request: httpx.Request) -> httpx.Response: + # Pull the calling task's principal out of the request header. + seen_per_task[request.url.path.lstrip("/")] = request.headers.get( + "x-delegated-principal", "" + ) + return httpx.Response(200) + + async def _inject(request: httpx.Request) -> None: + resolved = _resolve_context_templates("${context:TEST_PRINCIPAL}") + if resolved: + request.headers["X-Delegated-Principal"] = resolved + + transport = httpx.MockTransport(_record_handler) + async with httpx.AsyncClient( + transport=transport, event_hooks={"request": [_inject]}, + ) as client: + async def handler(path: str, principal: str, delay: float): + custom_var.set(principal) + await asyncio.sleep(delay) + await client.get(f"http://example.invalid/{path}") + + await asyncio.gather( + handler("alice-path", "alice@example.com", 0.02), + handler("bob-path", "bob@example.com", 0.01), + ) + + assert seen_per_task == { + "alice-path": "alice@example.com", + "bob-path": "bob@example.com", + } diff --git a/tools/mcp_tool.py b/tools/mcp_tool.py index e50efc05a0c28..4da824a863861 100644 --- a/tools/mcp_tool.py +++ b/tools/mcp_tool.py @@ -638,6 +638,120 @@ def _safe_numeric(value, default, coerce=int, minimum=1): return default +# --------------------------------------------------------------------------- +# Per-request context-template resolution for MCP-headers config. +# +# Distinguishing case from config-load-time substitution: +# +# ``${ENV_VAR}`` — substituted ONCE at config load time by an +# external envsubst pass (typically a +# deploy-time entrypoint script). Values are +# frozen for the lifetime of the MCP server task. +# +# ``${context:NAME}`` — left intact through config load, then +# resolved PER REQUEST from +# ``gateway.session_context``. Asyncio-task-local, +# so concurrent handler tasks each see their own +# session state and never bleed into each other. +# +# This split lets consumers attach per-request values (most commonly delegated +# user identity) to outbound MCP HTTP headers without monkey-patching httpx or +# routing through a sidecar process. Transport-specific behavior: +# +# * Streamable HTTP (default for remote MCP servers) — per-request hook on +# httpx.AsyncClient.event_hooks["request"] performs substitution. Empty +# resolved values cause the header to be dropped (not sent as an empty +# string). +# +# * SSE — the headers dict is passed to ``sse_client`` once at stream-open; +# per-request substitution is not possible. Templates resolve at +# stream-open time and are then frozen for the SSE connection's +# lifetime (typically minutes-to-hours). +# --------------------------------------------------------------------------- + +_CONTEXT_TEMPLATE_RE = re.compile(r'\$\{context:([A-Za-z_][A-Za-z0-9_]*)\}') + + +def _has_context_template(value: Any) -> bool: + """Return True if *value* is a string containing ``${context:NAME}``.""" + return isinstance(value, str) and bool(_CONTEXT_TEMPLATE_RE.search(value)) + + +def _resolve_context_templates(value: str) -> str: + """Replace every ``${context:NAME}`` occurrence in *value* with the + current resolution of ``NAME`` via ``gateway.session_context``. + + Unset names resolve to ``""``. Mixed literal text + templates is fine + (e.g. ``"Bearer ${context:JWT}"``). Non-template input is returned + unchanged. + + The final substituted string is stripped of leading/trailing + whitespace. An all-whitespace resolved value (``" "`` from a + paste artifact, env-injection mishap, etc.) is therefore equivalent + to empty — which downstream call sites treat as "drop this header" + rather than send ``X-Foo:`` with whitespace. Intentional inner + whitespace ("Bearer abc 123") is preserved. + """ + if not isinstance(value, str): + return value + # Lazy import — keeps the gateway -> tools dependency one-way and avoids + # paying the import cost when no headers use templates. + from gateway.session_context import get_session_env + + def _replace(match: re.Match) -> str: + name = match.group(1) + # str()-coerce in case a plugin registers a ContextVar whose value + # is not a string (int, bool, None). re.sub raises TypeError on + # non-string replacements; defensive coercion keeps the resolver + # robust against plugin authors who didn't read the str-only + # convention. + return str(get_session_env(name, "")) + + return _CONTEXT_TEMPLATE_RE.sub(_replace, value).strip() + + +def _split_static_and_templated_headers( + headers: Dict[str, Any], +) -> tuple[Dict[str, Any], Dict[str, str]]: + """Separate a headers dict into a static part and a templated part. + + Returns ``(static, templated)`` where *static* contains values with no + ``${context:NAME}`` placeholders (safe to set as httpx client defaults) + and *templated* contains the raw template strings (to be resolved + per-request via a request hook). + """ + static: Dict[str, Any] = {} + templated: Dict[str, str] = {} + for key, value in headers.items(): + if _has_context_template(value): + templated[key] = value + else: + static[key] = value + return static, templated + + +def _resolve_frozen_headers( + static_headers: Dict[str, Any], + templated_headers: Dict[str, str], +) -> Dict[str, Any]: + """Resolve templated headers ONCE (for SSE / legacy-HTTP transports + that don't support per-request injection) and merge with statics. + + Templated entries whose resolution is empty after stripping are + dropped (don't emit ``X-Foo:`` with blank content for the + connection's lifetime). Used by both the SSE and the legacy HTTP + branches; extracted as a helper so the freeze-time semantics are + testable without spinning up a real ``sse_client`` / + ``streamablehttp_client`` instance. + """ + resolved_templated = { + key: resolved + for key, template in templated_headers.items() + if (resolved := _resolve_context_templates(template)) + } + return {**static_headers, **resolved_templated} + + class SamplingHandler: """Handles sampling/createMessage requests for a single MCP server. @@ -1351,6 +1465,11 @@ async def _run_http(self, config: dict): # case-insensitive so conventional casing is preserved. if not any(key.lower() == "mcp-protocol-version" for key in headers): headers["mcp-protocol-version"] = LATEST_PROTOCOL_VERSION + # Separate headers carrying ``${context:NAME}`` templates from static + # ones. Static headers go on the httpx client as defaults; templated + # 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) connect_timeout = config.get("connect_timeout", _DEFAULT_CONNECT_TIMEOUT) ssl_verify = config.get("ssl_verify", True) @@ -1393,9 +1512,35 @@ async def _run_http(self, config: dict): # Streamable HTTP code path's httpx read timeout below. Original # observation from @amiller in PR #5981 (Router Teamwork, # Supermemory on Cloudflare Workers idle-disconnect at ~60s). + # SSE is one long-lived connection — per-request substitution is + # not possible. Resolve any ``${context:NAME}`` templates ONCE + # at stream-open and pass the resulting static dict. Sessions + # spun up before the session context vars are set will see empty + # values, which is consistent with non-templated unset env vars. + # + # Warn loudly: if the caller configured templated headers on an + # SSE server, they almost certainly meant per-request semantics. + # Frozen-at-open is a real footgun for identity-carrying headers + # (a stream opened by user A keeps user A's principal for every + # user-B turn over its lifetime). + if templated_headers: + logger.warning( + "MCP server '%s' (SSE transport): templated headers " + "%s are resolved ONCE at stream-open and frozen for " + "the connection's lifetime — per-request resolution is " + "not possible on SSE. Switch to Streamable HTTP " + "(transport: http) if you need per-request values.", + self.name, sorted(templated_headers.keys()), + ) + # Resolve+merge with empty-drop discipline — see + # ``_resolve_frozen_headers`` doc for why this matters on SSE + # (a stream opened before any context var is set would + # otherwise send ``X-Foo:`` with blank content for the + # connection's lifetime). + sse_headers = _resolve_frozen_headers(static_headers, templated_headers) _sse_kwargs: dict = { "url": url, - "headers": headers or None, + "headers": sse_headers or None, "timeout": float(connect_timeout), "sse_read_timeout": 300.0, } @@ -1426,9 +1571,24 @@ async def _run_http(self, config: dict): import httpx _original_url = httpx.URL(url) - - async def _strip_auth_on_cross_origin_redirect(response): - """Strip Authorization headers when redirected to a different origin.""" + # Bind the per-server templated_headers dict into the closures + # so concurrent MCP servers don't share resolution state, and + # so the cross-origin redirect stripper knows which identity- + # carrying headers must be dropped on origin change. + _templated_for_hook = dict(templated_headers) + + async def _strip_identity_on_cross_origin_redirect(response): + """Strip identity-carrying headers when redirected to a + different origin. + + Covers the static ``Authorization`` header (the original + concern) *and* any ``${context:NAME}``-templated headers + (delegated principal, signed assertions, etc.) injected + by the per-request hook below. Both classes follow the + same threat model: an attacker who controls a redirect + target could otherwise harvest identity tokens by 302'ing + to a server they control. + """ if response.is_redirect and response.next_request: target = response.next_request.url if (target.scheme, target.host, target.port) != ( @@ -1436,15 +1596,60 @@ async def _strip_auth_on_cross_origin_redirect(response): ): response.next_request.headers.pop("authorization", None) response.next_request.headers.pop("Authorization", None) + # Drop every templated identity header on cross-origin + # redirect. Header names are case-sensitive in the + # httpx.Headers store via __delitem__ but case- + # insensitive on read — pop() handles both casings. + for name in _templated_for_hook: + response.next_request.headers.pop(name, None) + + event_hooks: Dict[str, List[Any]] = { + "response": [_strip_identity_on_cross_origin_redirect], + } + if templated_headers: + async def _inject_templated_headers(request): + """Resolve ``${context:NAME}`` templates against the + *calling task's* session context on each outbound request. + + **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. + """ + 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] + + event_hooks["request"] = [_inject_templated_headers] client_kwargs: dict = { "follow_redirects": True, "timeout": httpx.Timeout(float(connect_timeout), read=300.0), "verify": ssl_verify, - "event_hooks": {"response": [_strip_auth_on_cross_origin_redirect]}, + "event_hooks": event_hooks, } - if headers: - client_kwargs["headers"] = headers + if static_headers: + client_kwargs["headers"] = static_headers if _oauth_auth is not None: client_kwargs["auth"] = _oauth_auth @@ -1466,9 +1671,28 @@ async def _strip_auth_on_cross_origin_redirect(response): "tearing down HTTP session", self.name, ) else: - # Deprecated API (mcp < 1.24.0): manages httpx client internally. + # Deprecated API (mcp < 1.24.0): manages httpx client internally, + # so we can't install a per-request event hook. Resolve any + # ``${context:NAME}`` templates once at startup and freeze — the + # session-context value at server-init time is what every + # subsequent MCP call will carry. Acceptable for the deprecated + # path; new HTTP path (mcp >= 1.24.0) uses per-request resolution. + # + # Warn loudly: same footgun as the SSE branch — frozen-at-open + # identity is wrong for any multi-user delegation use case. + if templated_headers: + logger.warning( + "MCP server '%s' (legacy HTTP, mcp < 1.24.0): templated " + "headers %s are resolved ONCE at startup and frozen — " + "the per-request resolution path requires mcp >= 1.24.0. " + "Upgrade the mcp package if you need per-request values.", + self.name, sorted(templated_headers.keys()), + ) + # Same resolve-and-merge with empty-drop discipline as the + # SSE branch above — see ``_resolve_frozen_headers``. + legacy_headers = _resolve_frozen_headers(static_headers, templated_headers) _http_kwargs: dict = { - "headers": headers, + "headers": legacy_headers, "timeout": float(connect_timeout), "verify": ssl_verify, }