From 056c4a06e250a9b8b6f664aee7cb6c3dfee70a07 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 4 Nov 2026 09:56:56 -0800 Subject: [PATCH] fix(env): the gateway's control routes take the deploy's control token Each gateway deploy draws a control token, as computer envs already do, passes it to the gateway as AGENTENV_CONTROL_TOKEN and keeps it on the deploy record. The gateway then refuses any call without it to the routes that change tool access, triggers, the clock and notifications. Their framework callers already send the record's control headers. Every other route stays open for now. A gateway started without a token checks none: an older framework, an export, or a create_gateway caller that passes none. An older gateway image ignores the token. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/constants.py | 9 +++ src/agent_env/env/gateway/entrypoint.py | 3 +- src/agent_env/env/gateway/gateway.py | 13 +++- .../env_providers/env_gateway_provider.py | 27 ++++++- tst/integration/env/gateway/gateway_test.py | 49 +++++++----- .../test_register_env_triggers_local.py | 12 +-- tst/unit/env/envs/test_multi_env_name.py | 5 +- tst/unit/env/gateway/control_token_test.py | 74 +++++++++++++++++++ tst/unit/env/gateway/gateway_card_test.py | 3 +- .../env_gateway_provider_test.py | 27 ++++++- 10 files changed, 185 insertions(+), 37 deletions(-) create mode 100644 tst/unit/env/gateway/control_token_test.py diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 30bd00031..3791cfbdf 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -68,6 +68,7 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: METHOD_ADD = "data/add" METHOD_GET = "data/get" MCP_TRANSPORT = "mcp" +CONTROL_TOKEN_ENV = "AGENTENV_CONTROL_TOKEN" # Env/gateway-level extensions the gateway itself owns (tool gating + triggers). EXT_DISABLE_TOOL_URI = "urn:agentenv:disable-tool/v1" @@ -152,6 +153,14 @@ def sources(names: list[str]) -> str: WILDCARD = "*" TOOL_DISABLE_ACTION = "disable" TOOL_ENABLE_ACTION = "enable" +# The harness's routes that take the deploy's control token: the ones that change tool access, triggers, the clock +# and notifications. Every other route stays open for now. +CONTROL_ROUTES = frozenset({ + f"/tools/{TOOL_DISABLE_ACTION}", f"/tools/{TOOL_ENABLE_ACTION}", + "/triggers/register", "/triggers/remove", "/triggers/clear", + "/clock/set-time", "/clock/advance", "/clock/clear", + "/notifications/config", "/notifications/take", "/notifications/clear", +}) class GatewayMode(str, Enum): diff --git a/src/agent_env/env/gateway/entrypoint.py b/src/agent_env/env/gateway/entrypoint.py index b0b936ae9..f5d2ed4b3 100644 --- a/src/agent_env/env/gateway/entrypoint.py +++ b/src/agent_env/env/gateway/entrypoint.py @@ -7,7 +7,7 @@ import httpx from . import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode, InternalMCPServer -from .constants import DEFAULT_MCP_SERVER_NAME +from .constants import CONTROL_TOKEN_ENV, DEFAULT_MCP_SERVER_NAME from .gateway import Gateway logger = logging.getLogger(__name__) @@ -152,5 +152,6 @@ def parse_name_url_map(env_str: str) -> dict[str, str]: rest_proxy_urls=rest_proxy_urls, gateway_mode=gateway_mode, service_db_url=service_db_url, + control_token=os.environ.get(CONTROL_TOKEN_ENV), ) gateway.run() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 73f6358be..63de66f88 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -4,6 +4,7 @@ import asyncio import atexit import contextlib +import hmac import json import logging import os @@ -47,6 +48,7 @@ CHANGES_PAGE_DEFAULT, CHANGES_PAGE_MAX, CHILD_DISCOVERY_TIMEOUT_S, + CONTROL_ROUTES, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, GATEWAY_COMPOSE_HOST, @@ -119,6 +121,7 @@ def __init__( rest_proxy_urls: dict[str, str] | None = None, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, service_db_url: str | None = None, + control_token: str | None = None, ): self.host = host self.port = port @@ -128,6 +131,7 @@ def __init__( self.rest_proxy_urls = rest_proxy_urls or {} self.gateway_mode = gateway_mode self._service_db_url = service_db_url + self._control_bearer = f"Bearer {control_token}".encode() if control_token else None self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: self._db_conn = psycopg2.connect(service_db_url) @@ -413,8 +417,8 @@ def run(self) -> None: # inherited). Process: child sessions and the time-trigger driver, outer to the MCP session manager's lifespan. def _asgi_app(self): - """The served app: FastMCP's Starlette app with the request's caller bound around it, and the process - lifespan wrapped around the MCP session manager's.""" + """The served app: FastMCP's Starlette app with the request's caller bound around it, the control routes behind + the deploy's token when it has one, and the process lifespan wrapped around the MCP session manager's.""" inner = self._mcp.streamable_http_app() inner.router.lifespan_context = self._process_lifespan(inner.router.lifespan_context) @@ -422,6 +426,11 @@ async def app(scope, receive, send): if scope["type"] != "http": await inner(scope, receive, send) return + if self._control_bearer and scope["path"] in CONTROL_ROUTES and not hmac.compare_digest( + dict(scope.get("headers", ())).get(b"authorization", b""), self._control_bearer): + response = self._json_response({"ok": False, "error": "this route takes the env's control token"}, 401) + await response(scope, receive, send) + return role = DEFAULT_ROLE for name, value in scope.get("headers", ()): if name == ROLE_HEADER_BYTES: diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 84ca8cb11..27768fb66 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -38,6 +38,8 @@ DB_WEB_PORT, PGWEB_SERVICE_NAME, ) +from agent_env.env.envs.computer.utils.boot import control_environment +from agent_env.env.envs.computer.utils.boot import control_token as new_control_token from agent_env.env.envs.playwright_mcp import PLAYWRIGHT_MCP_RECORDER_PORT from agent_env.providers.env_providers.constants import ( AGENT_ENV_WEBSITE_BACKEND_PORT, @@ -119,6 +121,8 @@ class DeployedGateway: # The env card the readiness probe read, and when (ISO 8601, UTC). environment_card: dict | None = None environment_card_read_at_utc: str | None = None + # The token the gateway's control routes take, which the record keeps; None for a gateway started without one. + control_token: str | None = None @dataclass @@ -283,6 +287,7 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, + control_token: str | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. @@ -491,6 +496,7 @@ def create_docker_compose( gateway_env.append(f" - REST_PROXY_URLS={','.join(rest_proxy_urls)}") gateway_env.append(f" - GATEWAY_MODE={gateway_mode.value}") gateway_env.append(f" - SERVICE_DB_URL={state_provider.base_url(state_instance)}") + gateway_env.extend(f" - {key}={value}" for key, value in control_environment(control_token).items()) publish = host_port or (lambda port: port) exposed_gateway_ports = [] if expose_gateway_port: @@ -625,6 +631,7 @@ async def deploy( mcp_server_images=topology.mcp_server_images, website_configs=topology.website_configs, website_images=topology.website_images, mcp_server_name=topology.mcp_server_name, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, attribution=attribution, + control_token=new_control_token(), ) return DeployedGatewayEnv( env_id=env.id, @@ -641,6 +648,7 @@ async def deploy( website_frontend_urls=deployed.website_frontend_urls, gateway_mode=gateway_mode.value, env_state_instance_ids=deployed.env_state_instance_ids, + control_token=deployed.control_token, ) async def create_gateway( @@ -663,6 +671,7 @@ async def create_gateway( mcp_server_name: str | None = None, *, attribution: Optional[Attribution] = None, + control_token: str | None = None, ) -> DeployedGateway: """Create and start a gateway from explicit topology, for a caller that builds its own record rather than calling ``deploy()``. @@ -676,13 +685,16 @@ async def create_gateway( ``mcp_server_name`` is the name the gateway presents on its card and as its MCP server; drawn as ``env`` + 4 random digits when the env declares none; the record takes it from the gateway's card. + + ``control_token`` is the token the gateway's control routes take, which the caller's record must keep; without one + they stay open. """ return await self._launch( sandbox_provider, env_id, state_instance=state_instance, mcp_servers=mcp_servers, mcp_server_images=mcp_server_images, gateway_port=gateway_port, website_configs=website_configs, website_images=website_images, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, mcp_server_name=mcp_server_name, - attribution=attribution, + attribution=attribution, control_token=control_token, ) async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, *, mcp_server_name: str | None = None, **gateway) -> DeployedGateway: @@ -741,6 +753,7 @@ async def _deploy_via_vm( existing_sandbox: Sandbox | None = None, sidecars: list[SidecarConfig] | None = None, mcp_server_name: str | None = None, + control_token: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env from agent_env.config import gateway_env_id, get_config, website_browser_env_id @@ -832,6 +845,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + control_token=control_token, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -925,7 +939,7 @@ async def _stand_up(_spec): gateway_url=gateway_url, mcp_url=mcp_url, db_web_url=db_web_url, db_mcp_url=db_mcp_url, website_frontend_urls=website_frontend_urls, mcp_server_name=mcp_server_name, env_state_instance_ids=[state_instance.instance_id], - environment_card=card, environment_card_read_at_utc=card_read_at_utc, + environment_card=card, environment_card_read_at_utc=card_read_at_utc, control_token=control_token, ) async def _deploy_via_containers( @@ -942,6 +956,7 @@ async def _deploy_via_containers( memory_mb: int | None = None, attribution: Optional[Attribution] = None, mcp_server_name: str | None = None, + control_token: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env from agent_env.config import gateway_env_id, get_config @@ -1038,6 +1053,7 @@ async def _deploy_via_containers( "GATEWAY_MODE": gateway_mode.value, "SERVICE_DB_URL": self._state_provider.base_url(state_instance), "MCP_SERVER_NAME": mcp_server_name or DEFAULT_MCP_SERVER_NAME, + **control_environment(control_token), } gateway_sb = await sandbox_provider.create_container( image_name=gateway_env.docker_image_artifact.image_name, @@ -1071,6 +1087,7 @@ async def _deploy_via_containers( env_state_instance_ids=[state_instance.instance_id], environment_card=card, environment_card_read_at_utc=card_read_at_utc, + control_token=control_token, ) except Exception: await self.close() @@ -1371,6 +1388,7 @@ async def _deploy_gateway( existing_sandbox: Sandbox | None = None, sidecars: list[SidecarConfig] | None = None, mcp_server_name: str | None = None, + control_token: str | None = None, attribution: Optional[Attribution] = None, ) -> DeployedGateway: """One gateway deploy on one (unchained) sandbox provider: a VM on one that creates VMs, else a container per @@ -1396,7 +1414,8 @@ async def _deploy_gateway( sandbox_provider, mcp_servers, mcp_server_images, gateway_port=gateway_port, website_configs=website_configs, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, - cpu=cpu, memory_mb=memory_mb, mcp_server_name=mcp_server_name, **deploy_kwargs, + cpu=cpu, memory_mb=memory_mb, mcp_server_name=mcp_server_name, control_token=control_token, + **deploy_kwargs, ) earlier_sandbox = self._sandbox try: @@ -1405,7 +1424,7 @@ async def _deploy_gateway( gateway_port=gateway_port, website_configs=website_configs, website_images=website_images, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, - mcp_server_name=mcp_server_name, **deploy_kwargs, + mcp_server_name=mcp_server_name, control_token=control_token, **deploy_kwargs, ) except BaseException: await self._terminate_new_sandbox(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) diff --git a/tst/integration/env/gateway/gateway_test.py b/tst/integration/env/gateway/gateway_test.py index 255aea263..e1829624c 100644 --- a/tst/integration/env/gateway/gateway_test.py +++ b/tst/integration/env/gateway/gateway_test.py @@ -763,7 +763,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) # Role-based tool filtering: disable list_emails for "default" role, verify it's hidden from # default callers but still accessible to a non-default role; then re-enable and verify recovery. - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: resp = await client.post( f"{result.gateway_url}/tools/{TOOL_DISABLE_ACTION}", json={"role": "default", "tools": ["list_emails"]}, @@ -839,7 +839,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) assert "alex.chen@techcorp.com" in (await cli_session.call_tool("list_emails", {"folder_name": "INBOX"})).text # Re-enable: list_emails accessible to default role again - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: resp = await client.post( f"{result.gateway_url}/tools/{TOOL_ENABLE_ACTION}", json={"role": "default", "tools": ["list_emails"]}, @@ -863,7 +863,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) # Global disable (role="*"): tool should be hidden from every role, including never-seen ones, # and become visible to everyone again on global enable. - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: resp = await client.post( f"{result.gateway_url}/tools/{TOOL_DISABLE_ACTION}", json={"role": "*", "tools": ["list_emails"]}, @@ -901,7 +901,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) tools_result = await session.list_tools() assert "list_emails" not in [t.name for t in tools_result.tools] - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: resp = await client.post( f"{result.gateway_url}/tools/{TOOL_ENABLE_ACTION}", json={"role": "*", "tools": ["list_emails"]}, @@ -920,7 +920,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) # Override scenario: globally disable everything via the wildcard tool sentinel ("*"), then # re-enable everything for one specific role. Verifies (a) per-role allow beats global deny, # and (b) wildcard rules apply to all tools without snapshotting (future-tool semantic). - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: resp = await client.post( f"{result.gateway_url}/tools/{TOOL_DISABLE_ACTION}", json={"role": "*", "tools": "*"}, @@ -2268,7 +2268,7 @@ async def assert_list_items_present(session, tools_result, tool_names): await _verify_list_tools(mcp_url=result.mcp_url, tool_verifier_callback=assert_list_items_present) - async with httpx.AsyncClient(timeout=30) as client: + async with httpx.AsyncClient(timeout=30, headers=result.control_headers) as client: r = await client.post( f"{result.gateway_url}{disable_params['endpoint']}", json={"role": "default", "tools": ["list_items"]}, @@ -2510,7 +2510,7 @@ async def _poll(client, gw, want_fired, timeout_s=120): "predicate": {"regex": r"alex\.chen@techcorp\.com"}}}]}, ] - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # Baseline: pre-disable send_email for role default so the grant has a visible effect. r = await client.post(f"{gw}/tools/{TOOL_DISABLE_ACTION}", json={"role": "default", "tools": ["send_email"]}, timeout=30) @@ -2677,7 +2677,7 @@ async def test_gateway_barrier_orders_a_templated_mirror(sandbox_provider, mcp_s result = await email_env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: reg = await client.post(f"{gw}/triggers/register", json={"watch_roles": ["default"], "triggers": [mirror]}, timeout=30) assert reg.status_code == 200, f"register failed: {reg.status_code} {reg.text}" @@ -2789,7 +2789,7 @@ async def test_gateway_trigger_rearms_after_a_failed_mirror(sandbox_provider, mc result = await email_env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: reg = await client.post(f"{gw}/triggers/register", json={"watch_roles": ["default"], "triggers": [breaks]}, timeout=30) assert reg.status_code == 200, f"register failed: {reg.status_code} {reg.text}" @@ -2861,7 +2861,7 @@ async def test_gateway_trigger_registration_is_fail_loud(sandbox_provider, mcp_s result = await email_env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # Force discovery, so the unknown-tool checks are live rather than skipped. lt = await client.post(f"{gw}/step", json={"action": "list_tools"}, timeout=30) assert "send_email" in [t["name"] for t in lt.json()["tools"]], lt.text @@ -2896,7 +2896,7 @@ async def test_gateway_website_trigger_fires_on_a_request_through_the_website_pr card = await protocol_v1.get_card(gw) assert TRIGGER_WEBSITE in protocol_v1.extension_params(card, EXT_TRIGGERS_URI)["features"], card - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: unknown = await client.post(f"{gw}/triggers/register", timeout=30, json={ "triggers": [{**visit, "when": {**visit["when"], "app": "calendar"}}]}) assert unknown.status_code == 400 and "unknown website 'calendar'" in unknown.text, unknown.text @@ -2950,7 +2950,7 @@ async def test_gateway_queued_trigger_is_not_reported_settled(sandbox_provider, result = await email_env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: reg = await client.post(f"{gw}/triggers/register", json={"watch_roles": ["default"], "triggers": [slow]}, timeout=30) assert reg.status_code == 200, f"register failed: {reg.status_code} {reg.text}" @@ -3048,7 +3048,7 @@ async def test_gateway_triggers_nl_executor(sandbox_provider, mcp_server_envs, e await email_env.load_environment_artifact(email_service_artifact) ctx = TaskStepContext(deployed_envs=[result], metadata={}) - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # Lock the executor role to exactly one tool (deny-all, then allow send_email). Wrong # realizations become structurally impossible, and the executor's own call stays outside # watch_roles (no cascade). @@ -3204,7 +3204,7 @@ async def _rate_check(client, gw, rate, sleep_s=5.0): result = await env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # Off-by-default + unconditional advertisement. assert (await _read(client, gw)).status_code == 404, "unarmed clock should 404" card = (await client.get(f"{gw}/.well-known/agent-env.json", timeout=30)).json() @@ -3229,6 +3229,19 @@ async def _rate_check(client, gw, rate, sleep_s=5.0): assert GET_TIME not in (await _get_time_surfaces(client, gw, result.mcp_url)).mcp assert (await _read(client, gw)).status_code == 404, "rejected arm must not arm the clock" + # Without the deploy's control token, every control route refuses and changes nothing; reads stay open. + async with httpx.AsyncClient() as agent: + for method, path in (("PUT", "/clock/set-time"), ("PUT", "/clock/advance"), ("POST", "/clock/clear"), + ("POST", "/tools/enable"), ("POST", "/tools/disable"), ("POST", "/triggers/register"), + ("POST", "/triggers/remove"), ("POST", "/triggers/clear"), + ("POST", "/notifications/config"), ("POST", "/notifications/take"), + ("POST", "/notifications/clear")): + refused = await agent.request(method, f"{gw}{path}", timeout=30, + json={"virtual_time": t0a, "virtual_seconds_per_real_second": 3600}) + assert refused.status_code == 401, (path, refused.text) + assert (await agent.get(f"{gw}/clock/time", timeout=30)).status_code == 404 + assert (await _read(client, gw)).status_code == 404, "a refused arm must not arm the clock" + # Arm at t0a, rate=3600 (1 real s -> 1 virtual h); baseline near t0a. armed_at = time.monotonic() arm = await client.put(f"{gw}/clock/set-time", json={"virtual_time": t0a, "virtual_seconds_per_real_second": 3600}, timeout=30) @@ -3398,7 +3411,7 @@ async def _vtime(client, gw): try: result = await env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: for rate in rates: arm = await client.put(f"{gw}/clock/set-time", json={"virtual_time": t0, "virtual_seconds_per_real_second": rate}, timeout=30) assert arm.status_code == 200 and arm.json()["virtual_seconds_per_real_second"] == rate, f"rate={rate}: {arm.text}" @@ -3471,7 +3484,7 @@ async def _drive_until(client, gw, done, timeout_s=90): try: result = await env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # Pre-disable send_email so the flagship's permission mutation is observable. await client.post(f"{gw}/tools/{TOOL_DISABLE_ACTION}", json={"role": "default", "tools": ["send_email"]}, timeout=30) @@ -3567,7 +3580,7 @@ def _dt(iso): try: result = await env.deploy() gw = result.gateway_url - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: # rate 3600: 1 real s = 1 virtual h, so PT2H lands in ~2 real s and PT1H arrives ~1/real s. arm = await client.put(f"{gw}/clock/set-time", json={"virtual_time": t0, "virtual_seconds_per_real_second": 3600}, timeout=30) @@ -5548,7 +5561,7 @@ async def _set_async_wait(result, tool_name, min_seconds, max_seconds=None, scal async def _clock_clear(result) -> None: """Disarm the gateway clock (POST /clock/clear); it then 404s on /clock/time.""" - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=result.control_headers) as client: r = await client.post(f"{result.gateway_url}/clock/clear", timeout=30) assert r.status_code == 200 and r.json()["armed"] is False, r.text diff --git a/tst/integration/task_step/test_register_env_triggers_local.py b/tst/integration/task_step/test_register_env_triggers_local.py index 0a695d011..c4ac0b9ed 100644 --- a/tst/integration/task_step/test_register_env_triggers_local.py +++ b/tst/integration/task_step/test_register_env_triggers_local.py @@ -45,7 +45,7 @@ async def suite(tmp_path_factory): async def items_env(suite): """The module's env as a fresh deploy leaves it: no triggers, executors or events, no clock, and no items.""" _, deployed, _ = suite - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: for path in ("/triggers/clear", "/clock/clear"): (await client.post(path)).raise_for_status() reset = (await client.post(RPC_PATH, json={"jsonrpc": "2.0", "id": "reset", "method": METHOD_RESET, @@ -71,7 +71,7 @@ async def test_each_nl_action_reaches_the_executor_it_names(items_env): context = TaskStepContext(deployed_envs=[deployed], deployed_agents=[alice.agent, bob.agent], metadata={}) await RegisterEnvTriggersStep(id="reg", version=None, env_id=env.id, triggers=triggers, watch_roles=["default"], executor_timeout_seconds=60).execute(context) - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: provoke = await client.post("/step", json={"action": "call_tool", "tool_name": "list_items", "arguments": {}}) assert provoke.status_code == 200, provoke.text state = await _quiet(client, {"to-alice": 1, "to-bob": 1, "on-add": 0}) @@ -113,7 +113,7 @@ async def test_personas_react_to_each_other_until_one_stays_quiet(items_env): context = TaskStepContext(deployed_envs=[deployed], deployed_agents=[denise.agent, bob.agent], metadata={}) await RegisterEnvTriggersStep(id="reg", version=None, env_id=env.id, triggers=triggers, watch_roles=["default"], executor_timeout_seconds=60).execute(context) - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: provoke = await client.post("/step", json={"action": "call_tool", "tool_name": "list_items", "arguments": {}}) assert provoke.status_code == 200, provoke.text state = await _quiet(client, {"denise-hears-solver": 1, "bob-hears-denise": 2, "denise-hears-bob": 1}) @@ -142,7 +142,7 @@ async def test_a_reply_waits_its_delay_and_hours_and_its_writes_name_the_fire(it try: bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started) context = TaskStepContext(deployed_envs=[deployed], deployed_agents=[bob.agent], metadata={}) - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: # The env is shared, so its trajectory holds the earlier tests' calls: read this test's, after them. earlier = len((await client.get("/trajectory")).raise_for_status().text.splitlines()) await RegisterEnvTriggersStep(id="reg", version=None, env_id=env.id, triggers=triggers, watch_roles=["default"], @@ -187,7 +187,7 @@ async def test_a_deadline_fires_on_silence_and_a_reply_cancels_it(items_env): try: bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started) context = TaskStepContext(deployed_envs=[deployed], deployed_agents=[bob.agent], metadata={}) - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: await _arm_clock(client, "2026-10-14T10:00:00Z", 60) await RegisterEnvTriggersStep(id="reg", version=None, env_id=env.id, triggers=triggers, watch_roles=["default"], executor_timeout_seconds=60).execute(context) @@ -216,7 +216,7 @@ async def test_the_step_can_wait_for_start_of_run_time_triggers(items_env): {"id": "later", "when": {"type": "time", "at": "2026-10-14T18:00:00Z"}, "actions": [{"type": "tool", "tool": "items_add_item", "args": {"item": "LATER"}}]}] context = TaskStepContext(deployed_envs=[deployed], deployed_agents=[], metadata={}) - async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60) as client: + async with httpx.AsyncClient(base_url=deployed.gateway_url, timeout=60, headers=deployed.control_headers) as client: await _arm_clock(client, "2026-10-14T10:00:00Z", 60) started = time.monotonic() await RegisterEnvTriggersStep(id="reg", version=None, env_id=env.id, triggers=triggers, watch_roles=["default"], diff --git a/tst/unit/env/envs/test_multi_env_name.py b/tst/unit/env/envs/test_multi_env_name.py index 0055bef90..35d78b764 100644 --- a/tst/unit/env/envs/test_multi_env_name.py +++ b/tst/unit/env/envs/test_multi_env_name.py @@ -67,7 +67,7 @@ async def test_a_deploy_records_the_mcp_url_and_name_its_gateway_card_gives(): card = {"name": "crm", "additionalInterfaces": [{"url": "/mcp", "transport": "mcp"}]} result = SimpleNamespace(gateway_url="https://gw.example", mcp_url="https://gw.example/mcp", db_web_url=None, db_mcp_url=None, website_frontend_urls=None, environment_card=card, environment_card_read_at_utc="t", - env_state_instance_ids=[], mcp_server_name="crm") + env_state_instance_ids=[], mcp_server_name="crm", control_token="t0k") async def deploy_gateway(self, sandbox_provider, **kwargs): self._sandbox = SimpleNamespace(sandbox_id="gw", type="modal_vm") @@ -80,4 +80,5 @@ async def deploy_gateway(self, sandbox_provider, **kwargs): patch("agent_env.providers.env_state.acquire_state_for_deploy", AsyncMock(return_value=None)), \ patch("agent_env.env.envs._deployment.register_env_instance", side_effect=lambda deployed, ttl: deployed): record = await MultiEnv(id="e", version=1, mcp_server_envs=[], name="crm").deploy() - assert (type(record), record.env_provider_type, record.mcp_url, record.mcp_server_name) == (DeployedGatewayEnv, "gateway", "https://gw.example/mcp", "crm") + assert (type(record), record.env_provider_type, record.mcp_url, record.mcp_server_name, record.control_token) == ( + DeployedGatewayEnv, "gateway", "https://gw.example/mcp", "crm", "t0k") diff --git a/tst/unit/env/gateway/control_token_test.py b/tst/unit/env/gateway/control_token_test.py new file mode 100644 index 000000000..39c5abd8e --- /dev/null +++ b/tst/unit/env/gateway/control_token_test.py @@ -0,0 +1,74 @@ +"""The gateway's control routes take the deploy's control token when it has one; every other route stays open.""" +from __future__ import annotations + +import httpx +import pytest + +from agent_env.env.gateway.constants import CONTROL_ROUTES +from agent_env.env.gateway.gateway import Gateway + +TOKEN = "s3cret" +CONTROL = { + "/tools/disable": "POST", "/tools/enable": "POST", + "/triggers/register": "POST", "/triggers/remove": "POST", "/triggers/clear": "POST", + "/clock/set-time": "PUT", "/clock/advance": "PUT", "/clock/clear": "POST", + "/notifications/config": "POST", "/notifications/take": "POST", "/notifications/clear": "POST", +} +OPEN = {"/state": "GET", "/step": "POST", "/clock/time": "GET", "/clock/state": "GET", "/triggers/state": "GET", + "/notifications/state": "GET", "/trajectory": "GET", "/changes": "GET", "/.well-known/agent-env.json": "GET"} + + +def _client(token: str | None) -> httpx.AsyncClient: + gw = Gateway(host="127.0.0.1", port=0, server_name="unit", internal_mcp_servers=[], control_token=token) + return httpx.AsyncClient(transport=httpx.ASGITransport(app=gw._asgi_app()), base_url="http://gateway") + + +def test_the_routes_under_test_are_the_ones_the_gateway_covers(): + assert set(CONTROL) == CONTROL_ROUTES + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(CONTROL)) +@pytest.mark.parametrize("headers", [{}, {"Authorization": "Bearer wrong"}, {"Authorization": TOKEN}], + ids=["none", "wrong", "not-bearer"]) +async def test_a_control_route_refuses_a_call_without_the_token(path, headers): + async with _client(TOKEN) as client: + response = await client.request(CONTROL[path], path, json={}, headers=headers) + state = (await client.get("/clock/state")).json() + assert response.status_code == 401 + assert response.json() == {"ok": False, "error": "this route takes the env's control token"} + assert state["armed"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(CONTROL)) +async def test_a_control_route_answers_a_call_with_the_token(path): + async with _client(TOKEN) as client: + response = await client.request(CONTROL[path], path, json={}, headers={"Authorization": f"Bearer {TOKEN}"}) + assert response.status_code != 401 + + +@pytest.mark.asyncio +async def test_the_token_arms_the_clock_and_its_reads_stay_open(): + async with _client(TOKEN) as client: + armed = await client.put("/clock/set-time", json={"virtual_time": "2019-06-15T12:00:00Z"}, + headers={"Authorization": f"Bearer {TOKEN}"}) + read = await client.get("/clock/time") + assert (armed.status_code, read.status_code) == (200, 200) + assert read.json()["virtual_time"].startswith("2019-06-15T12:00") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(OPEN)) +async def test_every_other_route_stays_open(path): + async with _client(TOKEN) as client: + response = await client.request(OPEN[path], path, json={} if OPEN[path] == "POST" else None) + assert response.status_code != 401 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", sorted(CONTROL)) +async def test_a_gateway_started_without_a_token_checks_none(path): + async with _client(None) as client: + response = await client.request(CONTROL[path], path, json={}) + assert response.status_code != 401 diff --git a/tst/unit/env/gateway/gateway_card_test.py b/tst/unit/env/gateway/gateway_card_test.py index 6b71a6b7c..530bc979c 100644 --- a/tst/unit/env/gateway/gateway_card_test.py +++ b/tst/unit/env/gateway/gateway_card_test.py @@ -249,7 +249,8 @@ def test_every_advertised_endpoint_is_a_served_route(): assert advertised <= served, advertised - served -@pytest.mark.parametrize("name", ["WELL_KNOWN_PATH", "RPC_PATH", "PROTOCOL_VERSION", "METHOD_RESET", "METHOD_ADD", "METHOD_GET", "MCP_TRANSPORT"]) +@pytest.mark.parametrize("name", ["WELL_KNOWN_PATH", "RPC_PATH", "PROTOCOL_VERSION", "METHOD_RESET", "METHOD_ADD", "METHOD_GET", "MCP_TRANSPORT", + "CONTROL_TOKEN_ENV"]) def test_gateway_constants_mirror_the_protocol(name): """constants.py copies these by hand because the gateway image has no agentenv_protocol dependency.""" assert getattr(gateway_constants, name) == getattr(agentenv_protocol, name) diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index a2f9958bc..27cddb2fc 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -373,11 +373,14 @@ async def _spy_acquire(self, ctx): ttl_seconds=60, disk_size_gb=10, mcp_server_name="crm", + control_token="t0k", ) # Last create_container call is the gateway. Inspect its env var. gateway_call = create_calls[-1] env = gateway_call["env"] + assert [c["image_name"] for c in create_calls if "AGENTENV_CONTROL_TOKEN" in (c.get("env") or {})] == ["agent-gateway"] + assert (env["AGENTENV_CONTROL_TOKEN"], result.control_token) == ("t0k", "t0k") assert "INTERNAL_MCP_SERVERS" in env spec = env["INTERNAL_MCP_SERVERS"] # Container mode puts every container on the private network, and the gateway reaches each server there. @@ -1132,6 +1135,7 @@ def _compose_with_sidecar(): gateway_image="agent-gateway", gateway_port=18765, sidecars=[_sidecar_with_secret()], + control_token="ctl-token-value", # Explicit ServiceDBConfig so rendering never resolves the default-db env # from the store (which would hit the network — blocked in unit tests). state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), @@ -1173,14 +1177,29 @@ def test_no_sidecars_renders_no_extra_service(): assert " relay:" not in compose +def test_compose_hands_the_control_token_to_the_gateway_alone(): + from agent_env.env.envs.service_db import ServiceDBConfig + + services = yaml.safe_load(_compose_with_sidecar())["services"] + holders = {name: [item for item in service.get("environment", []) if item.startswith("AGENTENV_CONTROL_TOKEN=")] + for name, service in services.items()} + assert {name: items for name, items in holders.items() if items} == { + GATEWAY_SERVICE_NAME: ["AGENTENV_CONTROL_TOKEN=ctl-token-value"]} + assert "AGENTENV_CONTROL_TOKEN" not in EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], gateway_image="agent-gateway", + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance()) + + def test_redact_compose_secrets_masks_only_secret_env_lines(): """The logged compose masks *SECRET/PASSWORD/TOKEN env values but leaves everything else (image refs, ports, non-secret env, DATABASE_URL) intact.""" compose = _compose_with_sidecar() redacted = _redact_compose_secrets(compose) # secret value gone, key preserved - assert "s3cr3t-token-value" not in redacted + assert "s3cr3t-token-value" not in redacted and "ctl-token-value" not in redacted assert " - UPSTREAM_CLIENT_SECRET=***" in redacted + assert " - AGENTENV_CONTROL_TOKEN=***" in redacted # non-secret lines untouched assert " - UPSTREAM_URL=https://upstream.sandbox.example.com/abc" in redacted assert " image: ecr/relay:1.0.4" in redacted @@ -1944,7 +1963,7 @@ async def test_create_gateway_hands_the_deploy_path_every_argument(attribution, gp, provider = EnvironmentGatewayProvider(), _vm_provider() args = dict(mcp_servers=[MagicMock()], mcp_server_images=[MagicMock()], gateway_port=18999, website_configs=[MagicMock()], website_images=[MagicMock()], gateway_mode=GatewayMode.CONSISTENT, ttl_seconds=61, disk_size_gb=11, cpu=1.5, - memory_mb=3072, existing_sandbox=MagicMock(), sidecars=[MagicMock()], mcp_server_name="crm") + memory_mb=3072, existing_sandbox=MagicMock(), sidecars=[MagicMock()], mcp_server_name="crm", control_token="t0k") with patch.object(gp, "_deploy_via_vm", AsyncMock(return_value=_probed())) as vm, \ patch.object(env_gateway_provider, "_tool_names", AsyncMock(return_value=set())): await gp.create_gateway(sandbox_provider=provider, env_id="crm-env", attribution=attribution, **args) @@ -1995,7 +2014,9 @@ async def test_deploy_returns_the_gateway_record_for_the_env(caplog): env_id="crm-env", env_version=3, gateway_url="https://vm.example", mcp_url="https://vm.example/mcp", db_web_url=None, sandbox_id="vm-1", sandbox_type="modal_vm", sandbox_ids={}, db_mcp_url=None, environment_card_url="https://vm.example/.well-known/agent-env.json", environment_card=card, - environment_card_read_at_utc=record.environment_card_read_at_utc, gateway_mode="consistent", env_state_instance_ids=["st-ext"])) + environment_card_read_at_utc=record.environment_card_read_at_utc, gateway_mode="consistent", env_state_instance_ids=["st-ext"], + control_token=record.control_token)) + assert len(record.control_token) >= 43 and compose.call_args.kwargs["control_token"] == record.control_token assert record.environment_card_read_at_utc and record.instance_id is None # unregistered: the env registers it assert (record.mcp_url, record.mcp_server_name, record.env_provider_type) == ("https://vm.example/mcp", "env1234", "gateway") messages = [r.getMessage() for r in caplog.records if r.getMessage().startswith(("env_deploy_provider_attempt ", "env_deploy_tools_probe "))]