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..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") @@ -96,6 +101,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..f57f45ff 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,7 +44,10 @@ # 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, + CHILD_DISCOVERY_TIMEOUT_S, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, GATEWAY_COMPOSE_HOST, @@ -66,12 +72,16 @@ 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 + 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 # client was still legitimately waiting on -- the ordering this module's constants file @@ -105,17 +115,21 @@ 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._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 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 +138,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 +306,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 +334,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 +343,29 @@ 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 and the scheduled trigger actions are cancelled, then the child sessions close.""" + @asynccontextmanager + async def lifespan(app): + try: + 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") + self._trigger_engine.start_driver() + try: + 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 def _listen_socket(self, family: int, host: str) -> socket.socket: sock = socket.socket(family, socket.SOCK_STREAM) @@ -450,18 +489,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 +528,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 @@ -654,7 +706,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 @@ -667,27 +720,115 @@ 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 + 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() + 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: + # 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 + + 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, 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), + timeout=self.TOOL_CALL_TIMEOUT_S) + except Exception as e: + if not _is_dead_session(e): + raise + 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) + 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 +838,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=CHILD_DISCOVERY_TIMEOUT_S) tools = tools_result.tools logger.info(f"Found {len(tools)} tool(s) from {server.name}") for tool in tools: @@ -717,7 +859,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 +924,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 +946,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 +970,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 +1077,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 +1091,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 +1211,63 @@ 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, 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. + + >>> [_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..79c52993 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,13 +539,22 @@ 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() 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: @@ -644,10 +654,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 +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) + 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 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'") @@ -1010,10 +1032,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 +1053,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 +1068,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 +1123,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 +1161,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..bb94734f --- /dev/null +++ b/tst/unit/env/gateway/child_sessions_test.py @@ -0,0 +1,407 @@ +"""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 (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 +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 import gateway as gateway_module +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=(), 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] = [] + self.exited: list[asyncio.Task] = [] + + 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 + 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") + + +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) + 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")), +], 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() + 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_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() + 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_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() + 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_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 __aexit__(self, *exc): + return None + + async def initialize(self): + await asyncio.Event().wait() + + @contextlib.asynccontextmanager + async def streams(url): + yield None, None, None + + 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_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() + 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_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() + 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 == 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: + 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/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..e6e12bcf --- /dev/null +++ b/tst/unit/env/gateway/role_forwarding_test.py @@ -0,0 +1,128 @@ +"""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, _call_never_reached_child, _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,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_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_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..f944683b 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,17 @@ 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"), (" ", "'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"}]}]}, + "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 +268,53 @@ 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_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()) @@ -348,14 +410,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" },