refactor(gateway): one Caller value for role and session; the three binding scopes side by side

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())

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,16 +64,14 @@
     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
 from .triggers import TriggerEngine, TriggerError
 
 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({

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

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()

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")

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]