diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index a04bd51f2..0a10bbeb2 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" @@ -159,6 +160,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 65cf7e3c3..c30970c86 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, GATEWAY_PROVIDER_HEADERS_ENV +from .constants import CONTROL_TOKEN_ENV, DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV from .gateway import Gateway logger = logging.getLogger(__name__) @@ -152,6 +152,7 @@ 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), provider_headers=os.environ.get(GATEWAY_PROVIDER_HEADERS_ENV, "").split(","), ) gateway.run() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 79f6d1911..c1a4b5954 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 @@ -50,6 +51,7 @@ CHANGES_PAGE_DEFAULT, CHANGES_PAGE_MAX, CHILD_DISCOVERY_TIMEOUT_S, + CONTROL_ROUTES, DATA_PLANE_LOAD_TIMEOUT_MAX_S, DEFAULT_ROLE, EXT_WEBSITE_URI, @@ -143,6 +145,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, provider_headers: Iterable[str] = (), ): """`provider_headers` names the headers the sandbox provider in front of the gateway is sent, its credentials, @@ -163,6 +166,7 @@ def __init__( log.addFilter(TokenRedaction()) 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._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: @@ -462,8 +466,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) if self.rest_proxy_urls: @@ -473,6 +477,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 b8df4c430..04009d327 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -45,6 +45,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, @@ -128,6 +130,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 @@ -292,6 +296,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. @@ -500,6 +505,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()) if provider_headers := sandbox_request_header_names(): gateway_env.append(f" - {GATEWAY_PROVIDER_HEADERS_ENV}={','.join(provider_headers)}") publish = host_port or (lambda port: port) @@ -637,6 +643,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, @@ -653,6 +660,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( @@ -675,6 +683,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()``. @@ -688,13 +697,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: @@ -753,6 +765,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 @@ -844,6 +857,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)}") @@ -939,7 +953,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( @@ -956,6 +970,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 @@ -1052,6 +1067,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_PROVIDER_HEADERS_ENV: ",".join(names)} if (names := sandbox_request_header_names()) else {}), } gateway_sb = await sandbox_provider.create_container( @@ -1086,6 +1102,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() @@ -1386,6 +1403,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 @@ -1411,7 +1429,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: @@ -1420,7 +1439,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 a2a023a8c..eb32ad96c 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}) @@ -116,7 +116,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}) @@ -145,7 +145,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"], @@ -190,7 +190,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) @@ -220,7 +220,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 48979ddfa..b26a632b8 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 8ed04f46c..7e2994096 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -377,11 +377,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. @@ -1137,6 +1140,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()), @@ -1178,14 +1182,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 @@ -1980,7 +1999,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) @@ -2031,7 +2050,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 "))]