From 7ace4a122b0fb9d8ffa87a057c154d8492d087be Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 11 Oct 2026 16:23:14 -0700 Subject: [PATCH] refactor(gateway): one Caller value for role and session; the three binding scopes side by side Behaviour-neutral. The role lived in one ContextVar bound per request, the session key in another bound per MCP session, and the _meta stamp was assembled by hand in the MCP proxy tool, in /step and in the trigger engine. A frozen Caller(role, session) in caller.py now carries both, Caller.current() reads it and Caller.meta() is the one place that produces the stamp, with the `default` sentinel documented where it is stamped: a _meta role entry of any value is what marks a call as proxied under the protocol, so an unforwarded caller is stamped `default` on purpose. _call_child_tool takes a Caller (None for an unattributed trigger action). The per-MCP-session lifespan leaves __init__ for a _session_lifespan method next to _asgi_app (request scope) and _process_lifespan (process scope), so the three scopes read in one place and __init__ only wires. The wire names stay literals in constants.py because the gateway container ships without agentenv_protocol; a unit test now pins them to the protocol's exports instead. Co-Authored-By: Claude Fable 5.1 --- src/agent_env/env/gateway/caller.py | 53 ++++++++++++++ src/agent_env/env/gateway/gateway.py | 74 ++++++++++---------- src/agent_env/env/gateway/triggers.py | 7 +- tst/unit/env/gateway/caller_test.py | 61 ++++++++++++++++ tst/unit/env/gateway/role_forwarding_test.py | 8 +-- tst/unit/env/gateway/triggers_test.py | 2 +- 6 files changed, 160 insertions(+), 45 deletions(-) create mode 100644 src/agent_env/env/gateway/caller.py create mode 100644 tst/unit/env/gateway/caller_test.py diff --git a/src/agent_env/env/gateway/caller.py b/src/agent_env/env/gateway/caller.py new file mode 100644 index 00000000..4c5c50ed --- /dev/null +++ b/src/agent_env/env/gateway/caller.py @@ -0,0 +1,53 @@ +"""Who the gateway acts for on one inbound request. + +A ``Caller`` is the role the request's ``AgentEnv-Role`` header named (``default`` when it named none) and the key of +the MCP session the request belongs to (empty on the REST and harness paths, which have no session). It is bound in +three scopes, which read in order in ``gateway.py``: per HTTP request by the served app (the role), per inbound MCP +session by the session lifespan (the key, kept for every tools/call of that session), and never per process. Every +child ``tools/call`` carries ``Caller.meta()``, so the three places that stamp one (the MCP proxy tool, ``/step`` and +the trigger engine) cannot disagree about the wire shape. +""" + +from __future__ import annotations + +import dataclasses +from contextvars import ContextVar, Token + +from .constants import AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE + + +@dataclasses.dataclass(frozen=True) +class Caller: + role: str = DEFAULT_ROLE + session: str = "" + + @staticmethod + def current() -> "Caller": + """The caller bound for the running task: an unforwarded one outside any request.""" + return _current.get() + + def bind(self) -> Token: + """Make this the current caller; ``Caller.reset`` the returned token to restore the previous one.""" + return _current.set(self) + + @staticmethod + def reset(token: Token) -> None: + _current.reset(token) + + def meta(self) -> dict[str, str]: + """The ``_meta`` stamp for a child ``tools/call``. + + The role entry is always present, ``default`` included, and that is load-bearing: under the protocol a + ``_meta`` role entry of any value marks the call as proxied, so the child reads no header of the gateway's + connection and takes the session from ``_meta`` or none. Stamping ``default`` is how an unforwarded caller + stays unforwarded behind a gateway whose own connection carries no identity. The session entry rides along + only when the call came in on an MCP session, so a child keying per-session state keeps one owner per agent + across the gateway's one shared connection to it. + """ + meta = {AGENT_ENV_ROLE_META_KEY: self.role} + if self.session: + meta[AGENT_ENV_SESSION_META_KEY] = self.session + return meta + + +_current: ContextVar[Caller] = ContextVar("agent_env_gateway_caller", default=Caller()) diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 3a4af677..62932269 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -12,7 +12,6 @@ import socket import subprocess from contextlib import asynccontextmanager -from contextvars import ContextVar from datetime import datetime, timezone from pathlib import Path from collections.abc import AsyncIterator @@ -65,6 +64,7 @@ WILDCARD, GatewayMode, ) +from .caller import Caller from .child_transport import child_http_client from .clock import Clock, ClockError, _iso from .get_time import GET_TIME_TOOL_NAME, build_get_time_tool @@ -72,9 +72,6 @@ logger = logging.getLogger(__name__) -_role_var: ContextVar[str] = ContextVar("agent_env_role", default=DEFAULT_ROLE) -# One key per inbound MCP session (set by the per-session lifespan); empty outside one (/step, triggers, startup). -_caller_session_var: ContextVar[str] = ContextVar("agent_env_caller_session", default="") class Gateway: @@ -139,20 +136,7 @@ def __init__( self._trajectory_file = open(GATEWAY_TRAJECTORY_FILE, "a") atexit.register(self._close_trajectory_file) - @asynccontextmanager - async def lifespan(mcp: FastMCP): - # Entered by mcp's Server.run once per inbound MCP session, in the session's server task before any - # handler task is spawned, so the key set here is inherited by every tools/call of that session. - # Process-scoped state (child sessions, time-trigger driver) lives in _process_lifespan; discovery - # here only lets a connect heal a failed startup discovery (tools/list is synchronous and cannot). - _caller_session_var.set(uuid4().hex) - try: - await self._ensure_tools_discovered() - except Exception: - logger.exception("Tool discovery failed at MCP session start; proxy tools stay unavailable until it succeeds") - yield - - self._mcp = FastMCP(server_name, lifespan=lifespan) + self._mcp = FastMCP(server_name, lifespan=self._session_lifespan) self._mcp.settings.host = host self._mcp.settings.port = port self._mcp.settings.transport_security.enable_dns_rebinding_protection = False @@ -325,9 +309,15 @@ def run(self) -> None: else: anyio.run(server.serve) + # The three scopes a Caller is bound in, innermost first. Request: the served app binds the role the + # AgentEnv-Role header named. Session: mcp's Server.run enters the session lifespan once per inbound MCP + # session, in the session's server task before any handler task is spawned, so the key bound there is inherited + # by every tools/call of that session (and the role is the one its initialize request carried, which that task + # inherited). Process: child sessions and the time-trigger driver, outer to the MCP session manager's lifespan. + def _asgi_app(self): - """The served app: FastMCP's Starlette app with the AgentEnv-Role header bound per request, and the process - lifespan (child sessions, time-trigger driver) wrapped around the MCP session manager's.""" + """The served app: FastMCP's Starlette app with the request's caller bound around it, and the process + lifespan wrapped around the MCP session manager's.""" inner = self._mcp.streamable_http_app() inner.router.lifespan_context = self._process_lifespan(inner.router.lifespan_context) @@ -340,14 +330,25 @@ async def app(scope, receive, send): if name == _ROLE_HEADER_BYTES: role = value.decode(errors="ignore").strip() or DEFAULT_ROLE break - token = _role_var.set(role) + token = Caller(role=role, session=Caller.current().session).bind() try: await inner(scope, receive, send) finally: - _role_var.reset(token) + Caller.reset(token) return app + @asynccontextmanager + async def _session_lifespan(self, mcp: FastMCP): + """Once per inbound MCP session: the session's key joins the caller, and a connect gets to heal a failed + startup discovery (tools/list is synchronous and cannot). Nothing process-scoped lives here.""" + Caller(role=Caller.current().role, session=uuid4().hex).bind() + try: + await self._ensure_tools_discovered() + except Exception: + logger.exception("Tool discovery failed at MCP session start; proxy tools stay unavailable until it succeeds") + yield + def _process_lifespan(self, mcp_lifespan): """Once per process (uvicorn's lifespan), outer to the MCP session manager's: at shutdown the MCP sessions end first, then the driver stops and the scheduled trigger actions are cancelled, then the child sessions close.""" @@ -390,17 +391,17 @@ def _is_disabled(self, role: str, name: str) -> bool: return False def _install_role_filters(self) -> None: - """Wrap the FastMCP tool manager so MCP tools/list and tools/call honor _role_var.""" + """Wrap the FastMCP tool manager so MCP tools/list and tools/call honor the current caller's role.""" tm = self._mcp._tool_manager original_list_tools = tm.list_tools original_call_tool = tm.call_tool def filtered_list_tools(): - role = _role_var.get() + role = Caller.current().role return [t for t in original_list_tools() if not self._is_disabled(role, t.name)] async def filtered_call_tool(name, arguments, context=None, convert_result=False): - role = _role_var.get() + role = Caller.current().role if self._is_disabled(role, name): raise ToolError(f"Tool '{name}' is disabled for role '{role}'") native = name not in self._tool_server_urls @@ -843,11 +844,13 @@ def _child_transport_failed(self, sessions: dict[str, ClientSession], exc: httpx task.add_done_callback(lambda _: self._retiring.discard(id(sessions))) async def _call_child_tool(self, server_url: str, tool_name: str, arguments: dict, *, - meta: dict | None = None) -> CallToolResult: + caller: Caller | None = None) -> CallToolResult: """Every child tools/call (MCP proxy, /step, trigger engine): one shared generation, the caller's `_meta` - stamp, and one reconnect when the session turns out dead, retrying the call only if it never reached the - child (a call the transport lost in flight may already have written; the next call reopens). Discovery is - not redone on a reconnect; the registered proxies look their session up by URL.""" + stamp (none for an unattributed trigger action), and one reconnect when the session turns out dead, + retrying the call only if it never reached the child (a call the transport lost in flight may already have + written; the next call reopens). Discovery is not redone on a reconnect; the registered proxies look their + session up by URL.""" + meta = caller.meta() if caller is not None else None sessions = await self._ensure_child_sessions() closed = self._child_close # this generation's; no await since the lookup, so they agree try: @@ -1046,8 +1049,7 @@ async def _step_call_tool(self, body: dict, role: str) -> Response: }) try: async with self._tool_call_guard(): - result = await self._call_child_tool(server_url, tool_name, arguments, - meta={AGENT_ENV_ROLE_META_KEY: role}) + result = await self._call_child_tool(server_url, tool_name, arguments, caller=Caller(role=role)) changelog_id = self._query_changelog_id() await self._log_event({ @@ -1153,11 +1155,9 @@ def _create_and_register_proxy_tool(self, server_url: str, tool: MCPTool) -> Non input_schema = tool.inputSchema or {"type": "object", "properties": {}} async def proxy(**kwargs) -> CallToolResult: - # The session's own role and key, never anything the inbound request wrote: _role_var is what the - # disabled check and trigger attribution already keyed this call on. - meta = {AGENT_ENV_ROLE_META_KEY: _role_var.get()} - if caller_session := _caller_session_var.get(): - meta[AGENT_ENV_SESSION_META_KEY] = caller_session + # The session's own caller, never anything the inbound request wrote: the same value the disabled check + # and trigger attribution already keyed this call on. + caller = Caller.current() tool_call_event_id = await self._log_event({ "event_type": "tool_call", "tool_call": { @@ -1167,7 +1167,7 @@ async def proxy(**kwargs) -> CallToolResult: }) async with self._tool_call_guard(): - result = await self._call_child_tool(server_url, tool_name, kwargs, meta=meta) + result = await self._call_child_tool(server_url, tool_name, kwargs, caller=caller) changelog_id = self._query_changelog_id() await self._log_event({ diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py index 79c52993..b7519ceb 100644 --- a/src/agent_env/env/gateway/triggers.py +++ b/src/agent_env/env/gateway/triggers.py @@ -20,7 +20,8 @@ from mcp.types import CallToolResult, TextContent from .clock import ClockError, _advance, _iso, _parse_duration, _parse_rfc3339 -from .constants import AGENT_ENV_ROLE_META_KEY, DEFAULT_ROLE, TRIGGER_IN_FLIGHT_STATUSES, TRIGGER_STATUSES +from .caller import Caller +from .constants import DEFAULT_ROLE, TRIGGER_IN_FLIGHT_STATUSES, TRIGGER_STATUSES logger = logging.getLogger(__name__) @@ -1068,8 +1069,8 @@ async def _internal_call(self, tool_name: str, arguments: dict, role: str | None return result text = result if isinstance(result, str) else json.dumps(result, default=str) return CallToolResult(content=[TextContent(type="text", text=text)], isError=False) - meta = {AGENT_ENV_ROLE_META_KEY: role} if role is not None else None - return await gw._call_child_tool(server_url, tool_name, arguments, meta=meta) + caller = Caller(role=role) if role is not None else None + return await gw._call_child_tool(server_url, tool_name, arguments, caller=caller) async def _fire(self, tid: str, trig: dict, ctx: dict) -> None: """Run the actions for one provoking call, then drain any calls queued while firing (repeat diff --git a/tst/unit/env/gateway/caller_test.py b/tst/unit/env/gateway/caller_test.py new file mode 100644 index 00000000..82cbe21f --- /dev/null +++ b/tst/unit/env/gateway/caller_test.py @@ -0,0 +1,61 @@ +"""The one caller value: its binding, the `_meta` stamp it produces, and the wire names it shares with the protocol.""" +from __future__ import annotations + +import asyncio +import contextvars + +import agentenv_protocol +import pytest + +from agent_env.env.gateway.caller import Caller +from agent_env.env.gateway.constants import ( + AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_META_KEY, + AGENT_ENV_SESSION_META_KEY, + DEFAULT_ROLE, +) + + +def test_the_wire_names_are_the_protocols(): + """The gateway container ships without agentenv_protocol, so these are literals; they must stay the protocol's.""" + assert (AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE) == ( + agentenv_protocol.ROLE_HEADER, agentenv_protocol.ROLE_META_KEY, agentenv_protocol.SESSION_META_KEY, + agentenv_protocol.DEFAULT_ROLE) + + +def test_an_unforwarded_caller_is_stamped_default_with_no_session(): + assert Caller.current() == Caller(DEFAULT_ROLE, "") + assert Caller().meta() == {AGENT_ENV_ROLE_META_KEY: DEFAULT_ROLE}, "the role entry is what marks a call as proxied" + + +def test_the_session_rides_along_only_when_there_is_one(): + assert Caller("alice@example.com").meta() == {AGENT_ENV_ROLE_META_KEY: "alice@example.com"} + assert Caller("alice@example.com", "s1").meta() == { + AGENT_ENV_ROLE_META_KEY: "alice@example.com", AGENT_ENV_SESSION_META_KEY: "s1"} + + +def test_binding_is_per_context_and_restorable(): + token = Caller("bob@example.com", "s2").bind() + try: + assert Caller.current() == Caller("bob@example.com", "s2") + inner = contextvars.copy_context() + inner.run(Caller("carol@example.com").bind) + assert inner.run(Caller.current) == Caller("carol@example.com") + assert Caller.current() == Caller("bob@example.com", "s2"), "a copied context binds on its own" + finally: + Caller.reset(token) + assert Caller.current() == Caller() + + +@pytest.mark.asyncio +async def test_a_task_inherits_the_caller_bound_when_it_was_spawned(): + token = Caller("dave@example.com", "s3").bind() + try: + seen = await asyncio.create_task(_current_in_task()) + finally: + Caller.reset(token) + assert seen == Caller("dave@example.com", "s3") + + +async def _current_in_task() -> Caller: + return Caller.current() diff --git a/tst/unit/env/gateway/role_forwarding_test.py b/tst/unit/env/gateway/role_forwarding_test.py index e6e12bcf..a00ceb93 100644 --- a/tst/unit/env/gateway/role_forwarding_test.py +++ b/tst/unit/env/gateway/role_forwarding_test.py @@ -15,6 +15,8 @@ from agent_env.env.gateway import AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, InternalMCPServer from agent_env.env.gateway import gateway as gateway_module +from agent_env.env.gateway.caller import Caller +from agent_env.env.gateway.constants import DEFAULT_ROLE from agent_env.env.gateway.gateway import GATEWAY_TRAJECTORY_FILE, Gateway, _call_never_reached_child, _is_dead_session from agent_env.env.gateway.get_time import GET_TIME_TOOL_NAME @@ -67,10 +69,8 @@ def test_the_mcp_proxy_stamps_the_session_s_role_and_key_never_the_request_s(): def call(role=None, key=None): ctx = contextvars.copy_context() # what the ASGI wrapper and the per-session lifespan would have set - if role is not None: - ctx.run(gateway_module._role_var.set, role) - if key is not None: - ctx.run(gateway_module._caller_session_var.set, key) + if role is not None or key is not None: + ctx.run(Caller(role=role or DEFAULT_ROLE, session=key or "").bind) ctx.run(asyncio.run, gw._mcp._tool_manager.call_tool(TOOL.name, {})) call("alice@example.com", "s1") diff --git a/tst/unit/env/gateway/triggers_test.py b/tst/unit/env/gateway/triggers_test.py index f944683b..4f311436 100644 --- a/tst/unit/env/gateway/triggers_test.py +++ b/tst/unit/env/gateway/triggers_test.py @@ -410,7 +410,7 @@ def _step_gateway(): internal_mcp_servers=[], rest_proxy_urls={}) gw._tool_server_urls = {"gdocs_create_document": "http://svc/mcp"} - async def _fake_call_child_tool(server_url, tool_name, arguments, *, meta=None): + async def _fake_call_child_tool(server_url, tool_name, arguments, *, caller=None): return _result() gw._call_child_tool = _fake_call_child_tool # type: ignore[assignment]