From 0f1b50365a71aa2925d0426885db869a8ac014e6 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 12:32:13 -0700 Subject: [PATCH 1/4] feat(gateway): forward the caller's role to child servers and share one child session generation per process Every child tools/call now carries MCP request _meta: the caller's role under "agentenv.io/role" (the gateway's own view of the caller, never copied from the inbound request) and, on the MCP path, a per-session key under "agentenv.io/session" so a server can keep per-agent session state once the gateway shares one connection to it. /step stamps the header role; a trigger tool action stamps a role only when it declares "as": "", validated at registration. The REST proxy drops every client AgentEnv-Role header and sets the gateway's own, and on a child's proxied MCP endpoint rewrites the role in the JSON-RPC _meta the same way, so the gateway is the trust boundary for identity without withdrawing the route the QC harness uses. Child sessions and the time-trigger driver become process-scoped. The per-session FastMCP lifespan used to open its own child sessions, overwrite the shared dict, and stop the driver when any inbound session ended, so a second agent's disconnect broke the first agent's calls and timers. One generation of child sessions is opened lazily under a lock, owned by one task, and shared by the MCP proxy, /step and triggers; a dead session (stream closed, session terminated after the child's idle timeout, transport gone mid-call) is reopened once and the call retried. Discovery is idempotent and bounded. mcp floor raised to 1.25, where ClientSession.call_tool takes meta. Co-Authored-By: Claude Fable 5.1 --- pyproject.toml | 2 +- src/agent_env/env/gateway/__init__.py | 4 + src/agent_env/env/gateway/constants.py | 4 + src/agent_env/env/gateway/gateway.py | 356 +++++++++++++----- src/agent_env/env/gateway/triggers.py | 48 ++- tst/unit/env/gateway/child_sessions_test.py | 261 +++++++++++++ tst/unit/env/gateway/import_graph_test.py | 9 +- .../env/gateway/process_scope_socket_test.py | 166 ++++++++ .../env/gateway/rest_proxy_stream_test.py | 89 +++++ tst/unit/env/gateway/role_forwarding_test.py | 127 +++++++ tst/unit/env/gateway/triggers_ctx_test.py | 2 +- tst/unit/env/gateway/triggers_test.py | 50 ++- uv.lock | 2 +- 13 files changed, 990 insertions(+), 130 deletions(-) create mode 100644 tst/unit/env/gateway/child_sessions_test.py create mode 100644 tst/unit/env/gateway/process_scope_socket_test.py create mode 100644 tst/unit/env/gateway/role_forwarding_test.py diff --git a/pyproject.toml b/pyproject.toml index 488b20f8..041030ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [ "cryptography>=42", # the local transfer CA "httpx>=0.27.0", "litellm>=1.96.0,<1.97", - "mcp>=1.0.0,<2", # 2.0.0 removed mcp.server.fastmcp, imported at module scope by env/gateway + "mcp>=1.25,<2", # same floor as env/gateway/requirements.txt and agentenv-protocol; 2.0.0 removed mcp.server.fastmcp, imported at module scope by env/gateway "packaging>=22", # evaluates a plugin's requirement markers with no extra "a2a-sdk==0.3.26", "pydantic>=2.12", # model_validate(extra=...) reads agent responses tolerantly diff --git a/src/agent_env/env/gateway/__init__.py b/src/agent_env/env/gateway/__init__.py index ff5c847e..3ca4778d 100644 --- a/src/agent_env/env/gateway/__init__.py +++ b/src/agent_env/env/gateway/__init__.py @@ -5,6 +5,8 @@ from .constants import ( AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_META_KEY, + AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE, GATEWAY_TRAJECTORY_FILE, TOOL_DISABLE_ACTION, @@ -27,6 +29,8 @@ class InternalMCPServer: __all__ = [ "AGENT_ENV_GATEWAY_MCP_PORT", "AGENT_ENV_ROLE_HEADER", + "AGENT_ENV_ROLE_META_KEY", + "AGENT_ENV_SESSION_META_KEY", "DEFAULT_ROLE", "GATEWAY_TRAJECTORY_FILE", "GatewayMode", diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 5a2256d0..d26271c7 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -96,6 +96,10 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: ) AGENT_ENV_ROLE_HEADER = "AgentEnv-Role" +# MCP request `_meta` keys the gateway stamps on every child tools/call: the caller's role, and on the MCP path a +# key naming the inbound session, so a child can keep per-caller state behind the gateway's one shared session. +AGENT_ENV_ROLE_META_KEY = "agentenv.io/role" +AGENT_ENV_SESSION_META_KEY = "agentenv.io/session" DEFAULT_ROLE = "default" WILDCARD = "*" TOOL_DISABLE_ACTION = "disable" diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 46b35760..7aad6dcf 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -8,6 +8,7 @@ import json import logging import os +import posixpath import socket import subprocess from contextlib import asynccontextmanager @@ -17,6 +18,7 @@ from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Any from urllib.parse import unquote +from uuid import uuid4 import anyio import httpx @@ -28,7 +30,8 @@ from mcp.server.fastmcp.exceptions import ToolError from mcp.server.fastmcp.tools import Tool as FastMCPTool from mcp.server.fastmcp.utilities.func_metadata import FuncMetadata, ArgModelBase -from mcp.types import CallToolResult, TextContent, Tool as MCPTool +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED, CallToolResult, TextContent, Tool as MCPTool from pydantic import create_model from starlette.requests import Request from starlette.responses import Response, StreamingResponse @@ -41,6 +44,8 @@ # Relative import — gateway runs as a standalone container with no agent_env.* from .constants import ( AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_META_KEY, + AGENT_ENV_SESSION_META_KEY, CARD_FETCH_TIMEOUT_S, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, @@ -66,12 +71,15 @@ 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: """Gateway that aggregates MCP tools and REST proxies for website backends.""" TOOL_CALL_TIMEOUT_S = 600.0 + CHILD_CLOSE_TIMEOUT_S = 3.0 # wait for the owner task to exit a retired child-session generation; past it, cancel it TRIGGER_BARRIER_TIMEOUT_S = 30.0 # a barrier may ask for less; past it the call is answered anyway # Must be >= the largest timeout any client can ask for, or the proxy severs a load the # client was still legitimately waiting on -- the ordering this module's constants file @@ -105,17 +113,20 @@ def __init__( if service_db_url: self._db_conn = psycopg2.connect(service_db_url) self._db_conn.autocommit = True - self._mcp_sessions: dict[str, ClientSession] = {} + # The process's one generation of child sessions, shared by the MCP proxy, /step and triggers. The client + # contexts are anyio task groups, so only the task that entered them may exit them: the owner task. + self._child_sessions: dict[str, ClientSession] = {} + self._child_owner: asyncio.Task | None = None + self._child_close: asyncio.Event | None = None + self._child_lock = asyncio.Lock() # open / retire / close of a generation (taken after _discover_lock) + self._discover_lock = asyncio.Lock() self._server_tools: dict[str, list[MCPTool]] = {} self._tools_discovered = False self._tool_server_urls: dict[str, str] = {} - self._step_sessions: dict[str, ClientSession] = {} - self._step_exit_stack: contextlib.AsyncExitStack | None = None self._tool_call_lock: asyncio.Lock | None = asyncio.Lock() if self.gateway_mode == GatewayMode.CONSISTENT else None self._role_rules: dict[str, dict[str, bool]] = {} self._role_rules_lock = asyncio.Lock() - self._init_lock = asyncio.Lock() self._event_counter = 0 self._event_lock = asyncio.Lock() Path(GATEWAY_TRAJECTORY_FILE).parent.mkdir(parents=True, exist_ok=True) @@ -124,35 +135,33 @@ def __init__( @asynccontextmanager async def lifespan(mcp: FastMCP): - async with contextlib.AsyncExitStack() as stack: - try: - self._mcp_sessions = await self._open_persistent_sessions(stack) - await self._discover_and_register_tools(self._mcp_sessions) - except Exception: - logger.exception("Failed to discover tools during MCP lifespan startup") - self._mcp_sessions = {} - self._server_tools = {} - self._tools_discovered = False - if self.website_urls: - @self._mcp.tool(name="list_website_urls", description="List available website URLs that can be browsed.") - async def list_website_urls() -> list[str]: - return list(self.website_urls.values()) - self._server_tools.setdefault("gateway", []).append( - MCPTool( - name="list_website_urls", - description="List available website URLs that can be browsed.", - inputSchema={"type": "object", "properties": {}}, - ) - ) - self._trigger_engine.start_driver() - stack.push_async_callback(self._trigger_engine.stop_driver) - yield + # 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.settings.host = host self._mcp.settings.port = port self._mcp.settings.transport_security.enable_dns_rebinding_protection = False self._trigger_engine = TriggerEngine(self) + if self.website_urls: + @self._mcp.tool(name="list_website_urls", description="List available website URLs that can be browsed.") + async def list_website_urls() -> list[str]: + return list(self.website_urls.values()) + self._server_tools.setdefault("gateway", []).append( + MCPTool( + name="list_website_urls", + description="List available website URLs that can be browsed.", + inputSchema={"type": "object", "properties": {}}, + ) + ) self._clock = Clock() self._env_get_time_url_cache: str | None = None self._install_role_filters() @@ -294,8 +303,27 @@ async def env_data_plane(request: Request) -> Response: def run(self) -> None: logger.info(f"Starting Gateway on {self.host}:{self.port} (mode={self.gateway_mode.value})") + config = uvicorn.Config( + self._asgi_app(), + host=self._mcp.settings.host, + port=self._mcp.settings.port, + log_level=self._mcp.settings.log_level.lower(), + ) + server = uvicorn.Server(config) + if self._resolve_i6pn_host(): + # On Modal's i6pn, serve both families: IPv4 for the orchestrator's HTTPS tunnel and + # IPv6 for backing servers reaching the gateway over the i6pn mesh. A lone `::` bind + # comes up IPv6-only on this image and starves the tunnel. + anyio.run(server.serve, [self._listen_socket(socket.AF_INET, "0.0.0.0"), + self._listen_socket(socket.AF_INET6, "::")]) + else: + anyio.run(server.serve) + + 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.""" inner = self._mcp.streamable_http_app() - role_header_bytes = AGENT_ENV_ROLE_HEADER.lower().encode() + inner.router.lifespan_context = self._process_lifespan(inner.router.lifespan_context) async def app(scope, receive, send): if scope["type"] != "http": @@ -303,7 +331,7 @@ async def app(scope, receive, send): return role = DEFAULT_ROLE for name, value in scope.get("headers", ()): - if name == role_header_bytes: + if name == _ROLE_HEADER_BYTES: role = value.decode(errors="ignore").strip() or DEFAULT_ROLE break token = _role_var.set(role) @@ -312,21 +340,27 @@ async def app(scope, receive, send): finally: _role_var.reset(token) - config = uvicorn.Config( - app, - host=self._mcp.settings.host, - port=self._mcp.settings.port, - log_level=self._mcp.settings.log_level.lower(), - ) - server = uvicorn.Server(config) - if self._resolve_i6pn_host(): - # On Modal's i6pn, serve both families: IPv4 for the orchestrator's HTTPS tunnel and - # IPv6 for backing servers reaching the gateway over the i6pn mesh. A lone `::` bind - # comes up IPv6-only on this image and starves the tunnel. - anyio.run(server.serve, [self._listen_socket(socket.AF_INET, "0.0.0.0"), - self._listen_socket(socket.AF_INET6, "::")]) - else: - anyio.run(server.serve) + return app + + 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, then the child sessions close.""" + @asynccontextmanager + async def lifespan(app): + try: + await self._ensure_tools_discovered() + except Exception: + # a raised startup makes uvicorn exit the process; discovery is retried on first use instead + logger.exception("Tool discovery failed at startup; retrying on first use") + self._trigger_engine.start_driver() + try: + async with mcp_lifespan(app): + yield + finally: + await self._trigger_engine.stop_driver() + await self._close_child_sessions() + + return lifespan def _listen_socket(self, family: int, host: str) -> socket.socket: sock = socket.socket(family, socket.SOCK_STREAM) @@ -450,18 +484,22 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] | base_url = url_map.get(service_name) if base_url is None: return self._json_response({"error": f"Unknown service: {service_name}"}, 404) + role = self._request_role(request) try: target_url = self._proxy_target_url(base_url, request) except ValueError as e: # an escaped slash inside the route prefix; refuse rather than forward a wrong path return self._json_response({"error": str(e)}, 400) + content = await request.body() # /svc and /website request bodies are small (JSON-RPC, forms) + headers = self._proxy_request_headers(request, role) + if (url_map is self.rest_proxy_urls and request.method == "POST" and service_name.startswith("mcp-") + and _is_child_mcp_path(request.path_params.get("path", ""))): + # The gateway is the trust boundary for identity: a child's proxied MCP endpoint stays reachable, but the + # role in its JSON-RPC `_meta` is the gateway's to write, like the header above. + content = _stamp_jsonrpc_role(content, role) + headers["content-length"] = str(len(content)) client = httpx.AsyncClient(timeout=self.REST_PROXY_TIMEOUT) - upstream_request = client.build_request( - request.method, - target_url, - content=await request.body(), # /svc and /website request bodies are small (JSON-RPC, forms) - headers=self._proxy_request_headers(request), - ) + upstream_request = client.build_request(request.method, target_url, content=content, headers=headers) try: upstream = await client.send(upstream_request, stream=True) # returns once the headers are in except Exception as e: @@ -485,9 +523,18 @@ def _proxy_target_url(base_url: str, request: Request) -> str: return url @staticmethod - def _proxy_request_headers(request: Request) -> httpx.Headers: - """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity.""" - headers = httpx.Headers([(k, v) for k, v in request.headers.raw if k.lower() not in _HOP_BY_HOP]) + def _request_role(request: Request) -> str: + """The role the gateway attributes a REST request to: the first AgentEnv-Role header, else the default role.""" + return request.headers.get(AGENT_ENV_ROLE_HEADER, "").strip() or DEFAULT_ROLE + + @staticmethod + def _proxy_request_headers(request: Request, role: str) -> httpx.Headers: + """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity. The role + header is the gateway's to set: the client's copies are dropped and the role the gateway attributed the + request to goes in their place, so the upstream never reads a value the gateway did not write.""" + headers = httpx.Headers([(k, v) for k, v in request.headers.raw + if k.lower() not in _HOP_BY_HOP and k.lower() != _ROLE_HEADER_BYTES]) + headers[AGENT_ENV_ROLE_HEADER] = role # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take. headers.setdefault("accept-encoding", "identity") return headers @@ -667,27 +714,106 @@ async def _open_persistent_sessions(self, exit_stack: contextlib.AsyncExitStack, raise return sessions + async def _ensure_child_sessions(self) -> dict[str, ClientSession]: + """The process's one generation of child sessions, opened on first need. Liveness is the owner task's, not + the dict's emptiness, so a gateway with no backing servers does not reopen on every call.""" + if self._child_owner is not None and not self._child_owner.done(): + return self._child_sessions + async with self._child_lock: + if self._child_owner is not None and not self._child_owner.done(): + return self._child_sessions + await self._close_child_sessions() # a dead or half-open generation + loop = asyncio.get_running_loop() + ready: asyncio.Future = loop.create_future() + close = asyncio.Event() + owner = loop.create_task(self._own_child_sessions(ready, close), name="gateway-child-sessions") + owner.add_done_callback(self._log_owner_exit) + try: + sessions = await ready + except BaseException: + close.set() # the owner exits whatever it opened (also a success that lands after we left) + raise + self._child_sessions, self._child_owner, self._child_close = sessions, owner, close + return sessions + + async def _own_child_sessions(self, ready: asyncio.Future, close: asyncio.Event) -> None: + """Holds one generation open: opens the client contexts, reports through `ready`, parks until `close` is set + (or it is cancelled) and exits them. Returns rather than raising on a failed open, which `ready` carries.""" + async with contextlib.AsyncExitStack() as stack: + try: + sessions = await self._open_persistent_sessions(stack) + except asyncio.CancelledError: + raise + except Exception as e: + if not ready.done(): + ready.set_exception(e) + return + if ready.cancelled(): + return # nobody is waiting any more; the stack closes what was opened + ready.set_result(sessions) + await close.wait() + + @staticmethod + def _log_owner_exit(owner: asyncio.Task) -> None: + """Retrieve a generation's exception when it dies on its own (a transport crash), so it is logged now rather + than as 'exception never retrieved' at garbage collection.""" + if not owner.cancelled() and owner.exception() is not None: + logger.warning("child session generation ended", exc_info=owner.exception()) + + async def _close_child_sessions(self) -> None: + """Retire the current generation and wait (bounded) for its owner to exit it. Callers hold _child_lock, + except the process lifespan, which runs after the last request.""" + owner, close = self._child_owner, self._child_close + self._child_sessions, self._child_owner, self._child_close = {}, None, None + if close is not None: + close.set() + if owner is not None and not owner.done(): + try: + await asyncio.wait_for(owner, timeout=self.CHILD_CLOSE_TIMEOUT_S) + except Exception: + logger.warning("child sessions did not close cleanly", exc_info=True) + + async def _reset_child_sessions(self, stale: dict[str, ClientSession]) -> None: + """Retire the generation a caller found dead; a concurrent caller that already did is a no-op.""" + async with self._child_lock: + if self._child_sessions is stale: + await self._close_child_sessions() + + async def _call_child_tool(self, server_url: str, tool_name: str, arguments: dict, *, + meta: dict | 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. Discovery is not redone on a reconnect; the + registered proxies look their session up by URL.""" + sessions = await self._ensure_child_sessions() + try: + return await asyncio.wait_for(sessions[server_url].call_tool(tool_name, arguments, meta=meta), + timeout=self.TOOL_CALL_TIMEOUT_S) + except Exception as e: + if not _is_dead_session(e): + raise + logger.warning(f"child session to {server_url} is dead ({e!r}), reconnecting...") + await self._reset_child_sessions(sessions) + sessions = await self._ensure_child_sessions() + return await asyncio.wait_for(sessions[server_url].call_tool(tool_name, arguments, meta=meta), + timeout=self.TOOL_CALL_TIMEOUT_S) + async def _ensure_tools_discovered(self) -> None: - """Lazy-init tool discovery for /step and /state endpoints.""" + """Idempotent discovery over the process's child sessions (startup, MCP session start, /step, /state, triggers).""" if self._tools_discovered: return - async with self._init_lock: + async with self._discover_lock: if self._tools_discovered: return + sessions = await self._ensure_child_sessions() try: - if self._step_exit_stack is None: - self._step_exit_stack = contextlib.AsyncExitStack() - if not self._step_sessions: - self._step_sessions = await self._open_persistent_sessions(self._step_exit_stack) - await self._discover_and_register_tools(self._step_sessions) - except Exception: - logger.exception("Tool discovery failed during lazy init") - self._server_tools = {} + await self._discover_and_register_tools(sessions) + except Exception as e: + logger.exception("Tool discovery failed") + for server in self.internal_mcp_servers: + self._server_tools.pop(server.name, None) # the gateway's own entry stays self._tools_discovered = False - self._step_sessions = {} - if self._step_exit_stack is not None: - await self._step_exit_stack.aclose() - self._step_exit_stack = None + if _is_dead_session(e): + await self._reset_child_sessions(sessions) raise async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) -> None: @@ -697,7 +823,8 @@ async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) for server in self.internal_mcp_servers: try: session = sessions[server.mcp_url] - tools_result = await session.list_tools() + # bounded: this runs under _discover_lock, which startup and every MCP connect wait on + tools_result = await asyncio.wait_for(session.list_tools(), timeout=CARD_FETCH_TIMEOUT_S) tools = tools_result.tools logger.info(f"Found {len(tools)} tool(s) from {server.name}") for tool in tools: @@ -717,7 +844,8 @@ async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) self._tools_discovered = True if self._mirror_get_time_descriptor(): logger.info(f"Re-mirrored '{GET_TIME_TOOL_NAME}' descriptor after tool discovery") - self._trigger_engine.invalidate_readonly_cache() + # list_website_urls sits in _server_tools from construction, so a memo may predate the backing tools + self._trigger_engine.invalidate_readonly_cache() @staticmethod def _json_response(data: dict | list, status_code: int) -> Response: @@ -781,7 +909,7 @@ async def _handle_step(self, request: Request) -> Response: except Exception: return self._json_response({"error": "Invalid JSON body"}, 400) action = body.get("action") - role = request.headers.get(AGENT_ENV_ROLE_HEADER, "").strip() or DEFAULT_ROLE + role = self._request_role(request) if action == "list_tools": return self._step_list_tools(role) elif action == "call_tool": @@ -803,14 +931,6 @@ def _step_list_tools(self, role: str) -> Response: ] return self._json_response({"tools": tools}, 200) - async def _get_step_session(self, server_url: str) -> ClientSession: - """Get or create a persistent MCP session for /step calls.""" - if not self._step_sessions: - if self._step_exit_stack is None: - self._step_exit_stack = contextlib.AsyncExitStack() - self._step_sessions = await self._open_persistent_sessions(self._step_exit_stack) - return self._step_sessions[server_url] - async def _step_call_tool(self, body: dict, role: str) -> Response: tool_name = body.get("tool_name") if not tool_name or not isinstance(tool_name, str): @@ -835,23 +955,8 @@ async def _step_call_tool(self, body: dict, role: str) -> Response: }) try: async with self._tool_call_guard(): - session = await self._get_step_session(server_url) - try: - result = await asyncio.wait_for( - session.call_tool(tool_name, arguments), - timeout=self.TOOL_CALL_TIMEOUT_S, - ) - except ClosedResourceError: - logger.warning(f"/step session to {server_url} closed, reconnecting...") - self._step_sessions.clear() - if self._step_exit_stack is not None: - await self._step_exit_stack.aclose() - self._step_exit_stack = None - session = await self._get_step_session(server_url) - result = await asyncio.wait_for( - session.call_tool(tool_name, arguments), - timeout=self.TOOL_CALL_TIMEOUT_S, - ) + result = await self._call_child_tool(server_url, tool_name, arguments, + meta={AGENT_ENV_ROLE_META_KEY: role}) changelog_id = self._query_changelog_id() await self._log_event({ @@ -957,6 +1062,11 @@ 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 tool_call_event_id = await self._log_event({ "event_type": "tool_call", "tool_call": { @@ -966,11 +1076,7 @@ async def proxy(**kwargs) -> CallToolResult: }) async with self._tool_call_guard(): - session = self._mcp_sessions[server_url] - result = await asyncio.wait_for( - session.call_tool(tool_name, kwargs), - timeout=self.TOOL_CALL_TIMEOUT_S, - ) + result = await self._call_child_tool(server_url, tool_name, kwargs, meta=meta) changelog_id = self._query_changelog_id() await self._log_event({ @@ -1090,6 +1196,54 @@ def model_dump_one_level(self) -> dict[str, Any]: b"te", b"trailer", b"upgrade", b"proxy-authorization", b"proxy-authenticate", }) +# The role header as it appears in the raw ASGI header list (lower-cased by the server). +_ROLE_HEADER_BYTES = AGENT_ENV_ROLE_HEADER.lower().encode() + + +def _is_dead_session(exc: BaseException) -> bool: + """Whether a child call failed because its session is gone: the stream closed under it, the child answered 404 + for the session id (idle expiry or a restart, which the client reports as code 32600), or the transport died + with the call in flight (CONNECTION_CLOSED).""" + if isinstance(exc, (ClosedResourceError, anyio.BrokenResourceError)): + return True + return isinstance(exc, McpError) and exc.error.code in (32600, CONNECTION_CLOSED) + + +def _is_child_mcp_path(decoded_rest: str) -> bool: + """Whether a /svc/{service}/{path} remainder reaches the backing server's MCP endpoint (`mcp` or anything under + it), judged after dot-segment normalisation, since httpx collapses `x/../mcp` to `/mcp` on the way upstream. + + >>> [_is_child_mcp_path(p) for p in ("mcp", "mcp/", "mcp/x", "x/../mcp", "./mcp", "/mcp", "mcpx", "agentenv", "")] + [True, True, True, True, True, True, False, False, False] + """ + normalized = posixpath.normpath("/" + decoded_rest.lstrip("/")) + return normalized == "/mcp" or normalized.startswith("/mcp/") + + +def _stamp_jsonrpc_role(body: bytes, role: str) -> bytes: + """A JSON-RPC body (one request or a batch) with every request's `params._meta` carrying the gateway's role + stamp in place of whatever the client wrote there; a body that is not JSON-RPC is returned as it came. + + >>> json.loads(_stamp_jsonrpc_role(b'{"method": "tools/call", "params": {"_meta": {"agentenv.io/role": "x", "agentenv.io/session": "s"}}}', "r"))["params"] + {'_meta': {'agentenv.io/role': 'r'}} + >>> _stamp_jsonrpc_role(b"not json", "r") + b'not json' + """ + try: + parsed = json.loads(body) + except ValueError: + return body + stamped = False + for item in parsed if isinstance(parsed, list) else [parsed]: + params = item.get("params") if isinstance(item, dict) else None + if isinstance(params, dict): + meta = dict(params["_meta"]) if isinstance(params.get("_meta"), dict) else {} + meta.pop(AGENT_ENV_SESSION_META_KEY, None) + meta[AGENT_ENV_ROLE_META_KEY] = role + params["_meta"] = meta + stamped = True + return json.dumps(parsed).encode() if stamped else body + class _UpstreamResponse(StreamingResponse): """Streams an httpx response through untouched and closes it when the response ends, iterated or not. diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py index a2b07f77..3e97b4b3 100644 --- a/src/agent_env/env/gateway/triggers.py +++ b/src/agent_env/env/gateway/triggers.py @@ -20,7 +20,7 @@ from mcp.types import CallToolResult, TextContent from .clock import ClockError, _advance, _iso, _parse_duration, _parse_rfc3339 -from .constants import TRIGGER_IN_FLIGHT_STATUSES, TRIGGER_STATUSES +from .constants import AGENT_ENV_ROLE_META_KEY, DEFAULT_ROLE, TRIGGER_IN_FLIGHT_STATUSES, TRIGGER_STATUSES logger = logging.getLogger(__name__) @@ -520,7 +520,8 @@ def _recompute_time_flag(self) -> None: self._time_trigger_gate.clear() def start_driver(self) -> None: - """Start the background poller (idempotent; requires a running loop).""" + """Start the background poller (idempotent; requires a running loop). The gateway's process lifespan starts + it once at startup.""" if self._driver_task is None or self._driver_task.done(): self._driver_task = asyncio.get_running_loop().create_task(self._drive_time()) @@ -538,7 +539,7 @@ async def _drive_time(self) -> None: logger.exception("clock time-driver tick failed") async def stop_driver(self) -> None: - """Cancel the background poller on gateway teardown (idempotent).""" + """Cancel the background poller at process shutdown (idempotent).""" task, self._driver_task = self._driver_task, None if task is not None: task.cancel() @@ -644,10 +645,12 @@ def _normalize_action_when(when: dict, tid: str) -> None: @staticmethod def _normalize_action(action: Any, w: str, has_ctx: bool, executor: dict | None) -> None: - """Validate one fire-action in place, plus its optional `verify` check.""" + """Validate one fire-action in place, plus its optional `verify` check and `as` role.""" atype = action.get("type") if isinstance(action, dict) else None if atype not in _ACTION_TYPES: raise TriggerError(f"{w}: type must be one of {_ACTION_TYPES}") + if "as" in action and atype != "tool": + raise TriggerError(f"{w}: 'as' is only valid on tool actions") if atype == "nl": if not isinstance(action.get("instruction"), str) or not action["instruction"]: raise TriggerError(f"{w}: instruction must be a non-empty string") @@ -664,6 +667,13 @@ def _normalize_action(action: Any, w: str, has_ctx: bool, executor: dict | None) if not isinstance(action.get("args", {}), dict): raise TriggerError(f"{w}: args must be an object") _validate_placeholders(action.get("args", {}), f"{w}.args", allow_ctx=has_ctx) + acting = action.get("as") + if acting is not None: + # the default role means "not forwarded" to the child, so naming it can only be a mistake + if not isinstance(acting, str) or not acting.strip() or acting == DEFAULT_ROLE: + raise TriggerError(f"{w}: 'as' must be a non-empty AgentEnv-Role other than {DEFAULT_ROLE!r}") + if _VAR_LOOSE_RE.search(acting): # never run through _template, so a placeholder would be sent verbatim + raise TriggerError(f"{w}: 'as' is not templated; name the role literally") else: if action.get("action") not in ("enable", "disable"): raise TriggerError(f"{w}: action must be 'enable' or 'disable'") @@ -1010,10 +1020,10 @@ async def _eval_state(self, tid: str, trig: dict, provoking: dict) -> None: if matched and trig["status"] == "armed": self._mark_detected(tid, trig, provoking) - async def _run_check(self, check: dict, ctx: dict | None = None) -> bool: + async def _run_check(self, check: dict, ctx: dict | None = None, role: str | None = None) -> bool: bindings: dict[str, Any] = {} for step in check["steps"]: - result = await self._internal_call(step["tool"], _template(step.get("args", {}), bindings, ctx)) + result = await self._internal_call(step["tool"], _template(step.get("args", {}), bindings, ctx), role=role) if getattr(result, "isError", False): return False raw_text = "\n".join(c.text for c in result.content if getattr(c, "type", None) == "text") @@ -1031,7 +1041,9 @@ async def _run_check(self, check: dict, ctx: dict | None = None) -> bool: return False return True - async def _internal_call(self, tool_name: str, arguments: dict) -> CallToolResult: + async def _internal_call(self, tool_name: str, arguments: dict, role: str | None = None) -> CallToolResult: + """A child call stamped as `role` when one is given (a tool action's `as`); the pin otherwise. Gateway-native + tools have no child and no identity, so the role is ignored there.""" gw = self._gw await gw._ensure_tools_discovered() server_url = gw._tool_server_urls.get(tool_name) @@ -1044,8 +1056,8 @@ async def _internal_call(self, tool_name: str, arguments: dict) -> CallToolResul return result text = result if isinstance(result, str) else json.dumps(result, default=str) return CallToolResult(content=[TextContent(type="text", text=text)], isError=False) - session = await gw._get_step_session(server_url) - return await asyncio.wait_for(session.call_tool(tool_name, arguments), timeout=gw.TOOL_CALL_TIMEOUT_S) + 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) 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 @@ -1099,33 +1111,36 @@ async def _run_action(self, tid: str, index: int, action: dict, ctx: dict) -> bo self._gw._apply_rule(action["role"], tool, action["action"] == "disable") detail = {"role": action["role"], "action": action["action"], "tools": tools} elif action["type"] == "tool": + role = action.get("as") + acting = {"as": role} if role else {} try: args = _template(action.get("args", {}), {}, ctx) except TemplateError as e: self._emit("action_failed", tid, action_index=index, detail=str(e)[:300], tool=action["tool"]) return False - result = await self._internal_call(action["tool"], args) + result = await self._internal_call(action["tool"], args, role=role) text = _result_text(result) failed = bool(getattr(result, "isError", False)) # The write itself, not just the trigger_fired marker: a Harbor grader has no timeline. await self._gw._log_event({ "event_type": "internal_tool_call", "source": "trigger_engine", "trigger_id": tid, - "action_index": index, "tool": action["tool"], "arguments": _echo_args(args), + "action_index": index, "tool": action["tool"], "arguments": _echo_args(args), "as": role, "ok": not failed, "result_len": len(text), "result_sha256": hashlib.sha256(text.encode("utf-8", "replace")).hexdigest()[:_ECHO_DIGEST_CHARS], "changelog_id_before": changelog_before, "changelog_id_after": self._gw._query_changelog_id()}) if failed: # The refusal text is the only account of why a mirror broke; counters cannot replace it. self._emit("action_failed", tid, action_index=index, detail=text[:300], - tool=action["tool"], args=_echo_args(args)) + tool=action["tool"], args=_echo_args(args), **acting) return False - detail = {"tool": action["tool"], "args": _echo_args(args)} + detail = {"tool": action["tool"], "args": _echo_args(args), **acting} else: if not await self._run_nl(tid, index, action, ctx): return False detail = {"nl": action["instruction"][:120]} if "verify" in action and action["type"] != "nl": - if not await self._run_verify(tid, index, action["verify"], ctx): + # only an action's own verify inherits its `as`; when.check state probes read as the pin + if not await self._run_verify(tid, index, action["verify"], ctx, role=action.get("as")): return False self._emit("action_ok", tid, action_index=index, detail=detail, changelog_id_before=changelog_before, changelog_id_after=self._gw._query_changelog_id()) @@ -1134,8 +1149,9 @@ async def _run_action(self, tid: str, index: int, action: dict, ctx: dict) -> bo self._emit("action_failed", tid, action_index=index, detail=str(e)[:300]) return False - async def _run_verify(self, tid: str, index: int, verify: dict, ctx: dict | None = None) -> bool: - if await self._run_check(verify, ctx): + async def _run_verify(self, tid: str, index: int, verify: dict, ctx: dict | None = None, + role: str | None = None) -> bool: + if await self._run_check(verify, ctx, role=role): self._emit("verify_ok", tid, action_index=index) return True self._emit("verify_failed", tid, action_index=index) diff --git a/tst/unit/env/gateway/child_sessions_test.py b/tst/unit/env/gateway/child_sessions_test.py new file mode 100644 index 00000000..01db6c08 --- /dev/null +++ b/tst/unit/env/gateway/child_sessions_test.py @@ -0,0 +1,261 @@ +"""One process-wide generation of child MCP sessions: opened on first need by one owner task (the client contexts are +anyio task groups, exited only by the task that entered them), shared by the MCP proxy, /step and the trigger engine, +reopened once when it turns out dead and never re-discovered; and the process lifespan that brackets it.""" +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import anyio +import pytest +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent, Tool as MCPTool, ToolAnnotations + +from agent_env.env.gateway import InternalMCPServer +from agent_env.env.gateway.gateway import Gateway + +URL = "http://backing/mcp" +TOOLS = [MCPTool(name="svc_thing", description="", inputSchema={"type": "object", "properties": {}}, + annotations=ToolAnnotations(readOnlyHint=True))] + + +def _result() -> CallToolResult: + return CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + + +class _FakeSession: + def __init__(self): + self.calls: list[str] = [] + self.failures: list[BaseException] = [] # raised by call_tool, in order + self.list_failures: list[BaseException] = [] # raised by list_tools, in order + + async def list_tools(self): + if self.list_failures: + raise self.list_failures.pop(0) + return SimpleNamespace(tools=TOOLS) + + async def call_tool(self, name, arguments, *args, meta=None, **kwargs): + self.calls.append(name) + if self.failures: + raise self.failures.pop(0) + return _result() + + +class _Opener: + """Stands in for _open_persistent_sessions: counts attempts and generations, records the entering/exiting task.""" + + def __init__(self, session: _FakeSession, *, delay: float = 0.0, failures=()): + self.session, self.delay, self.failures = session, delay, list(failures) + self.attempts = 0 + self.opened = 0 + self.entered: list[asyncio.Task] = [] + self.exited: list[asyncio.Task] = [] + + async def __call__(self, stack): + self.attempts += 1 + await asyncio.sleep(self.delay) + if self.failures: + raise self.failures.pop(0) + self.opened += 1 + self.entered.append(asyncio.current_task()) + + async def exited(): + self.exited.append(asyncio.current_task()) + + stack.push_async_callback(exited) + return {URL: self.session} + + +def _gateway(opener: _Opener, **kw) -> Gateway: + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[InternalMCPServer("svc", URL)], **kw) + gw._open_persistent_sessions = opener # type: ignore[method-assign] + return gw + + +def _explode(*a, **kw): + raise AssertionError("a completed pass must neither reopen nor rediscover") + + +@pytest.mark.asyncio +async def test_concurrent_first_use_opens_one_generation(): + opener = _Opener(_FakeSession(), delay=0.02) + gw = _gateway(opener) + try: + results = await asyncio.gather(*(gw._ensure_child_sessions() for _ in range(5))) + assert all(r is results[0] for r in results) and opener.opened == 1 + assert gw._child_sessions is results[0] and not gw._child_owner.done() + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_the_generation_is_exited_by_the_task_that_entered_it(): + opener = _Opener(_FakeSession()) + gw = _gateway(opener) + await gw._ensure_child_sessions() + owner = gw._child_owner + await gw._close_child_sessions() + assert opener.entered == [owner] and opener.exited == [owner] and owner is not asyncio.current_task() + assert owner.done() and gw._child_sessions == {} and gw._child_owner is None and gw._child_close is None + + +@pytest.mark.parametrize("death", [ + anyio.ClosedResourceError(), + anyio.BrokenResourceError(), + McpError(ErrorData(code=32600, message="Session terminated")), + McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), +], ids=["closed", "broken", "session_terminated", "connection_closed"]) +@pytest.mark.asyncio +async def test_a_dead_session_is_reopened_once_and_the_call_retried_without_rediscovery(death): + session = _FakeSession() + session.failures = [death] + opener = _Opener(session) + gw = _gateway(opener) + try: + await gw._ensure_tools_discovered() + first_owner, proxy = gw._child_owner, gw._mcp._tool_manager._tools["svc_thing"] + result = await gw._call_child_tool(URL, "svc_thing", {}) + assert not result.isError and session.calls == ["svc_thing", "svc_thing"] + assert opener.opened == 2 and first_owner.done() + assert gw._child_owner is not first_owner and not gw._child_owner.done() + assert gw._tools_discovered and gw._mcp._tool_manager._tools["svc_thing"] is proxy + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_any_other_error_is_the_caller_s_and_keeps_the_generation(): + session = _FakeSession() + session.failures = [McpError(ErrorData(code=-32602, message="Invalid params"))] + opener = _Opener(session) + gw = _gateway(opener) + try: + with pytest.raises(McpError, match="Invalid params"): + await gw._call_child_tool(URL, "svc_thing", {}) + assert session.calls == ["svc_thing"] and opener.opened == 1 and not gw._child_owner.done() + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_concurrent_dead_session_errors_reset_the_generation_once(): + session = _FakeSession() + opener = _Opener(session) + gw = _gateway(opener) + + async def call_tool(name, arguments, *args, meta=None, **kwargs): # every call on generation 1 dies + session.calls.append(name) + if opener.opened == 1: + raise anyio.ClosedResourceError() + return _result() + + session.call_tool = call_tool + try: + results = await asyncio.gather(gw._call_child_tool(URL, "a", {}), gw._call_child_tool(URL, "b", {})) + assert all(not r.isError for r in results) and opener.opened == 2 + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_a_failed_open_leaves_no_generation_and_is_retried_on_the_next_use(): + opener = _Opener(_FakeSession(), failures=[RuntimeError("refused")]) + gw = _gateway(opener) + with pytest.raises(RuntimeError, match="refused"): + await gw._ensure_child_sessions() + assert gw._child_owner is None and gw._child_sessions == {} + try: + assert (await gw._ensure_child_sessions())[URL] is opener.session + assert opener.attempts == 2 and opener.opened == 1 + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_a_failed_discovery_keeps_the_gateway_s_own_entry_and_heals_on_the_next_attempt(): + session = _FakeSession() + session.list_failures = [RuntimeError("boom")] + opener = _Opener(session) + gw = _gateway(opener, website_urls={"web": "http://web"}) + engine = gw._trigger_engine + # registered at construction, not per MCP session + assert [t.name for t in gw._server_tools["gateway"]] == ["list_website_urls"] + assert "list_website_urls" in gw._mcp._tool_manager._tools + assert engine._is_readonly("svc_thing") is False # a memo built before the backing tools are known + try: + with pytest.raises(RuntimeError, match="boom"): + await gw._ensure_tools_discovered() + assert gw._tools_discovered is False and "svc" not in gw._server_tools + assert [t.name for t in gw._server_tools["gateway"]] == ["list_website_urls"] + assert not gw._child_owner.done() # a failed tools/list is not a dead transport + + await gw._ensure_tools_discovered() + assert gw._tools_discovered and [t.name for t in gw._server_tools["svc"]] == ["svc_thing"] + assert engine._is_readonly("svc_thing") is True # discovery dropped the early memo + assert opener.opened == 1 + gw._open_persistent_sessions = _explode # type: ignore[method-assign] + await gw._ensure_tools_discovered() + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_discovery_on_a_dead_session_retires_the_generation(): + session = _FakeSession() + session.list_failures = [McpError(ErrorData(code=32600, message="Session terminated"))] + opener = _Opener(session) + gw = _gateway(opener) + try: + with pytest.raises(McpError): + await gw._ensure_tools_discovered() + assert gw._child_owner is None + await gw._ensure_tools_discovered() + assert gw._tools_discovered and opener.opened == 2 + finally: + await gw._close_child_sessions() + + +@pytest.mark.asyncio +async def test_the_process_lifespan_starts_the_driver_tolerates_a_failed_startup_discovery_and_unwinds(): + opener = _Opener(_FakeSession(), failures=[RuntimeError("refused")]) + gw = _gateway(opener) + app = gw._asgi_app() + scope = {"type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, "state": {}} + + async def run(release: asyncio.Event, sent: list[dict]) -> None: + pending = [{"type": "lifespan.startup"}] + + async def receive(): + if pending: + return pending.pop() + await release.wait() + return {"type": "lifespan.shutdown"} + + async def send(message): + sent.append(message) + + await app(scope, receive, send) + + release, sent = asyncio.Event(), [] + task = asyncio.create_task(run(release, sent)) + for _ in range(500): + if sent: + break + await asyncio.sleep(0.01) + assert sent == [{"type": "lifespan.startup.complete"}] # fail-open: a raised startup would exit uvicorn + driver = gw._trigger_engine._driver_task + assert driver is not None and not driver.done() + assert gw._child_owner is None and opener.attempts == 1 # the failed open left nothing behind + await gw._ensure_tools_discovered() # the next use heals it + assert gw._tools_discovered and not gw._child_owner.done() + release.set() + await asyncio.wait_for(task, 5) + assert sent[-1] == {"type": "lifespan.shutdown.complete"} + assert gw._trigger_engine._driver_task is None and gw._child_owner is None and gw._child_sessions == {} + + # the MCP session manager runs once per app, so a second lifespan fails its startup; the process state still unwinds + sent2: list[dict] = [] + with pytest.raises(RuntimeError, match="only be called once"): + await run(asyncio.Event(), sent2) + assert sent2[0]["type"] == "lifespan.startup.failed" + assert gw._trigger_engine._driver_task is None and gw._child_owner is None diff --git a/tst/unit/env/gateway/import_graph_test.py b/tst/unit/env/gateway/import_graph_test.py index 678e7fee..a0bf577a 100644 --- a/tst/unit/env/gateway/import_graph_test.py +++ b/tst/unit/env/gateway/import_graph_test.py @@ -13,6 +13,8 @@ class from ``gateway.gateway`` and the vocabulary from its root.""" import agent_env.env.gateway as gateway_pkg from agent_env.env.gateway import ( AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_META_KEY, + AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE, GATEWAY_TRAJECTORY_FILE, TOOL_DISABLE_ACTION, @@ -41,14 +43,19 @@ def test_importing_the_library_loads_neither_the_server_module_nor_its_driver(): def test_the_package_vocabulary_matches_the_server_module(): assert GatewayMode is server_module.GatewayMode - assert (AGENT_ENV_ROLE_HEADER, DEFAULT_ROLE, WILDCARD, TOOL_DISABLE_ACTION, TOOL_ENABLE_ACTION) == ( + assert (AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE, WILDCARD, + TOOL_DISABLE_ACTION, TOOL_ENABLE_ACTION) == ( server_module.AGENT_ENV_ROLE_HEADER, + server_module.AGENT_ENV_ROLE_META_KEY, + server_module.AGENT_ENV_SESSION_META_KEY, server_module.DEFAULT_ROLE, server_module.WILDCARD, server_module.TOOL_DISABLE_ACTION, server_module.TOOL_ENABLE_ACTION, ) assert GATEWAY_TRAJECTORY_FILE == server_module.GATEWAY_TRAJECTORY_FILE + # the literals a child with no agent_env dependency mirrors by hand + assert (AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY) == ("agentenv.io/role", "agentenv.io/session") def test_the_server_class_is_not_a_package_export(): diff --git a/tst/unit/env/gateway/process_scope_socket_test.py b/tst/unit/env/gateway/process_scope_socket_test.py new file mode 100644 index 00000000..85c0a3e4 --- /dev/null +++ b/tst/unit/env/gateway/process_scope_socket_test.py @@ -0,0 +1,166 @@ +"""Two agents on one gateway process, over real loopback sockets and the app `run()` serves: one child-session +generation shared by their MCP sessions, /step and the time-trigger driver; every child call stamped with its own +session's role and key; and one agent leaving taking nothing away from the other.""" +from __future__ import annotations + +import asyncio +import socket +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import httpx +import pytest +import uvicorn +from mcp import ClientSession +from mcp.client.streamable_http import streamable_http_client +from mcp.shared.exceptions import McpError +from mcp.types import CallToolResult, ErrorData, TextContent, Tool as MCPTool +from pytest_socket import enable_socket + +from agent_env.env.gateway import ( + AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_META_KEY, + AGENT_ENV_SESSION_META_KEY, + InternalMCPServer, +) +from agent_env.env.gateway.gateway import Gateway + +URL = "http://backing/mcp" +T0 = "2026-01-01T00:00:00Z" + + +@pytest.fixture(autouse=True) +def _disable_network(): + """Override the suite-wide socket block (tst/unit/conftest.py) for this module: loopback is the point here.""" + enable_socket() + yield + + +class _RecordingSession: + def __init__(self): + self.calls: list[tuple[str, dict, dict | None]] = [] + self.failures: list[BaseException] = [] + + async def list_tools(self): + return SimpleNamespace(tools=[MCPTool(name="svc_thing", description="", + inputSchema={"type": "object", "properties": {"v": {"type": "string"}}})]) + + async def call_tool(self, name, arguments, *args, meta=None, **kwargs): + self.calls.append((name, dict(arguments or {}), meta)) + if self.failures: + raise self.failures.pop(0) + return CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + + +async def _serve(app) -> tuple[uvicorn.Server, str]: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(("127.0.0.1", 0)) + # lifespan on (the default): startup runs the gateway's process lifespan, shutdown unwinds it + config = uvicorn.Config(app, log_level="critical", access_log=False, timeout_graceful_shutdown=3) + config.load() + server = uvicorn.Server(config) + server.lifespan = config.lifespan_class(config) + await server.startup(sockets=[sock]) + return server, f"http://127.0.0.1:{sock.getsockname()[1]}" + + +@asynccontextmanager +async def _stack(): + """The real gateway app on a loopback uvicorn server, its one backing server a recording fake session.""" + child = _RecordingSession() + opened: list[asyncio.Task] = [] + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[InternalMCPServer("svc", URL)], + website_urls={"web": "http://web"}) + + async def _open(stack): + opened.append(asyncio.current_task()) + return {URL: child} + + gw._open_persistent_sessions = _open # type: ignore[method-assign] + gw._trigger_engine._driver_interval = 0.01 + server, url = await _serve(gw._asgi_app()) + try: + yield gw, url, child, opened + finally: + await server.shutdown() + + +@asynccontextmanager +async def _client(url: str, role: str | None = None): + headers = {AGENT_ENV_ROLE_HEADER: role} if role else {} + async with ( + httpx.AsyncClient(headers=headers, timeout=httpx.Timeout(10, read=60)) as http, + streamable_http_client(f"{url}/mcp", http_client=http) as (read, write, _), + ClientSession(read, write) as session, + ): + await session.initialize() + yield session + + +@pytest.mark.asyncio +async def test_two_mcp_clients_share_one_generation_and_one_leaving_keeps_the_other_working(): + async with _stack() as (gw, url, child, opened): + owner0 = gw._child_owner + assert owner0 is not None and not owner0.done() and len(opened) == 1 # opened once, at startup + forged = {AGENT_ENV_ROLE_META_KEY: "victim@example.com", AGENT_ENV_SESSION_META_KEY: "forged"} + async with _client(url, "alice@example.com") as alice: + async with _client(url, "bob@example.com") as bob: + assert not (await alice.call_tool("svc_thing", {"v": "a1"}, meta=forged)).isError + assert not (await bob.call_tool("svc_thing", {"v": "b1"})).isError + # bob's DELETE ended his server session, which used to close alice's child sessions with it + assert not (await alice.call_tool("svc_thing", {"v": "a2"})).isError + async with httpx.AsyncClient() as http: + step = await http.post(f"{url}/step", headers={AGENT_ENV_ROLE_HEADER: "grader"}, + json={"action": "call_tool", "tool_name": "svc_thing", "arguments": {"v": "s1"}}) + assert step.status_code == 200, step.text + for _ in range(3): + async with _client(url): + pass + state = (await http.get(f"{url}/state")).json() + metas = [meta for _, _, meta in child.calls] + assert [m[AGENT_ENV_ROLE_META_KEY] for m in metas] == ["alice@example.com", "bob@example.com", + "alice@example.com", "grader"] + keys = [m.get(AGENT_ENV_SESSION_META_KEY) for m in metas] + assert keys[0] == keys[2] != keys[1] and keys[3] is None # one key per inbound MCP session; /step has none + assert "forged" not in keys and all(len(k) == 32 for k in keys[:3]) + assert gw._child_owner is owner0 and len(opened) == 1 + gateway_entry = next(s for s in state["mcp_servers"] if s["name"] == "gateway") + assert [t["name"] for t in gateway_entry["tools"]].count("list_website_urls") == 1 # not once per session + assert gw._child_owner is None and gw._trigger_engine._driver_task is None # process shutdown closed both + + +@pytest.mark.asyncio +async def test_an_armed_time_trigger_keeps_firing_after_a_client_leaves(): + async with _stack() as (gw, url, child, opened): + async with httpx.AsyncClient() as http: + r = await http.put(f"{url}/clock/set-time", json={"virtual_time": T0, "virtual_seconds_per_real_second": 3600}) + assert r.status_code == 200, r.text + r = await http.post(f"{url}/triggers/register", json={"watch_roles": ["default"], "triggers": [ + {"id": "rec", "when": {"type": "time", "every": "PT10M"}, + "actions": [{"type": "tool", "tool": "svc_thing", "args": {"v": "tick"}, "as": "carol@example.com"}]}]}) + assert r.status_code == 200, r.text + async with _client(url, "bob@example.com") as bob: + assert not (await bob.call_tool("svc_thing", {"v": "b1"})).isError + await asyncio.sleep(0.3) # 18 virtual minutes: at least one PT10M mark, after bob has left + state = (await http.get(f"{url}/triggers/state")).json() + rec = next(t for t in state["triggers"] if t["id"] == "rec") + assert rec["fire_count"] >= 1 and rec["status"] == "armed", state["events"][-5:] + driver = gw._trigger_engine._driver_task + assert driver is not None and not driver.done() + ticks = [meta for _, args, meta in child.calls if args == {"v": "tick"}] + assert ticks and all(meta == {AGENT_ENV_ROLE_META_KEY: "carol@example.com"} for meta in ticks) + assert len(opened) == 1 + assert gw._trigger_engine._driver_task is None and gw._child_owner is None and gw._child_sessions == {} + + +@pytest.mark.asyncio +async def test_a_dead_child_session_is_reopened_once_under_a_live_mcp_client(): + async with _stack() as (gw, url, child, opened): + async with _client(url, "alice@example.com") as alice: + assert not (await alice.call_tool("svc_thing", {"v": "1"})).isError + owner0 = gw._child_owner + child.failures.append(McpError(ErrorData(code=32600, message="Session terminated"))) # idle expiry / restart + assert not (await alice.call_tool("svc_thing", {"v": "2"})).isError # alice's own session was never dropped + assert len(opened) == 2 and owner0.done() and gw._child_owner is not owner0 and not gw._child_owner.done() + assert [args for _, args, _ in child.calls] == [{"v": "1"}, {"v": "2"}, {"v": "2"}] diff --git a/tst/unit/env/gateway/rest_proxy_stream_test.py b/tst/unit/env/gateway/rest_proxy_stream_test.py index 4b4e151f..85621f87 100644 --- a/tst/unit/env/gateway/rest_proxy_stream_test.py +++ b/tst/unit/env/gateway/rest_proxy_stream_test.py @@ -8,11 +8,13 @@ import asyncio import doctest +import json from urllib.parse import unquote, urlsplit import httpx import pytest +from agent_env.env.gateway import AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY from agent_env.env.gateway import gateway as gateway_module from agent_env.env.gateway.gateway import Gateway @@ -385,3 +387,90 @@ async def test_connect_timeout_maps_to_the_502_body_naming_the_error(monkeypatch up.install(monkeypatch) status, _, chunks = await _call(_gateway(), "/svc/mcp-svc/x") assert status == 502 and b"ConnectTimeout" in b"".join(chunks) + + +# --- The role header and the JSON-RPC `_meta` are the gateway's to write, not the client's --- + + +@pytest.mark.asyncio +async def test_the_client_s_role_headers_are_replaced_by_the_gateway_s_attribution(monkeypatch): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + status, _, _ = await _call(_gateway(), "/svc/mcp-svc/x", raw_headers=[ + (b"agentenv-role", b"alice@example.com"), (b"AgentEnv-Role", b"bob@example.com"), (b"cookie", b"a=1")]) + assert status == 200 + assert up.requests[0].headers.get_list("agentenv-role") == ["alice@example.com"] # first value wins; one entry survives + assert up.requests[0].headers["cookie"] == "a=1" + + +@pytest.mark.asyncio +async def test_an_absent_or_blank_role_header_is_the_default_role_as_step_attributes_it(monkeypatch): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + gw = _gateway() + await _call(gw, "/svc/mcp-svc/x") + await _call(gw, "/svc/mcp-svc/x", headers={"AgentEnv-Role": " "}) + assert [r.headers.get_list("agentenv-role") for r in up.requests] == [["default"], ["default"]] + seen: list[str] = [] + monkeypatch.setattr(gw, "_step_list_tools", lambda role: (seen.append(role), gw._json_response({"tools": []}, 200))[1]) + try: + status, _, _ = await _call(gw, "/step", method="POST", body=b'{"action": "list_tools"}', headers={"AgentEnv-Role": " "}) + finally: + await gw._close_child_sessions() + assert status == 200 and seen == ["default"] + + +@pytest.mark.asyncio +async def test_the_website_proxy_stamps_the_role_too(monkeypatch): + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/assets/app.js", headers={"AgentEnv-Role": "viewer"}) + await _call(gw, "/website/web") + assert [r.headers.get_list("agentenv-role") for r in up.requests] == [["viewer"], ["default"]] + + +def _rpc(**params) -> dict: + return {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": params} + + +@pytest.mark.asyncio +async def test_a_proxied_child_mcp_post_carries_the_gateway_s_role_in_every_request_s_meta(monkeypatch): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + gw = _gateway() + forged = {AGENT_ENV_ROLE_META_KEY: "victim@example.com", AGENT_ENV_SESSION_META_KEY: "forged"} + one = json.dumps(_rpc(name="t", arguments={}, _meta=forged)).encode() + batch = json.dumps([_rpc(name="t", arguments={}), {"jsonrpc": "2.0", "method": "notifications/initialized"}]).encode() + for path, body in [("/svc/mcp-svc/mcp", one), ("/svc/mcp-svc/mcp/", batch), ("/svc/mcp-svc/x/../mcp", one)]: + status, _, _ = await _call(gw, path, method="POST", body=body, headers={ + "AgentEnv-Role": "alice@example.com", "content-type": "application/json", "content-length": str(len(body))}) + assert status == 200 + sent = [json.loads(r.content) for r in up.requests] + stamp = {AGENT_ENV_ROLE_META_KEY: "alice@example.com"} + assert sent[0]["params"]["_meta"] == stamp # the forged role and session key are gone + assert sent[1][0]["params"]["_meta"] == stamp and "params" not in sent[1][1] # each request of a batch + assert sent[2]["params"]["_meta"] == stamp and up.requests[2].url.path == "/mcp" # dot segments collapse upstream + assert all(r.headers["content-length"] == str(len(r.content)) for r in up.requests) # re-serialised, re-measured + assert all(r.headers.get_list("agentenv-role") == ["alice@example.com"] for r in up.requests) + + +@pytest.mark.asyncio +async def test_bodies_that_are_not_a_child_mcp_request_are_forwarded_untouched(monkeypatch): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + rest_proxy_urls={"mcp-svc": UPSTREAM, "backend": UPSTREAM}) + rpc = json.dumps(_rpc(name="t", arguments={}, _meta={AGENT_ENV_ROLE_META_KEY: "x"})).encode() + cases = [ + ("/svc/mcp-svc/mcp", "POST", b"not json"), # unparsable: as it came + ("/svc/mcp-svc/mcp", "GET", b""), # not a POST + ("/svc/mcp-svc/mcpx", "POST", rpc), # not the MCP endpoint + ("/svc/mcp-svc/agentenv", "POST", rpc), # the data plane keeps its body + ("/svc/backend/mcp", "POST", rpc), # not an mcp-* service + ("/svc/mcp-svc/mcp", "POST", b'{"jsonrpc": "2.0", "id": 1, "result": {}}'), # no params: nothing to stamp + ] + for path, method, body in cases: + status, _, _ = await _call(gw, path, method=method, body=body, headers={"AgentEnv-Role": "alice@example.com"}) + assert status == 200, path + assert [r.content for r in up.requests] == [body for _, _, body in cases] diff --git a/tst/unit/env/gateway/role_forwarding_test.py b/tst/unit/env/gateway/role_forwarding_test.py new file mode 100644 index 00000000..dd069dff --- /dev/null +++ b/tst/unit/env/gateway/role_forwarding_test.py @@ -0,0 +1,127 @@ +"""Who is calling rides on every child tools/call as MCP request `_meta`, written by the gateway from what it +attributed the caller itself: the inbound MCP session's role and key on the proxy path, the AgentEnv-Role header on +/step, a tool action's literal `as` in the trigger engine. Nothing a caller wrote in its own `_meta` is copied.""" +from __future__ import annotations + +import asyncio +import contextvars +import json +from types import SimpleNamespace + +import anyio +import pytest +from mcp.shared.exceptions import McpError +from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent, Tool as MCPTool + +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.gateway import GATEWAY_TRAJECTORY_FILE, Gateway, _is_dead_session +from agent_env.env.gateway.get_time import GET_TIME_TOOL_NAME + +URL = "http://backing/mcp" +TOOL = MCPTool(name="svc_thing", description="", inputSchema={"type": "object", "properties": {}}) + + +class _RecordingSession: + """Stands in for the child's ClientSession: records each call's `_meta`, failing as told first.""" + + def __init__(self, failures=()): + self.calls: list[tuple[str, dict, dict | None]] = [] + self.failures = list(failures) + + async def list_tools(self): + return SimpleNamespace(tools=[TOOL]) + + async def call_tool(self, name, arguments, *args, meta=None, **kwargs): + self.calls.append((name, dict(arguments or {}), meta)) + if self.failures: + raise self.failures.pop(0) + return CallToolResult(content=[TextContent(type="text", text="ok")], isError=False) + + +def _gateway(session: _RecordingSession) -> Gateway: + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[InternalMCPServer("svc", URL)]) + + async def _open(stack): + return {URL: session} + + gw._open_persistent_sessions = _open # type: ignore[method-assign] + return gw + + +def _metas(session: _RecordingSession) -> list[dict | None]: + return [meta for _, _, meta in session.calls] + + +def _trajectory(gw: Gateway) -> list[dict]: + gw._trajectory_file.flush() + with open(GATEWAY_TRAJECTORY_FILE) as f: + return [json.loads(line) for line in f] + + +def test_the_mcp_proxy_stamps_the_session_s_role_and_key_never_the_request_s(): + session = _RecordingSession() + gw = _gateway(session) + gw._create_and_register_proxy_tool(URL, TOOL) + gw._tool_server_urls[TOOL.name] = URL + + 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) + ctx.run(asyncio.run, gw._mcp._tool_manager.call_tool(TOOL.name, {})) + + call("alice@example.com", "s1") + call() + assert _metas(session) == [ + {AGENT_ENV_ROLE_META_KEY: "alice@example.com", AGENT_ENV_SESSION_META_KEY: "s1"}, + {AGENT_ENV_ROLE_META_KEY: "default"}, # stamped for the default role too; no key outside an MCP session + ] + + +@pytest.mark.asyncio +async def test_step_stamps_the_header_role_in_meta_including_on_the_reconnect(): + session = _RecordingSession(failures=[anyio.ClosedResourceError()]) + gw = _gateway(session) + gw._tool_server_urls[TOOL.name] = URL + before = len(_trajectory(gw)) + try: + resp = await gw._step_call_tool({"tool_name": TOOL.name, "arguments": {"a": 1}}, "bob@example.com") + finally: + await gw._close_child_sessions() + assert resp.status_code == 200 + assert _metas(session) == [{AGENT_ENV_ROLE_META_KEY: "bob@example.com"}] * 2 # the retry is stamped too + calls = [e for e in _trajectory(gw)[before:] if e["event_type"] == "tool_call"] + assert len(calls) == 1 and calls[0]["tool_call"]["function_name"] == TOOL.name # logged once, not per attempt + + +@pytest.mark.asyncio +async def test_trigger_internal_calls_are_stamped_only_when_an_as_role_is_given(): + session = _RecordingSession() + gw = _gateway(session) + engine = gw._trigger_engine + try: + await engine._internal_call(TOOL.name, {}, role="carol@example.com") + await engine._internal_call(TOOL.name, {}) + gw._clock.set_time("2019-06-15T12:00:00Z", 1) + gw._register_get_time_tool() + native = await engine._internal_call(GET_TIME_TOOL_NAME, {}, role="carol@example.com") + finally: + await gw._close_child_sessions() + assert _metas(session) == [{AGENT_ENV_ROLE_META_KEY: "carol@example.com"}, None] + assert "current_time" in native.content[0].text # gateway-native: no child, no identity, no session touched + + +@pytest.mark.parametrize("exc,dead", [ + (anyio.ClosedResourceError(), True), + (anyio.BrokenResourceError(), True), + (McpError(ErrorData(code=32600, message="Session terminated")), True), + (McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), True), + (McpError(ErrorData(code=-32602, message="Invalid params")), False), + (asyncio.TimeoutError(), False), + (ValueError("x"), False), +], ids=["closed", "broken", "session_terminated", "connection_closed", "invalid_params", "timeout", "other"]) +def test_the_dead_session_predicate(exc, dead): + assert _is_dead_session(exc) is dead diff --git a/tst/unit/env/gateway/triggers_ctx_test.py b/tst/unit/env/gateway/triggers_ctx_test.py index 3209625a..903fafc7 100644 --- a/tst/unit/env/gateway/triggers_ctx_test.py +++ b/tst/unit/env/gateway/triggers_ctx_test.py @@ -55,7 +55,7 @@ def _make_engine(monkeypatch, responses: dict, known_tools=None, delays: dict | gw = _FakeGateway(responses, known_tools) delays = delays or {} - async def fake_internal_call(self, tool_name, arguments): + async def fake_internal_call(self, tool_name, arguments, role=None): gw.calls.append((tool_name, dict(arguments) if isinstance(arguments, dict) else arguments)) if tool_name in delays: await asyncio.sleep(delays[tool_name]) diff --git a/tst/unit/env/gateway/triggers_test.py b/tst/unit/env/gateway/triggers_test.py index 7f16d956..2ae6568b 100644 --- a/tst/unit/env/gateway/triggers_test.py +++ b/tst/unit/env/gateway/triggers_test.py @@ -15,6 +15,8 @@ class _FakeGateway: def __init__(self, responses: dict) -> None: self.responses = responses self.calls: list[tuple[str, dict]] = [] + self.roles: list[tuple[str, str | None]] = [] # the `role` each _internal_call carried + self.events: list[dict] = [] self.rules: list[tuple[str, str, bool]] = [] self._role_rules_lock = asyncio.Lock() self._server_tools: dict = {} @@ -29,7 +31,8 @@ def _query_changelog_id(self): return 42 async def _log_event(self, event): - return "event_test" + self.events.append(dict(event)) + return f"event_{len(self.events)}" async def _ensure_tools_discovered(self): pass @@ -53,8 +56,9 @@ def engine(monkeypatch): } gw = _FakeGateway(responses) - async def fake_internal_call(self, tool_name, arguments): + async def fake_internal_call(self, tool_name, arguments, role=None): gw.calls.append((tool_name, dict(arguments))) + gw.roles.append((tool_name, role)) fn = gw.responses.get(tool_name) if fn is None: raise RuntimeError(f"unknown tool {tool_name}") @@ -113,6 +117,16 @@ def _registration(**overrides): ({"executor": {"a2a_url": "http://e"}, "triggers": []}, "executor.role must be"), ({"watch_roles": ["default", "executor"], "executor": {"a2a_url": "http://e", "role": "executor"}, "triggers": []}, "must not be in watch_roles"), + *[({"triggers": [{"id": "x", "when": {"type": "action", "tool": "t"}, + "actions": [{"type": "tool", "tool": "t2", "as": acting}]}]}, fragment) + for acting, fragment in [("", "'as' must be a non-empty"), ("default", "'as' must be a non-empty"), + (5, "'as' must be a non-empty"), ("${args.owner}", "not templated")]], + ({"triggers": [{"id": "x", "when": {"type": "action", "tool": "t"}, + "actions": [{"type": "permission", "action": "enable", "role": "r", "tools": ["t"], "as": "x"}]}]}, + "only valid on tool actions"), + ({"executor": {"a2a_url": "http://e", "role": "executor"}, + "triggers": [{"id": "x", "when": {"type": "action", "tool": "t"}, + "actions": [{"type": "nl", "instruction": "hi", "as": "x"}]}]}, "only valid on tool actions"), ]) def test_registration_fails_loud(engine, bad, fragment): with pytest.raises(TriggerError, match=".*"): @@ -253,6 +267,28 @@ async def test_state_trigger_pipeline_tool_action_and_verify(engine): assert any(e["kind"] == "verify_ok" for e in engine._events) +@pytest.mark.asyncio +async def test_a_tool_action_s_as_role_reaches_its_call_and_verify_but_not_the_state_check(engine): + reg = _registration() + rate = reg["triggers"][1] + rate["actions"][0]["as"] = "carol@example.com" + rate["actions"].append({"type": "tool", "tool": "slack_send_message", "args": {"channel": "#c", "text": "plain"}}) + engine.register(reg) + engine._test_state["doc_text"] = "Budget Fit: fine." + engine.on_tool_call("default", "gdocs_batch_update", {"documentId": "doc-1"}, _result()) + await asyncio.sleep(0.3) + assert engine._triggers["rate"]["status"] == "fired" + gw = engine._test_gw + assert all(role is None for tool, role in gw.roles if tool.startswith("gdocs_")) # when.check reads as the pin + assert [role for tool, role in gw.roles if tool == "slack_send_message"] == ["carol@example.com", None] + assert [role for tool, role in gw.roles if tool == "slack_conversations_history"] == ["carol@example.com"] + oks = [e for e in engine._events if e["kind"] == "action_ok"] + assert [e["detail"].get("as") for e in oks] == ["carol@example.com", None] + assert any(e["kind"] == "verify_ok" for e in engine._events) + internal = [e for e in gw.events if e["event_type"] == "internal_tool_call"] + assert [(e["as"], e["ok"]) for e in internal] == [("carol@example.com", True), (None, True)] + + @pytest.mark.asyncio async def test_sensor_trigger_and_state_shape(engine): engine.register(_registration()) @@ -348,14 +384,10 @@ def _step_gateway(): internal_mcp_servers=[], rest_proxy_urls={}) gw._tool_server_urls = {"gdocs_create_document": "http://svc/mcp"} - class _FakeSession: - async def call_tool(self, name, arguments): - return _result() - - async def _fake_get_step_session(server_url): - return _FakeSession() + async def _fake_call_child_tool(server_url, tool_name, arguments, *, meta=None): + return _result() - gw._get_step_session = _fake_get_step_session # type: ignore[assignment] + gw._call_child_tool = _fake_call_child_tool # type: ignore[assignment] return gw diff --git a/uv.lock b/uv.lock index ba7cc51d..38c4a29a 100644 --- a/uv.lock +++ b/uv.lock @@ -125,7 +125,7 @@ requires-dist = [ { name = "griffelib", marker = "extra == 'dev'", specifier = "==2.3.0" }, { name = "httpx", specifier = ">=0.27.0" }, { name = "litellm", specifier = ">=1.96.0,<1.97" }, - { name = "mcp", specifier = ">=1.0.0,<2" }, + { name = "mcp", specifier = ">=1.25,<2" }, { name = "modal", specifier = ">=1.5.2.dev22" }, { name = "moto", marker = "extra == 'dev'", specifier = ">=5.0.0" }, { name = "packaging", specifier = ">=22" }, From 7da8d9d5a87afa9a82e268203e1e73943873ebcf Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 13:22:04 -0700 Subject: [PATCH 2/4] fix(gateway): retry only calls that never reached the child, bound startup, cancel scheduled actions at shutdown, normalise the action role A call the transport lost in flight may already have performed its write, so a dead session now always retires the generation but the call is retried only when the stream closed before the request was written or the child rejected the session id. Startup discovery and each child initialize are bounded, so a child that accepts the connection and never answers cannot hold the whole gateway back; discovery fails open and is retried on first use. Shutdown cancels the scheduled trigger actions before the child sessions close and refuses to open a new generation afterwards. A tool action's "as" role is stripped once at registration and forwarded exactly as the header path would. Co-Authored-By: Claude Fable 5.1 --- src/agent_env/env/gateway/gateway.py | 38 ++++- src/agent_env/env/gateway/triggers.py | 18 +- tst/unit/env/gateway/child_sessions_test.py | 168 ++++++++++++++++--- tst/unit/env/gateway/role_forwarding_test.py | 21 +-- tst/unit/env/gateway/triggers_test.py | 28 +++- 5 files changed, 224 insertions(+), 49 deletions(-) diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 7aad6dcf..c6858483 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -79,6 +79,7 @@ class Gateway: """Gateway that aggregates MCP tools and REST proxies for website backends.""" TOOL_CALL_TIMEOUT_S = 600.0 + STARTUP_DISCOVERY_TIMEOUT_S = 60.0 # startup discovery past it fails open, like a failed one; uvicorn must come up CHILD_CLOSE_TIMEOUT_S = 3.0 # wait for the owner task to exit a retired child-session generation; past it, cancel it TRIGGER_BARRIER_TIMEOUT_S = 30.0 # a barrier may ask for less; past it the call is answered anyway # Must be >= the largest timeout any client can ask for, or the proxy severs a load the @@ -119,6 +120,7 @@ def __init__( self._child_owner: asyncio.Task | None = None self._child_close: asyncio.Event | None = None self._child_lock = asyncio.Lock() # open / retire / close of a generation (taken after _discover_lock) + self._closing = False # set at process shutdown; no generation is opened past it self._discover_lock = asyncio.Lock() self._server_tools: dict[str, list[MCPTool]] = {} self._tools_discovered = False @@ -344,11 +346,11 @@ async def app(scope, receive, send): 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, then the child sessions close.""" + first, then the driver stops and the scheduled trigger actions are cancelled, then the child sessions close.""" @asynccontextmanager async def lifespan(app): try: - await self._ensure_tools_discovered() + await asyncio.wait_for(self._ensure_tools_discovered(), timeout=self.STARTUP_DISCOVERY_TIMEOUT_S) except Exception: # a raised startup makes uvicorn exit the process; discovery is retried on first use instead logger.exception("Tool discovery failed at startup; retrying on first use") @@ -357,7 +359,9 @@ async def lifespan(app): async with mcp_lifespan(app): yield finally: + self._closing = True await self._trigger_engine.stop_driver() + await self._trigger_engine.cancel_scheduled() await self._close_child_sessions() return lifespan @@ -701,7 +705,8 @@ async def _open_persistent_sessions(self, exit_stack: contextlib.AsyncExitStack, try: read_stream, write_stream, _ = await exit_stack.enter_async_context(streamable_http_client(server.mcp_url)) session = await exit_stack.enter_async_context(ClientSession(read_stream, write_stream)) - await session.initialize() + # a child that accepts the connection but never answers initialize is a failed attempt, not a hang + await asyncio.wait_for(session.initialize(), timeout=CARD_FETCH_TIMEOUT_S) sessions[server.mcp_url] = session logger.info(f"Opened persistent session to {server.name}") break @@ -722,6 +727,8 @@ async def _ensure_child_sessions(self) -> dict[str, ClientSession]: async with self._child_lock: if self._child_owner is not None and not self._child_owner.done(): return self._child_sessions + if self._closing: + raise RuntimeError("gateway is shutting down") await self._close_child_sessions() # a dead or half-open generation loop = asyncio.get_running_loop() ready: asyncio.Future = loop.create_future() @@ -782,8 +789,9 @@ async def _reset_child_sessions(self, stale: dict[str, ClientSession]) -> None: async def _call_child_tool(self, server_url: str, tool_name: str, arguments: dict, *, meta: dict | 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. Discovery is not redone on a reconnect; the - registered proxies look their session up by URL.""" + 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.""" sessions = await self._ensure_child_sessions() try: return await asyncio.wait_for(sessions[server_url].call_tool(tool_name, arguments, meta=meta), @@ -791,8 +799,11 @@ async def _call_child_tool(self, server_url: str, tool_name: str, arguments: dic except Exception as e: if not _is_dead_session(e): raise - logger.warning(f"child session to {server_url} is dead ({e!r}), reconnecting...") await self._reset_child_sessions(sessions) + if not _call_never_reached_child(e): + logger.warning(f"child session to {server_url} died with {tool_name} in flight ({e!r}); not retried") + raise + logger.warning(f"child session to {server_url} is dead ({e!r}), reconnecting...") sessions = await self._ensure_child_sessions() return await asyncio.wait_for(sessions[server_url].call_tool(tool_name, arguments, meta=meta), timeout=self.TOOL_CALL_TIMEOUT_S) @@ -1201,14 +1212,23 @@ def model_dump_one_level(self) -> dict[str, Any]: def _is_dead_session(exc: BaseException) -> bool: - """Whether a child call failed because its session is gone: the stream closed under it, the child answered 404 - for the session id (idle expiry or a restart, which the client reports as code 32600), or the transport died - with the call in flight (CONNECTION_CLOSED).""" + """Whether a child call failed because its session is gone, so the generation is retired: the stream closed under + it, the child answered 404 for the session id (idle expiry or a restart, which the client reports as code 32600), + or the transport died with the call in flight (CONNECTION_CLOSED).""" if isinstance(exc, (ClosedResourceError, anyio.BrokenResourceError)): return True return isinstance(exc, McpError) and exc.error.code in (32600, CONNECTION_CLOSED) +def _call_never_reached_child(exc: BaseException) -> bool: + """Whether a dead-session failure may be retried: the stream was closed before the request was written, or the + child rejected the session id at the transport. Not CONNECTION_CLOSED: the call was in flight when the transport + died, and the child may already have performed its write.""" + if isinstance(exc, (ClosedResourceError, anyio.BrokenResourceError)): + return True + return isinstance(exc, McpError) and exc.error.code == 32600 + + def _is_child_mcp_path(decoded_rest: str) -> bool: """Whether a /svc/{service}/{path} remainder reaches the backing server's MCP endpoint (`mcp` or anything under it), judged after dot-segment normalisation, since httpx collapses `x/../mcp` to `/mcp` on the way upstream. diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py index 3e97b4b3..79c52993 100644 --- a/src/agent_env/env/gateway/triggers.py +++ b/src/agent_env/env/gateway/triggers.py @@ -546,6 +546,15 @@ async def stop_driver(self) -> None: with contextlib.suppress(asyncio.CancelledError): await task + async def cancel_scheduled(self) -> None: + """Cancel the scheduled trigger actions and wait for them to exit. The gateway's process lifespan calls it at + shutdown, before the child sessions close, so no action outlives them.""" + tasks = list(self._tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + self._tasks.clear() + def _resolve_watch_roles(self, body: dict) -> set[str]: raw = body.get("watch_roles") if raw is None: @@ -667,13 +676,16 @@ def _normalize_action(action: Any, w: str, has_ctx: bool, executor: dict | None) if not isinstance(action.get("args", {}), dict): raise TriggerError(f"{w}: args must be an object") _validate_placeholders(action.get("args", {}), f"{w}.args", allow_ctx=has_ctx) - acting = action.get("as") - if acting is not None: + if action.get("as") is not None: + if not isinstance(action["as"], str): + raise TriggerError(f"{w}: 'as' must be a non-empty AgentEnv-Role other than {DEFAULT_ROLE!r}") + acting = action["as"].strip() # the default role means "not forwarded" to the child, so naming it can only be a mistake - if not isinstance(acting, str) or not acting.strip() or acting == DEFAULT_ROLE: + if not acting or acting == DEFAULT_ROLE: raise TriggerError(f"{w}: 'as' must be a non-empty AgentEnv-Role other than {DEFAULT_ROLE!r}") if _VAR_LOOSE_RE.search(acting): # never run through _template, so a placeholder would be sent verbatim raise TriggerError(f"{w}: 'as' is not templated; name the role literally") + action["as"] = acting # forwarded as-is by _run_action, so stripped here like the header path strips else: if action.get("action") not in ("enable", "disable"): raise TriggerError(f"{w}: action must be 'enable' or 'disable'") diff --git a/tst/unit/env/gateway/child_sessions_test.py b/tst/unit/env/gateway/child_sessions_test.py index 01db6c08..b3e0f015 100644 --- a/tst/unit/env/gateway/child_sessions_test.py +++ b/tst/unit/env/gateway/child_sessions_test.py @@ -1,9 +1,11 @@ """One process-wide generation of child MCP sessions: opened on first need by one owner task (the client contexts are anyio task groups, exited only by the task that entered them), shared by the MCP proxy, /step and the trigger engine, -reopened once when it turns out dead and never re-discovered; and the process lifespan that brackets it.""" +reopened once when it turns out dead (the call retried only if it never reached the child) and never re-discovered; +and the process lifespan that brackets it.""" from __future__ import annotations import asyncio +import contextlib from types import SimpleNamespace import anyio @@ -12,6 +14,7 @@ from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent, Tool as MCPTool, ToolAnnotations from agent_env.env.gateway import InternalMCPServer +from agent_env.env.gateway import gateway as gateway_module from agent_env.env.gateway.gateway import Gateway URL = "http://backing/mcp" @@ -44,8 +47,8 @@ async def call_tool(self, name, arguments, *args, meta=None, **kwargs): class _Opener: """Stands in for _open_persistent_sessions: counts attempts and generations, records the entering/exiting task.""" - def __init__(self, session: _FakeSession, *, delay: float = 0.0, failures=()): - self.session, self.delay, self.failures = session, delay, list(failures) + def __init__(self, session: _FakeSession, *, delay: float = 0.0, failures=(), gate: asyncio.Event | None = None): + self.session, self.delay, self.failures, self.gate = session, delay, list(failures), gate self.attempts = 0 self.opened = 0 self.entered: list[asyncio.Task] = [] @@ -54,6 +57,8 @@ def __init__(self, session: _FakeSession, *, delay: float = 0.0, failures=()): async def __call__(self, stack): self.attempts += 1 await asyncio.sleep(self.delay) + if self.gate is not None: + await self.gate.wait() if self.failures: raise self.failures.pop(0) self.opened += 1 @@ -76,6 +81,30 @@ def _explode(*a, **kw): raise AssertionError("a completed pass must neither reopen nor rediscover") +async def _run_lifespan(app, release: asyncio.Event, sent: list[dict]) -> None: + """Drive one ASGI lifespan: startup now, shutdown once `release` is set; what the app sends lands in `sent`.""" + scope = {"type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, "state": {}} + pending = [{"type": "lifespan.startup"}] + + async def receive(): + if pending: + return pending.pop() + await release.wait() + return {"type": "lifespan.shutdown"} + + async def send(message): + sent.append(message) + + await app(scope, receive, send) + + +async def _startup_answered(sent: list[dict]) -> None: + for _ in range(500): + if sent: + return + await asyncio.sleep(0.01) + + @pytest.mark.asyncio async def test_concurrent_first_use_opens_one_generation(): opener = _Opener(_FakeSession(), delay=0.02) @@ -103,8 +132,7 @@ async def test_the_generation_is_exited_by_the_task_that_entered_it(): anyio.ClosedResourceError(), anyio.BrokenResourceError(), McpError(ErrorData(code=32600, message="Session terminated")), - McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), -], ids=["closed", "broken", "session_terminated", "connection_closed"]) +], ids=["closed", "broken", "session_terminated"]) @pytest.mark.asyncio async def test_a_dead_session_is_reopened_once_and_the_call_retried_without_rediscovery(death): session = _FakeSession() @@ -123,6 +151,25 @@ async def test_a_dead_session_is_reopened_once_and_the_call_retried_without_redi await gw._close_child_sessions() +@pytest.mark.asyncio +async def test_a_call_lost_in_flight_retires_the_generation_but_is_not_retried(): + session = _FakeSession() + session.failures = [McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed"))] + opener = _Opener(session) + gw = _gateway(opener) + try: + await gw._ensure_child_sessions() + first_owner = gw._child_owner + with pytest.raises(McpError, match="Connection closed"): + await gw._call_child_tool(URL, "svc_thing", {}) + assert session.calls == ["svc_thing"] # the child may already have written; a retry could write twice + assert first_owner.done() and gw._child_owner is None and opener.opened == 1 + assert not (await gw._call_child_tool(URL, "svc_thing", {})).isError # the next call reopens + assert opener.opened == 2 and session.calls == ["svc_thing", "svc_thing"] + finally: + await gw._close_child_sessions() + + @pytest.mark.asyncio async def test_any_other_error_is_the_caller_s_and_keeps_the_generation(): session = _FakeSession() @@ -216,32 +263,41 @@ async def test_discovery_on_a_dead_session_retires_the_generation(): @pytest.mark.asyncio -async def test_the_process_lifespan_starts_the_driver_tolerates_a_failed_startup_discovery_and_unwinds(): - opener = _Opener(_FakeSession(), failures=[RuntimeError("refused")]) - gw = _gateway(opener) - app = gw._asgi_app() - scope = {"type": "lifespan", "asgi": {"version": "3.0", "spec_version": "2.0"}, "state": {}} +async def test_a_child_that_never_answers_initialize_is_a_failed_connect_attempt(monkeypatch): + class _HungSession: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self - async def run(release: asyncio.Event, sent: list[dict]) -> None: - pending = [{"type": "lifespan.startup"}] + async def __aexit__(self, *exc): + return None - async def receive(): - if pending: - return pending.pop() - await release.wait() - return {"type": "lifespan.shutdown"} + async def initialize(self): + await asyncio.Event().wait() - async def send(message): - sent.append(message) + @contextlib.asynccontextmanager + async def streams(url): + yield None, None, None - await app(scope, receive, send) + monkeypatch.setattr(gateway_module, "streamable_http_client", streams) + monkeypatch.setattr(gateway_module, "ClientSession", _HungSession) + monkeypatch.setattr(gateway_module, "CARD_FETCH_TIMEOUT_S", 0.02) + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[InternalMCPServer("svc", URL)]) + async with contextlib.AsyncExitStack() as stack: + with pytest.raises(TimeoutError): + await gw._open_persistent_sessions(stack, max_retries=2, retry_delay=0) + +@pytest.mark.asyncio +async def test_the_process_lifespan_starts_the_driver_tolerates_a_failed_startup_discovery_and_unwinds(): + opener = _Opener(_FakeSession(), failures=[RuntimeError("refused")]) + gw = _gateway(opener) + app = gw._asgi_app() release, sent = asyncio.Event(), [] - task = asyncio.create_task(run(release, sent)) - for _ in range(500): - if sent: - break - await asyncio.sleep(0.01) + task = asyncio.create_task(_run_lifespan(app, release, sent)) + await _startup_answered(sent) assert sent == [{"type": "lifespan.startup.complete"}] # fail-open: a raised startup would exit uvicorn driver = gw._trigger_engine._driver_task assert driver is not None and not driver.done() @@ -256,6 +312,66 @@ async def send(message): # the MCP session manager runs once per app, so a second lifespan fails its startup; the process state still unwinds sent2: list[dict] = [] with pytest.raises(RuntimeError, match="only be called once"): - await run(asyncio.Event(), sent2) + await _run_lifespan(app, asyncio.Event(), sent2) assert sent2[0]["type"] == "lifespan.startup.failed" assert gw._trigger_engine._driver_task is None and gw._child_owner is None + + +@pytest.mark.asyncio +async def test_startup_discovery_is_bounded_and_a_hung_open_leaves_no_generation(monkeypatch): + monkeypatch.setattr(Gateway, "STARTUP_DISCOVERY_TIMEOUT_S", 0.05) + gate = asyncio.Event() + opener = _Opener(_FakeSession(), gate=gate) # the open hangs until released + gw = _gateway(opener) + app = gw._asgi_app() + release, sent = asyncio.Event(), [] + task = asyncio.create_task(_run_lifespan(app, release, sent)) + await _startup_answered(sent) + assert sent == [{"type": "lifespan.startup.complete"}] and not gw._trigger_engine._driver_task.done() + assert gw._child_owner is None and gw._child_sessions == {} and not gw._tools_discovered + assert not gw._child_lock.locked() and not gw._discover_lock.locked() + gate.set() # the hung open now returns to an owner nobody waits on; it exits what it opened + try: + assert (await gw._ensure_child_sessions())[URL] is opener.session + await asyncio.sleep(0.01) + assert opener.attempts == 2 and opener.opened == 2 and opener.exited == [opener.entered[0]] + assert opener.entered[1] is gw._child_owner and not gw._child_owner.done() + await gw._ensure_tools_discovered() + assert gw._tools_discovered + finally: + release.set() + await asyncio.wait_for(task, 5) + assert sent[-1] == {"type": "lifespan.shutdown.complete"} and gw._child_owner is None + + +@pytest.mark.asyncio +async def test_shutdown_cancels_scheduled_actions_and_nothing_reopens_the_child_sessions(): + session = _FakeSession() + opener = _Opener(session) + gw = _gateway(opener) + app = gw._asgi_app() + release, sent = asyncio.Event(), [] + task = asyncio.create_task(_run_lifespan(app, release, sent)) + await _startup_answered(sent) + assert sent == [{"type": "lifespan.startup.complete"}] and not gw._child_owner.done() + outcome: list[str] = [] + + async def late_action(): # a trigger action still waiting when the process shuts down + try: + await asyncio.sleep(60) + await gw._call_child_tool(URL, "svc_thing", {}) + except asyncio.CancelledError: + outcome.append("cancelled") + raise + + scheduled = gw._trigger_engine._schedule(late_action()) + await asyncio.sleep(0) + release.set() + await asyncio.wait_for(task, 5) + assert sent[-1] == {"type": "lifespan.shutdown.complete"} + assert scheduled.cancelled() and outcome == ["cancelled"] and gw._trigger_engine._tasks == set() + assert session.calls == [] and opener.opened == 1 and opener.exited == opener.entered + assert gw._child_owner is None and gw._child_sessions == {} + with pytest.raises(RuntimeError, match="shutting down"): + await gw._ensure_child_sessions() + assert gw._child_owner is None and opener.attempts == 1 diff --git a/tst/unit/env/gateway/role_forwarding_test.py b/tst/unit/env/gateway/role_forwarding_test.py index dd069dff..e6e12bcf 100644 --- a/tst/unit/env/gateway/role_forwarding_test.py +++ b/tst/unit/env/gateway/role_forwarding_test.py @@ -15,7 +15,7 @@ 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.gateway import GATEWAY_TRAJECTORY_FILE, Gateway, _is_dead_session +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 URL = "http://backing/mcp" @@ -114,14 +114,15 @@ async def test_trigger_internal_calls_are_stamped_only_when_an_as_role_is_given( assert "current_time" in native.content[0].text # gateway-native: no child, no identity, no session touched -@pytest.mark.parametrize("exc,dead", [ - (anyio.ClosedResourceError(), True), - (anyio.BrokenResourceError(), True), - (McpError(ErrorData(code=32600, message="Session terminated")), True), - (McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), True), - (McpError(ErrorData(code=-32602, message="Invalid params")), False), - (asyncio.TimeoutError(), False), - (ValueError("x"), False), +@pytest.mark.parametrize("exc,dead,retry", [ + (anyio.ClosedResourceError(), True, True), + (anyio.BrokenResourceError(), True, True), + (McpError(ErrorData(code=32600, message="Session terminated")), True, True), + (McpError(ErrorData(code=CONNECTION_CLOSED, message="Connection closed")), True, False), # in flight: may have written + (McpError(ErrorData(code=-32602, message="Invalid params")), False, False), + (asyncio.TimeoutError(), False, False), + (ValueError("x"), False, False), ], ids=["closed", "broken", "session_terminated", "connection_closed", "invalid_params", "timeout", "other"]) -def test_the_dead_session_predicate(exc, dead): +def test_the_dead_session_predicates(exc, dead, retry): assert _is_dead_session(exc) is dead + assert _call_never_reached_child(exc) is retry diff --git a/tst/unit/env/gateway/triggers_test.py b/tst/unit/env/gateway/triggers_test.py index 2ae6568b..f944683b 100644 --- a/tst/unit/env/gateway/triggers_test.py +++ b/tst/unit/env/gateway/triggers_test.py @@ -119,7 +119,8 @@ def _registration(**overrides): "triggers": []}, "must not be in watch_roles"), *[({"triggers": [{"id": "x", "when": {"type": "action", "tool": "t"}, "actions": [{"type": "tool", "tool": "t2", "as": acting}]}]}, fragment) - for acting, fragment in [("", "'as' must be a non-empty"), ("default", "'as' must be a non-empty"), + for acting, fragment in [("", "'as' must be a non-empty"), (" ", "'as' must be a non-empty"), + ("default", "'as' must be a non-empty"), (" default ", "'as' must be a non-empty"), (5, "'as' must be a non-empty"), ("${args.owner}", "not templated")]], ({"triggers": [{"id": "x", "when": {"type": "action", "tool": "t"}, "actions": [{"type": "permission", "action": "enable", "role": "r", "tools": ["t"], "as": "x"}]}]}, @@ -289,6 +290,31 @@ async def test_a_tool_action_s_as_role_reaches_its_call_and_verify_but_not_the_s assert [(e["as"], e["ok"]) for e in internal] == [("carol@example.com", True), (None, True)] +@pytest.mark.asyncio +async def test_a_tool_action_s_as_role_is_stripped_once_at_registration(engine): + reg = _registration() + reg["triggers"][1]["actions"][0]["as"] = " carol@example.com " + engine.register(reg) + assert engine._triggers["rate"]["spec"]["actions"][0]["as"] == "carol@example.com" + engine._test_state["doc_text"] = "Budget Fit: fine." + engine.on_tool_call("default", "gdocs_batch_update", {"documentId": "doc-1"}, _result()) + await asyncio.sleep(0.3) + assert engine._triggers["rate"]["status"] == "fired" + assert [role for tool, role in engine._test_gw.roles if tool.startswith("slack_")] == ["carol@example.com"] * 2 + assert [e["detail"]["as"] for e in engine._events if e["kind"] == "action_ok"] == ["carol@example.com"] + + +@pytest.mark.asyncio +async def test_cancel_scheduled_cancels_every_pending_action_and_waits_for_it(engine): + async def pending(): + await asyncio.sleep(60) + + tasks = [engine._schedule(pending()) for _ in range(2)] + await asyncio.sleep(0) + await engine.cancel_scheduled() + assert all(t.cancelled() for t in tasks) and engine._tasks == set() + + @pytest.mark.asyncio async def test_sensor_trigger_and_state_shape(engine): engine.register(_registration()) From 9e7fd9feba6f7099aaa52a19d442bee6ee53f941 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 13:29:08 -0700 Subject: [PATCH 3/4] fix(gateway): cancel an abandoned child-session open instead of letting it finish for nobody When the startup bound (or any waiter) gives up on a generation still being opened, the owner task is cancelled too, so its stack unwinds what it entered and no connections are held for a generation nobody will use. Co-Authored-By: Claude Fable 5.1 --- src/agent_env/env/gateway/gateway.py | 5 ++++- tst/unit/env/gateway/child_sessions_test.py | 10 +++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index c6858483..ebd71de4 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -738,7 +738,10 @@ async def _ensure_child_sessions(self) -> dict[str, ClientSession]: try: sessions = await ready except BaseException: - close.set() # the owner exits whatever it opened (also a success that lands after we left) + # Nobody will use this generation: an open still in progress is abandoned (its stack unwinds what + # it entered), not finished in the background to hold connections nobody waits on. + close.set() + owner.cancel() raise self._child_sessions, self._child_owner, self._child_close = sessions, owner, close return sessions diff --git a/tst/unit/env/gateway/child_sessions_test.py b/tst/unit/env/gateway/child_sessions_test.py index b3e0f015..0b1aba09 100644 --- a/tst/unit/env/gateway/child_sessions_test.py +++ b/tst/unit/env/gateway/child_sessions_test.py @@ -330,12 +330,16 @@ async def test_startup_discovery_is_bounded_and_a_hung_open_leaves_no_generation assert sent == [{"type": "lifespan.startup.complete"}] and not gw._trigger_engine._driver_task.done() assert gw._child_owner is None and gw._child_sessions == {} and not gw._tools_discovered assert not gw._child_lock.locked() and not gw._discover_lock.locked() - gate.set() # the hung open now returns to an owner nobody waits on; it exits what it opened + await asyncio.sleep(0) + # the abandoned open was cancelled, not left running to finish for nobody + assert not [t for t in asyncio.all_tasks() if t.get_name() == "gateway-child-sessions" and not t.done()] + assert opener.attempts == 1 and opener.opened == 0 + gate.set() try: assert (await gw._ensure_child_sessions())[URL] is opener.session await asyncio.sleep(0.01) - assert opener.attempts == 2 and opener.opened == 2 and opener.exited == [opener.entered[0]] - assert opener.entered[1] is gw._child_owner and not gw._child_owner.done() + assert opener.attempts == 2 and opener.opened == 1 and opener.exited == [] + assert opener.entered[0] is gw._child_owner and not gw._child_owner.done() await gw._ensure_tools_discovered() assert gw._tools_discovered finally: From d4c51800deea5afb3bfa8eaffb1ccbb3764682bb Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 9 Oct 2026 18:43:00 -0700 Subject: [PATCH 4/4] gateway: give a child 30 s to answer tools/list during discovery The per-child bound reused the 5 s card-fetch budget. The child answers on the connection every caller shares, so a large service busy with another call can take longer than that, and a timeout fails the whole discovery for the session that triggered it. The initialize handshake keeps the short bound, which already retries ten times. Co-Authored-By: Claude Fable 5.1 --- src/agent_env/env/gateway/constants.py | 5 ++++ src/agent_env/env/gateway/gateway.py | 3 ++- tst/unit/env/gateway/child_sessions_test.py | 26 +++++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index d26271c7..13cf6148 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -84,6 +84,11 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: # Per-child card fetch budget (seconds) when composing the env card. CARD_FETCH_TIMEOUT_S = 5.0 +# Per-child tools/list budget (seconds) during discovery. Wider than a card fetch: the child answers on the +# connection every caller shares, so a large service busy with another call can take a while, and a timeout +# here fails the whole discovery, to be retried on the next connect. +CHILD_DISCOVERY_TIMEOUT_S = 30.0 + # Trigger status vocabulary, part of the GET /triggers/state contract; `state` asserts against it. TRIGGER_STATUSES = ("armed", "firing", "queued", "fired", "failed") diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index ebd71de4..f57f45ff 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -47,6 +47,7 @@ AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, CARD_FETCH_TIMEOUT_S, + CHILD_DISCOVERY_TIMEOUT_S, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, GATEWAY_COMPOSE_HOST, @@ -838,7 +839,7 @@ async def _discover_and_register_tools(self, sessions: dict[str, ClientSession]) try: session = sessions[server.mcp_url] # bounded: this runs under _discover_lock, which startup and every MCP connect wait on - tools_result = await asyncio.wait_for(session.list_tools(), timeout=CARD_FETCH_TIMEOUT_S) + tools_result = await asyncio.wait_for(session.list_tools(), timeout=CHILD_DISCOVERY_TIMEOUT_S) tools = tools_result.tools logger.info(f"Found {len(tools)} tool(s) from {server.name}") for tool in tools: diff --git a/tst/unit/env/gateway/child_sessions_test.py b/tst/unit/env/gateway/child_sessions_test.py index 0b1aba09..bb94734f 100644 --- a/tst/unit/env/gateway/child_sessions_test.py +++ b/tst/unit/env/gateway/child_sessions_test.py @@ -246,6 +246,32 @@ async def test_a_failed_discovery_keeps_the_gateway_s_own_entry_and_heals_on_the await gw._close_child_sessions() +@pytest.mark.asyncio +async def test_a_child_that_hangs_on_tools_list_fails_discovery_within_its_bound_and_keeps_the_generation(monkeypatch): + class _HangsOnList(_FakeSession): + hang = True + + async def list_tools(self): + if self.hang: + await asyncio.Event().wait() + return await super().list_tools() + + session = _HangsOnList() + opener = _Opener(session) + gw = _gateway(opener) + monkeypatch.setattr(gateway_module, "CHILD_DISCOVERY_TIMEOUT_S", 0.02) + try: + with pytest.raises(TimeoutError): + await gw._ensure_tools_discovered() + assert gw._tools_discovered is False + session.hang = False + await gw._ensure_tools_discovered() + assert gw._tools_discovered is True + assert opener.opened == 1 # a slow child is not a dead one: the generation stays + finally: + await gw._close_child_sessions() + + @pytest.mark.asyncio async def test_discovery_on_a_dead_session_retires_the_generation(): session = _FakeSession()