From 8389f76db9419e5c087ae048ef536b7c2a39a5f7 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 13:34:07 -0800 Subject: [PATCH] feat(gateway): a keyed deploy's role key, forwarded headers, and triggers by route (FD-4683) This PR's change, rebased as one commit onto #429. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../examples/website/notes/README.md | 6 +- .../agentenv_protocol/website/environment.py | 21 +- .../src/agentenv_protocol/website/routes.py | 24 +- .../tests/test_website_api.py | 13 ++ .../tests/test_website_environment.py | 12 + src/agent_env/a2a_agent/a2a_agent.py | 12 +- src/agent_env/cli/env/deploy.py | 7 +- src/agent_env/env/env.py | 9 +- src/agent_env/env/envs/mcp_server.py | 4 +- src/agent_env/env/envs/multi_env.py | 7 +- src/agent_env/env/envs/website.py | 4 +- src/agent_env/env/gateway/__init__.py | 2 + src/agent_env/env/gateway/caller.py | 3 +- src/agent_env/env/gateway/constants.py | 6 +- src/agent_env/env/gateway/entrypoint.py | 3 +- src/agent_env/env/gateway/gateway.py | 139 +++++++++--- src/agent_env/env/gateway/rest_proxy.py | 129 ++++++++++- src/agent_env/env/gateway/triggers.py | 36 ++- .../env_providers/env_gateway_provider.py | 39 +++- src/agent_env/task/episode/run.py | 9 +- src/agent_env/task/episode/steps.py | 5 +- .../task_step/task_steps/deploy_agent.py | 13 +- .../task_step/task_steps/deploy_env.py | 11 +- .../task_step/task_steps/load_artifact.py | 3 +- .../task_steps/mcp_cli_builder/codegen.py | 17 +- .../task_step/task_steps/post_turn_hooks.py | 3 +- .../task_steps/register_env_triggers.py | 6 +- tst/integration/env/gateway/gateway_test.py | 44 ++-- .../test_next_env_event_hook_local.py | 3 +- .../test_register_env_triggers_local.py | 25 ++- .../envs/test_mcp_server_plugin_provider.py | 2 +- tst/unit/env/envs/test_multi_env_name.py | 2 +- .../env/gateway/rest_proxy_stream_test.py | 212 ++++++++++++++++++ tst/unit/env/gateway/triggers_website_test.py | 188 +++++++++++++++- tst/unit/env/test_env.py | 10 +- .../env_gateway_provider_test.py | 52 ++++- tst/unit/task_step/deploy_env_test.py | 36 +++ .../task_step/local_container_urls_test.py | 32 ++- .../test_deploy_agent_role_delivery.py | 26 ++- tst/unit/task_step/test_load_artifact_file.py | 16 ++ tst/unit/task_step/test_mcp_cli_builder.py | 10 + .../test_prompt_agent_post_turn_hooks.py | 22 ++ .../task_step/test_register_env_triggers.py | 6 +- 43 files changed, 1074 insertions(+), 155 deletions(-) diff --git a/packages/agentenv-protocol/examples/website/notes/README.md b/packages/agentenv-protocol/examples/website/notes/README.md index 67d54de43..eade87bde 100644 --- a/packages/agentenv-protocol/examples/website/notes/README.md +++ b/packages/agentenv-protocol/examples/website/notes/README.md @@ -20,8 +20,10 @@ them. The page shows the universe's time. `edit_note`, `toggle_note` and `delete_note`. Errors are raised as `NotFound`, `BadRequest` and `Unauthorized`. - **Roles.** `all_notes` (`GET /api/admin/notes`) is declared with `roles=["admin"]`. It answers only a caller that sends `AgentEnv-Role: admin` with `AgentEnv-Role-Key` equal to the container's `WEBSITE_ROLE_KEY`; any other caller - gets 403. The gateway passes on whatever role a browser sends, so without the key a page could claim `admin`; with no - `WEBSITE_ROLE_KEY` set, every caller is `default` and the route is closed. + gets 403. A deploy with `keyed_roles` gives the site its own role key as `WEBSITE_ROLE_KEY` and its gateway sends it + on only beside a role it verified with the same key, so a page cannot claim `admin` through it. Any other deploy sets + no key, and its gateway passes on whatever role a browser sends; with no `WEBSITE_ROLE_KEY` set, every caller is + `default` and the route is closed. - **Users.** `ApiContext.user` is the user in the signed session cookie. With no cookie, it is the acting user (the store's `current_user`), and `GET /api/session` signs that user in. `POST /api/sign-in` switches to another seeded user. A user's notes are invisible to every other user, and another user's note answers 404. diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/environment.py b/packages/agentenv-protocol/src/agentenv_protocol/website/environment.py index e9eaee9fa..ac05ba6be 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/website/environment.py @@ -16,12 +16,14 @@ import re from pathlib import Path from typing import Any, ClassVar, NamedTuple, Optional +from urllib.parse import unquote from starlette.applications import Starlette from starlette.datastructures import MutableHeaders from starlette.middleware import Middleware from starlette.requests import Request from starlette.responses import FileResponse, JSONResponse, Response +from starlette._utils import get_route_path from starlette.routing import Route from ..agent_env_environment import ( @@ -39,7 +41,8 @@ from .clock import CLOCK_EXTENSION_URI, UNIVERSE_NOW_HEADER, UniverseClock, clock_patch_script, iso from .errors import NotFound from .readiness import WEBSITE_EXTENSION_URI, WEBSITE_READY_OP, WebsiteReadiness -from .routes import ApiContext, ApiRoute, endpoint, error_response, openapi, openapi_path, request_role, reserved_path, routes_of +from .routes import (ApiContext, ApiRoute, endpoint, error_response, forwarded_prefix, openapi, openapi_path, request_role, + reserved_path, routes_of) from .session import SESSION_COOKIE, SessionSigner, expires_at from .store import Collections @@ -284,9 +287,11 @@ def _check_health(cls: type) -> None: class _Forwarded: - """Puts the address the browser reached the site at in the scope, from ``X-Forwarded-Host`` and - ``X-Forwarded-Proto``, before anything routes it, so redirects, ``request.url`` and a mounted app's ``url_for`` - name that address and not the internal one the proxy reached the server at.""" + """Puts the address the browser reached the site at in the scope, from ``X-Forwarded-Host``, + ``X-Forwarded-Proto`` and ``X-Forwarded-Prefix``, before anything routes it, so redirects, ``request.url`` and a + mounted app's ``url_for`` name that address and not the internal one the proxy reached the server at. A prefix + becomes the scope's ``root_path`` and leads its ``path``, as ASGI has a mounted app's, so routing still sees the + site's own path.""" def __init__(self, app: Any) -> None: self.app = app @@ -301,6 +306,12 @@ async def __call__(self, scope: dict, receive: Any, send: Any) -> None: scope["headers"] = [(k, v) for k, v in scope["headers"] if k != b"host"] + [(b"host", host.encode())] if proto in ("http", "https"): scope["scheme"] = proto if scope["type"] == "http" else {"http": "ws", "https": "wss"}[proto] + prefix = forwarded_prefix(_first(headers.get(b"x-forwarded-prefix"))) + if prefix: + scope["root_path"] = unquote(prefix) + scope.get("root_path", "") + scope["path"] = unquote(prefix) + scope["path"] + if scope.get("raw_path") is not None: + scope["raw_path"] = prefix.encode("latin-1") + scope["raw_path"] await self.app(scope, receive, send) @@ -381,7 +392,7 @@ async def __call__(self, scope: dict, receive: Any, send: Any) -> None: if self.app is not None: await self.app(scope, receive, send) return - path, reading = scope["path"], scope["method"] in ("GET", "HEAD") + path, reading = get_route_path(scope) or "/", scope["method"] in ("GET", "HEAD") if reserved_path(path): return await error_response(NotFound("not found"))(scope, receive, send) if reading and (file := self._file(path)) is not None: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/routes.py b/packages/agentenv-protocol/src/agentenv_protocol/website/routes.py index 3059e78bb..939803342 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/routes.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/website/routes.py @@ -53,8 +53,9 @@ class ApiRoute(NamedTuple): fn: Callable def advertised(self) -> dict: - return {"method": self.method, "path": openapi_path(self.path), "name": self.name, - "roles": list(self.roles) or None} + """The route as the card names it: its path is the Starlette path it is served at, converters and all + (``/files/{rest:path}``), so a gateway matches the requests it answers, nested paths included.""" + return {"method": self.method, "path": self.path, "name": self.name, "roles": list(self.roles) or None} class _Api: @@ -159,11 +160,23 @@ def apply(self, response: Response) -> Response: def base_url(request: Request) -> str: """The site's address as the browser reached it: ``X-Forwarded-Proto`` and ``X-Forwarded-Host`` when a proxy - sent them, the request's own scheme and ``Host`` otherwise.""" + sent them, the request's own scheme and ``Host`` otherwise, under the path ``X-Forwarded-Prefix`` names when the + proxy serves the site below one (the gateway's ``/website/``).""" proto = request.headers.get("x-forwarded-proto", "").split(",")[0].strip() or request.url.scheme host = (request.headers.get("x-forwarded-host", "").split(",")[0].strip() or request.headers.get("host") or request.url.netloc) - return f"{proto}://{host}" + return f"{proto}://{host}{forwarded_prefix(request.headers.get('x-forwarded-prefix', ''))}" + + +def forwarded_prefix(value: str) -> str: + """The path an ``X-Forwarded-Prefix`` value names, without its trailing slash; empty for none, and for one that is + not an absolute path or would make the site's URLs name another host (``//host``). + + >>> [forwarded_prefix(v) for v in ("/website/notes/", "/a, /b", "", "/", "website", "//evil.example", "/a//b")] + ['/website/notes', '/a', '', '', '', '', ''] + """ + prefix = value.split(",")[0].strip().rstrip("/") + return prefix if prefix.startswith("/") and "//" not in prefix and prefix.isprintable() else "" ROLE_KEY_ENV = "WEBSITE_ROLE_KEY" @@ -172,7 +185,8 @@ def base_url(request: Request) -> str: def request_role(request: Request) -> str: """The role ``AgentEnv-Role`` names, believed only from a caller that also sends ``AgentEnv-Role-Key`` equal to - ``WEBSITE_ROLE_KEY``; every other caller, and every caller of a site with no key set, is ``default``. The gateway + ``WEBSITE_ROLE_KEY``; every other caller, and every caller of a site with no key set, is ``default``. A deploy with + ``keyed_roles`` sets the key and its gateway sends it on beside a role it verified the same way; any other gateway forwards whatever role a browser sends, so without the key a page could claim any role.""" claimed = caller_from_request(None, request.headers).role or DEFAULT_ROLE key = os.environ.get(ROLE_KEY_ENV, "") diff --git a/packages/agentenv-protocol/tests/test_website_api.py b/packages/agentenv-protocol/tests/test_website_api.py index b7754cf9d..e461ce097 100644 --- a/packages/agentenv-protocol/tests/test_website_api.py +++ b/packages/agentenv-protocol/tests/test_website_api.py @@ -320,6 +320,19 @@ def test_a_redirect_names_the_forwarded_address(client: TestClient) -> None: assert client.get("/api/items/", follow_redirects=False).headers["location"] == "http://testserver/api/items" +def test_a_site_behind_a_forwarded_prefix_names_it_in_its_address_and_redirects_and_routes_its_own_paths( + client: TestClient) -> None: + behind = {"Host": "internal-x.invalid", "X-Forwarded-Host": "gw.example", "X-Forwarded-Proto": "https", + "X-Forwarded-Prefix": "/website/shop/"} + assert client.get("/api/whoami", headers=behind).json()["base_url"] == "https://gw.example/website/shop" + assert client.get("/api/items/1", headers=behind).json()["title"] == "mug" + redirected = client.get("/api/items/", follow_redirects=False, headers=behind) + assert (redirected.status_code, redirected.headers["location"]) == (307, "https://gw.example/website/shop/api/items") + for unusable in ("//evil.example", "website", ""): + seen = client.get("/api/whoami", headers={**behind, "X-Forwarded-Prefix": unusable}).json() + assert seen["base_url"] == "https://gw.example", unusable + + def test_parameterized_routes_never_take_the_data_plane_or_card() -> None: @website(domains=["wiki.example"]) class Wiki(WebsiteEnvironment): diff --git a/packages/agentenv-protocol/tests/test_website_environment.py b/packages/agentenv-protocol/tests/test_website_environment.py index 5920b3995..05c118591 100644 --- a/packages/agentenv-protocol/tests/test_website_environment.py +++ b/packages/agentenv-protocol/tests/test_website_environment.py @@ -102,6 +102,18 @@ def test_static_files_and_index_fallback_answer_for_any_host(client: TestClient) assert client.post("/orders").status_code == 404 +def test_behind_a_forwarded_prefix_the_site_serves_its_own_paths_and_the_fallback_still_guards_reserved_ones( + client: TestClient) -> None: + behind = {"X-Forwarded-Prefix": "/website/shop"} + assert "home" in client.get("/", headers=behind).text + assert client.get("/assets/app.js", headers=behind).text == "console.log(1)" + assert "home" in client.get("/orders/42/edit", headers=behind).text + assert client.get("/api/orders", headers=behind).json() == [] + assert client.get("/.well-known/agent-env.json", headers=behind).json()["name"] == "Shop" + assert client.get("/website/shop/assets/app.js", headers=behind).status_code == 404 + assert client.post("/agentenv/x", headers=behind).status_code in (404, 405) + + def test_a_link_out_of_the_static_folder_is_never_served(dist: Path) -> None: secret = dist.parent / "secret.txt" (dist / "leak.txt").symlink_to(secret) diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 24c643be7..baa2ccf8f 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -275,21 +275,25 @@ def _skill_extension(deployed: "DeployedA2AAgent") -> dict: @staticmethod async def install_cli( deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", gateway_url: str, container: str | None = None, - on_host: bool = False, + on_host: bool = False, role_key: str | None = None, ) -> str: """Install a CliArtifact bundle into the agent's filesystem at /opt/cli//, in ``container`` when install_agent put the agent in one, or on the VM host, ``on_host``, when it put it there. Writes a sidecar - bin/.env with the gateway URL and chmod +x's the entrypoint. Idempotent. Returns the path to the entrypoint.""" + bin/.env with the gateway URL, and with ``role_key``, the key of that gateway's deploy, when it has one, so + each installed CLI sends its own deploy's key beside its role; then chmod +x's the entrypoint. Idempotent. + Returns the path to the entrypoint.""" from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container, stage_files_into_container + from agent_env.task_step.task_steps.mcp_cli_builder.codegen import GATEWAY_URL_ENV_VAR, ROLE_KEY_ENV_VAR from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider provider = build_sandbox_provider(deployed.sandbox_type) if deployed.sandbox_type else get_agent_sandbox_provider() sandbox = await provider.get_sandbox(deployed.sandbox_id) target_root = f"/opt/cli/{cli_artifact.command_name}" entrypoint_path = f"{target_root}/{cli_artifact.entrypoint}" env_dir = posixpath.dirname(entrypoint_path) + sidecar = f"{GATEWAY_URL_ENV_VAR}={gateway_url}\n" + (f"{ROLE_KEY_ENV_VAR}={role_key}\n" if role_key else "") if on_host: - await sandbox.write_host_file(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n".encode(), f"{env_dir}/{GATEWAY_ENV_FILE}") + await sandbox.write_host_file(sidecar.encode(), f"{env_dir}/{GATEWAY_ENV_FILE}") await stage_files_into_container( sandbox, cli_artifact.get_cli_files().get_file_artifacts(), target_root, on_host=True) await sandbox.exec_script(f"chmod +x {shlex.quote(entrypoint_path)}") @@ -297,7 +301,7 @@ async def install_cli( return entrypoint_path # write .env first — write_file_from_text creates the parent dir; write_file_from_object does not. - await sandbox.write_file_from_text(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n", f"{env_dir}/{GATEWAY_ENV_FILE}", + await sandbox.write_file_from_text(sidecar, f"{env_dir}/{GATEWAY_ENV_FILE}", **into_container(container)) for rel_path, file_artifact in cli_artifact.get_cli_files().get_file_artifacts().items(): diff --git a/src/agent_env/cli/env/deploy.py b/src/agent_env/cli/env/deploy.py index fe1984c28..f7afb1040 100644 --- a/src/agent_env/cli/env/deploy.py +++ b/src/agent_env/cli/env/deploy.py @@ -48,7 +48,10 @@ @click.option("--env-state-instance-id", default=None, help="Attach against an EXISTING env state instance instead of creating a fresh store " "(its own type selects the provider).") -def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: str, sandbox: str, service_db_env_id: str | None, gateway_env_id: str | None, website_browser_env_id: str | None, disk_size_gb: float, cpu: float | None, memory_mb: int | None, env_state_type: str | None, env_state_instance_id: str | None): +@click.option("--keyed-roles", is_flag=True, default=False, + help="Draw a role key for the deploy: its gateway then believes a role other than default only from a " + "caller that sends the key beside it.") +def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: str, sandbox: str, service_db_env_id: str | None, gateway_env_id: str | None, website_browser_env_id: str | None, disk_size_gb: float, cpu: float | None, memory_mb: int | None, env_state_type: str | None, env_state_instance_id: str | None, keyed_roles: bool): """Deploy an MCP server environment.""" from agent_env.providers import build_sandbox_provider from agent_env.config import get_config @@ -89,6 +92,8 @@ def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: if env_state_instance_id: deploy_kwargs["env_state_instance_id"] = env_state_instance_id click.echo(f"Env state instance id: {env_state_instance_id}") + if keyed_roles: + deploy_kwargs["keyed_roles"] = True try: with run_scope(env.id): deployed_env = asyncio.run(env.deploy(**deploy_kwargs)) diff --git a/src/agent_env/env/env.py b/src/agent_env/env/env.py index ef3373e2c..9e50ebabf 100644 --- a/src/agent_env/env/env.py +++ b/src/agent_env/env/env.py @@ -17,7 +17,7 @@ from agent_env.store.routing import namespace_routing_enabled from agent_env.utils.paths import validate_relative_filename -from .gateway import GatewayMode +from .gateway import AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_KEY_HEADER, GatewayMode from .gateway.constants import WELL_KNOWN_PATH from .store import get_env_store @@ -83,6 +83,7 @@ class DeployedEnv: # The provider that made the record, so the class it loads as (see from_dict); capabilities still come from the card. env_provider_type: Optional[str] = None environment_card_url: Optional[str] = None + role_key: Optional[str] = None # The env card as last read, and when (ISO 8601, UTC); written at deploy from the readiness probe. environment_card: Optional[dict] = None environment_card_read_at_utc: Optional[str] = None @@ -122,6 +123,11 @@ def control_headers_for(self, url: str) -> dict[str, str]: """What a call to ``url`` sends of the env's control token: ``control_headers`` when ``url`` is the env's address or under it, as a step's grant gives them, and none anywhere else, an override's address say.""" return protocol_v1.granted_headers(url, {self.environment_url: self.control_headers}) + def role_headers(self, role: str) -> dict[str, str]: + """The headers a request to the env sends to act as `role`: the role, and the deploy's `role_key` beside it, + without which a gateway that has one takes the request for `default`'s. A record from before the key, of a deploy without + `keyed_roles`, or of one with no gateway, has none, and its role header alone is believed.""" + return {AGENT_ENV_ROLE_HEADER: role, **({AGENT_ENV_ROLE_KEY_HEADER: self.role_key} if self.role_key else {})} def supports(self, uri: str, method: str) -> bool: """Whether the stored env card offers `method` on the extension at `uri`; no network.""" @@ -200,6 +206,7 @@ def _fields_from(cls, data: dict) -> dict: created_at_utc=data.get("created_at_utc"), expires_at_utc=data.get("expires_at_utc"), control_token=data.get("control_token"), + role_key=data.get("role_key"), ) diff --git a/src/agent_env/env/envs/mcp_server.py b/src/agent_env/env/envs/mcp_server.py index 2f950f77a..96fad5ec0 100644 --- a/src/agent_env/env/envs/mcp_server.py +++ b/src/agent_env/env/envs/mcp_server.py @@ -107,13 +107,13 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: @environment_card when it isn't set. Raises BundleError listing every problem.""" return ctx.accept_env(data, cls, named_by="image") - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False) -> DeployedEnv: # In container mode the server runs in its own container, so loads stage there, as for a MultiEnv child. return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, - artifact_id=artifact_id, artifact_version=artifact_version, + artifact_id=artifact_id, artifact_version=artifact_version, **({"keyed_roles": True} if keyed_roles else {}), ) def deploy_refusal(self, *, sandbox_type: str | None = None, **options) -> str | None: diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index e193771d2..d1da07009 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -91,6 +91,7 @@ def __init__(self, id: str, version: Optional[int], mcp_server_envs: list[MCPSer self._gateway_mode = GatewayMode.PERFORMANCE self._env_provider = None self._mcp_server_name: Optional[str] = None + self._role_key: Optional[str] = None self._deployed: Optional[DeployedEnv] = None self._replaced: list[tuple] = [] # the provider and sandbox of each deployment a later deploy replaced, for close() @@ -156,15 +157,16 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: ctx.refuse(problems) return fields - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False) -> DeployedEnv: deployed_env = await deploy_through_provider( self, environment_name=None, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, - artifact_id=artifact_id, artifact_version=artifact_version, + artifact_id=artifact_id, artifact_version=artifact_version, **({"keyed_roles": True} if keyed_roles else {}), ) self._gateway_mode = gateway_mode self._mcp_server_name = deployed_env.mcp_server_name + self._role_key = deployed_env.role_key _hand_provider_to_child_envs(self) _hand_record_to_child_envs(self, deployed_env) return deployed_env @@ -190,6 +192,7 @@ async def from_deployed_env(cls, deployed: DeployedEnv) -> MultiEnv: _hand_record_to_child_envs(env, deployed) env._instance_id = deployed.instance_id env._mcp_server_name = deployed.mcp_server_name + env._role_key = deployed.role_key return env async def load_environment_artifact(self, environment_artifact: EnvironmentArtifact) -> None: diff --git a/src/agent_env/env/envs/website.py b/src/agent_env/env/envs/website.py index 64d4b2d5c..702dd4c61 100644 --- a/src/agent_env/env/envs/website.py +++ b/src/agent_env/env/envs/website.py @@ -169,12 +169,12 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: ctx.refuse([ctx.config_problem(problem)]) return ctx.accept_env(data, cls, named_by="backend_image") - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False) -> DeployedEnv: return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, - artifact_id=artifact_id, artifact_version=artifact_version, + artifact_id=artifact_id, artifact_version=artifact_version, **({"keyed_roles": True} if keyed_roles else {}), ) def deploy_refusal(self, *, sandbox_type: str | None = None, **options) -> str | None: diff --git a/src/agent_env/env/gateway/__init__.py b/src/agent_env/env/gateway/__init__.py index 6cd59de6f..cb94520d2 100644 --- a/src/agent_env/env/gateway/__init__.py +++ b/src/agent_env/env/gateway/__init__.py @@ -6,6 +6,7 @@ from .constants import ( AGENT_ENV_HIDDEN_BY_DEFAULT_META_KEY, AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_KEY_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, DEFAULT_ROLE, @@ -31,6 +32,7 @@ class InternalMCPServer: "AGENT_ENV_GATEWAY_MCP_PORT", "AGENT_ENV_HIDDEN_BY_DEFAULT_META_KEY", "AGENT_ENV_ROLE_HEADER", + "AGENT_ENV_ROLE_KEY_HEADER", "AGENT_ENV_ROLE_META_KEY", "AGENT_ENV_SESSION_META_KEY", "DEFAULT_ROLE", diff --git a/src/agent_env/env/gateway/caller.py b/src/agent_env/env/gateway/caller.py index 4c5c50edb..1546fe2c9 100644 --- a/src/agent_env/env/gateway/caller.py +++ b/src/agent_env/env/gateway/caller.py @@ -1,6 +1,7 @@ """Who the gateway acts for on one inbound request. -A ``Caller`` is the role the request's ``AgentEnv-Role`` header named (``default`` when it named none) and the key of +A ``Caller`` is the role the request's ``AgentEnv-Role`` header named (``default`` when it named none, or named another +without the deploy's ``AgentEnv-Role-Key`` when the deploy has one) and the key of the MCP session the request belongs to (empty on the REST and harness paths, which have no session). It is bound in three scopes, which read in order in ``gateway.py``: per HTTP request by the served app (the role), per inbound MCP session by the session lifespan (the key, kept for every tools/call of that session), and never per process. Every diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 30bd00031..90f5a609d 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -126,9 +126,10 @@ def sources(names: list[str]) -> str: TRIGGER_EXECUTOR_HOURS_TIME_ZONE = "executor_hours_time_zone" TRIGGER_NEXT_EVENT_AT = "next_event_at" TRIGGER_NOTIFY = "notify" +TRIGGER_WEBSITE_OP = "website_op" TRIGGER_FEATURES = (TRIGGER_NAMED_EXECUTORS, TRIGGER_PROVOKING_ROLES, TRIGGER_ACTION_DELAY, TRIGGER_EXECUTOR_HOURS, TRIGGER_UNLESS, TRIGGER_WEBSITE, TRIGGER_EXECUTOR_HOURS_TIME_ZONE, TRIGGER_NEXT_EVENT_AT, - TRIGGER_NOTIFY) + TRIGGER_NOTIFY, TRIGGER_WEBSITE_OP) # The one `source` a /step call may name instead of step_api, with the id of the hook artifact that made it. STEP_SOURCE_POST_TURN_HOOK = "post_turn_hook" @@ -142,6 +143,9 @@ def sources(names: list[str]) -> str: ) AGENT_ENV_ROLE_HEADER = "AgentEnv-Role" +AGENT_ENV_ROLE_KEY_HEADER = "AgentEnv-Role-Key" +GATEWAY_ROLE_KEY_ENV = "AGENTENV_ROLE_KEY" +WEBSITE_ROLE_KEY_ENV = "WEBSITE_ROLE_KEY" # 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" diff --git a/src/agent_env/env/gateway/entrypoint.py b/src/agent_env/env/gateway/entrypoint.py index b0b936ae9..b6d469a17 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 DEFAULT_MCP_SERVER_NAME, GATEWAY_ROLE_KEY_ENV 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, + role_key=os.environ.get(GATEWAY_ROLE_KEY_ENV) or None, ) gateway.run() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index b7dd22091..94dfd2366 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 @@ -41,6 +42,7 @@ from .constants import ( AGENT_ENV_HIDDEN_BY_DEFAULT_META_KEY, AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_KEY_HEADER, AGENT_ENV_ROLE_META_KEY, # noqa: F401 (the server module serves the whole vocabulary, used in this file or not) AGENT_ENV_SESSION_META_KEY, # noqa: F401 CARD_FETCH_TIMEOUT_S, @@ -72,7 +74,19 @@ from .child_transport import child_http_client from .clock import Clock, ClockError, _iso from .get_time import GET_TIME_TOOL_NAME, build_get_time_tool -from .rest_proxy import ROLE_HEADER_BYTES, UpstreamResponse, fetch_child_card, forward, operations, rewrite_child_card, upstream_url +from .rest_proxy import ( + ROLE_HEADER_BYTES, + ROLE_KEY_HEADER_BYTES, + UpstreamResponse, + WebsiteRoute, + fetch_child_card, + forward, + operations, + rewrite_child_card, + route_name, + upstream_url, + website_routes, +) from .triggers import TriggerEngine, TriggerError from .notifications import CHECK_TOOL_NAMES, NotificationError, Notifications, build_check_tools @@ -119,7 +133,10 @@ def __init__( rest_proxy_urls: dict[str, str] | None = None, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, service_db_url: str | None = None, + role_key: str | None = None, ): + """`role_key` is the deploy's secret: with one, a request is believed to act as a role other than `default` + only when it carries the key beside the role; without one, as before it, the role header alone is believed.""" self.host = host self.port = port self.server_name = server_name @@ -129,6 +146,8 @@ def __init__( self._site_reserved: dict[str, frozenset[str]] = {} self.gateway_mode = gateway_mode self._service_db_url = service_db_url + self._role_key = role_key or None + self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: self._db_conn = psycopg2.connect(service_db_url) @@ -240,13 +259,15 @@ async def get_state(request: Request) -> Response: async def step_endpoint(request: Request) -> Response: return await self._handle_step(request) - async def _trigger_write(request: Request, method) -> Response: - """Shared body for the trigger write routes: parse JSON, invoke the engine method, - map a TriggerError (bad config) to a 400.""" + async def _trigger_write(request: Request, method, prepare=None) -> Response: + """Shared body for the trigger write routes: parse JSON, await `prepare` with it, invoke the engine + method, map a TriggerError (bad config) to a 400.""" try: body = await request.json() except Exception: return self._json_response({"ok": False, "error": "Invalid JSON body"}, 400) + if prepare is not None: + await prepare(body) try: return self._json_response(method(body), 200) except TriggerError as e: @@ -254,7 +275,7 @@ async def _trigger_write(request: Request, method) -> Response: @self._mcp.custom_route("/triggers/register", methods=["POST"]) async def triggers_register(request: Request) -> Response: - return await _trigger_write(request, self._trigger_engine.register) + return await _trigger_write(request, self._trigger_engine.register, prepare=self._read_trigger_website_routes) @self._mcp.custom_route("/triggers/remove", methods=["POST"]) async def triggers_remove(request: Request) -> Response: @@ -408,7 +429,7 @@ def run(self) -> None: anyio.run(server.serve) # The three scopes a Caller is bound in, innermost first. Request: the served app binds the role the - # AgentEnv-Role header named. Session: mcp's Server.run enters the session lifespan once per inbound MCP + # AgentEnv-Role header named, as `_attributed_role` believes it. Session: mcp's Server.run enters the session lifespan once per inbound MCP # session, in the session's server task before any handler task is spawned, so the key bound there is inherited # by every tools/call of that session (and the role is the one its initialize request carried, which that task # inherited). Process: child sessions and the time-trigger driver, outer to the MCP session manager's lifespan. @@ -423,11 +444,13 @@ async def app(scope, receive, send): if scope["type"] != "http": await inner(scope, receive, send) return - role = DEFAULT_ROLE + claimed, given_key = None, None for name, value in scope.get("headers", ()): - if name == ROLE_HEADER_BYTES: - role = value.decode(errors="ignore").strip() or DEFAULT_ROLE - break + if name == ROLE_HEADER_BYTES and claimed is None: + claimed = value.decode(errors="ignore") + elif name == ROLE_KEY_HEADER_BYTES and given_key is None: + given_key = value.decode(errors="ignore") + role = self._attributed_role(claimed or "", given_key or "") token = Caller(role=role, session=Caller.current().session).bind() try: await inner(scope, receive, send) @@ -603,7 +626,8 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] | role = self._request_role(request) mcp_server = url_map is self.rest_proxy_urls and service_name.startswith("mcp-") if request.method in _SAFE_METHODS: - return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server) + return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server, + role_key=self._role_key) changelog_before = self._query_changelog_id() attribution = self._attribution(role) # at the start: a fire may be gone by the time the response ends @@ -611,12 +635,16 @@ async def log_write(status: int) -> None: changelog_after = self._query_changelog_id() if changelog_after is None or changelog_after <= (changelog_before or 0): return # only a request that wrote changelog rows is a write window + path = "/" + request.path_params.get("path", "") + op = (self._website_op(service_name, request.method, site_request_path(request) or path) + if url_map is self.website_urls else None) await self._log_event({ "event_type": "http_request", "route": "website" if url_map is self.website_urls else "svc", "service": service_name, "method": request.method, - "path": "/" + request.path_params.get("path", ""), + "path": path, + **({"op": op} if op else {}), "status": status, **attribution, "changelog_id_before": changelog_before, @@ -624,7 +652,7 @@ async def log_write(status: int) -> None: }) return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server, - on_done=log_write) + role_key=self._role_key, on_done=log_write) async def _proxy_website_request(self, request: Request) -> Response: """Proxy a website request; with a website or state trigger registered, feed it to the engine and hold barriers. @@ -640,9 +668,10 @@ async def _proxy_website_request(self, request: Request) -> Response: app = request.path_params["service_name"] engine = self._trigger_engine if app in self.website_urls and (engine.has_website_triggers or engine.has_state_triggers): + path = site_request_path(request) or "/" + request.path_params.get("path", "") if isinstance(response, UpstreamResponse): - if not engine.may_hold_website_request(self._request_role(request), app, request.method, - "/" + request.path_params.get("path", ""), response.status_code): + if not engine.may_hold_website_request(self._request_role(request), app, request.method, path, + response.status_code, self._website_op(app, request.method, path)): response.after_body(lambda status: self._fire_website_triggers(request, app, status)) # answered at its last byte return response await response.read_ahead_if_sized(self.BARRIER_READ_AHEAD_MAX_BYTES) # its window ends before its triggers run @@ -656,16 +685,55 @@ async def _fire_website_triggers(self, request: Request, app: str, try: # Only a website trigger reads the request's args; a state re-check needs just the method and status. args = _request_args(request, await request.body()) if self._trigger_engine.has_website_triggers else {} - return self._trigger_engine.on_website_request(self._request_role(request), app, request.method, - "/" + request.path_params.get("path", ""), status, args) + path = site_request_path(request) or "/" + request.path_params.get("path", "") + return self._trigger_engine.on_website_request(self._request_role(request), app, request.method, path, + status, args, self._website_op(app, request.method, path)) except Exception: logger.exception("Trigger evaluation failed (fail-open)") return [], None - @staticmethod - 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 + def _request_role(self, request: Request) -> str: + """The role the gateway attributes a REST request to: the first AgentEnv-Role header, as `_attributed_role` + believes it, else the default role.""" + return self._attributed_role(request.headers.get(AGENT_ENV_ROLE_HEADER, ""), + request.headers.get(AGENT_ENV_ROLE_KEY_HEADER, "")) + + def _attributed_role(self, claimed: str, given_key: str) -> str: + """The role a request claiming `claimed` acts as: `default` when it claims none, and when the deploy has a + role key the request did not send. Without a key (a deploy without keyed_roles) the claim alone is believed.""" + role = claimed.strip() or DEFAULT_ROLE + if self._role_key is None or role == DEFAULT_ROLE: + return role + return role if hmac.compare_digest(given_key.encode(), self._role_key.encode()) else DEFAULT_ROLE + + async def _read_trigger_website_routes(self, body: Any) -> None: + """Read the routes of each website a website trigger in a register `body` names an `op` of, so registration + can check the op and requests can be matched by it; one whose card could not be read before is read again.""" + triggers = body.get("triggers") if isinstance(body, dict) else None + apps = {trigger["when"]["app"] for trigger in triggers if isinstance(trigger, dict) + and isinstance(trigger.get("when"), dict) and trigger["when"].get("type") == "website" + and trigger["when"].get("op") is not None and isinstance(trigger["when"].get("app"), str) + and trigger["when"]["app"] in self.website_urls} if isinstance(triggers, list) else set() + await asyncio.gather(*(self._read_website_routes(app, again=True) for app in apps)) + + async def _read_website_routes(self, app: str, *, again: bool = False) -> None: + """Read website `app`'s routes from its card once (`again`: once more, when its card could not be read). The + env card's composition reads them too, so a deploy, which reads that card, leaves them read.""" + if app in self._website_routes and not (again and self._website_routes[app] is None): + return + async with httpx.AsyncClient() as client: + card = await fetch_child_card(client, app, self.rest_proxy_urls.get(app) or self.website_urls[app]) + self._website_routes[app] = website_routes(card) if card is not None else None + + def website_route_names(self, app: str) -> set[str] | None: + """The names of the routes website `app`'s card names, as last read; None when it has not been read.""" + routes = self._website_routes.get(app) + return None if routes is None else {route.name for route in routes if route.name is not None} + + def _website_op(self, app: str, method: str, path: str) -> str | None: + """The name of the route of website `app` a request reached, by the routes last read from its card; never a + read of its own, so answering a request waits on no card.""" + return route_name(self._website_routes.get(app) or [], method, path) async def _site_reserved_paths(self, app: str) -> frozenset[str] | None: """The paths ``app``'s card advertises for agentenv (``site_reserved_paths``), read once from its server, the one @@ -702,8 +770,11 @@ def _rewrite_child_card(self, key: str, card: dict) -> dict: return rewrite_child_card(key, card, self._gateway_extensions()) async def _fetch_child_card(self, client: httpx.AsyncClient, key: str, base_url: str) -> dict | None: - """Fetch + rewrite one backing server's card; None (omit) on 404 / error / timeout / malformed JSON.""" + """Fetch + rewrite one backing server's card; None (omit) on 404 / error / timeout / malformed JSON. A + website's card read here also refreshes the routes its requests are named by.""" child = await fetch_child_card(client, key, base_url) + if child is not None and key in self.website_urls: + self._website_routes[key] = website_routes(child) return None if child is None else self._rewrite_child_card(key, child) def _gateway_extensions(self) -> list[dict]: @@ -1563,16 +1634,22 @@ def site_reserved_paths(card: dict) -> frozenset[str]: return frozenset(_normalized(p) for p in paths if isinstance(p, str) and p.startswith("/") and _normalized(p) != "/") +def site_request_path(request: Request) -> str | None: + """The path a site's server receives for a request on its route: the raw path forwarded, its dot segments removed + by httpx and then decoded. It differs from the path as routed here when an encoded slash sits beside a dot segment + (``public%2Fx/../api`` reaches ``/api``). None for a request-target the proxy can't build a URL for: ``forward()`` + refuses it.""" + try: + return httpx.URL(upstream_url("http://site", request)).path + except (ValueError, httpx.InvalidURL): + return None + + def reserved_site_request(request: Request, reserved: frozenset[str] = frozenset()) -> bool: """Whether a request on a site's route reaches its data plane or env card, or one of the paths ``reserved`` names - (``site_reserved_paths``): by the path as routed here, or by the one its server receives, which is the raw path - forwarded, its dot segments removed by httpx and then decoded. The two differ when an encoded slash sits beside a - dot segment (``public%2Fx/../agentenv``). A request-target the proxy can't build a URL for isn't reserved here: - ``forward()`` refuses it.""" + (``site_reserved_paths``): by the path as routed here, or by the one its server receives (``site_request_path``). + A request-target the proxy can't build a URL for isn't reserved here: ``forward()`` refuses it.""" if reserved_site_path("/" + request.path_params.get("path", ""), reserved): return True - try: - sent = httpx.URL(upstream_url("http://site", request)).path - except (ValueError, httpx.InvalidURL): - return False - return reserved_site_path(sent, reserved) + sent = site_request_path(request) + return sent is not None and reserved_site_path(sent, reserved) diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index 0812440ec..10a3b3ba2 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -10,12 +10,17 @@ round-trip through str), minus the hop-by-hop set, in both directions; - the role header is the gateway's to write: the client's copies are dropped and the role the gateway attributed the request to goes in their place, and a POST to a backing MCP server's own endpoint gets the same treatment - inside its JSON-RPC ``params._meta``; + inside its JSON-RPC ``params._meta``; so is the role key, which goes upstream only beside a role other than + ``default`` the key verified, so a site holding the same key believes the role; +- the upstream learns the address the client reached it at: ``X-Forwarded-Host`` and ``X-Forwarded-Proto`` (a front's + own, when one before the gateway sent them, else the request's ``Host`` and scheme) and ``X-Forwarded-Prefix``, the + route prefix (``/svc/{service}`` or ``/website/{service}``) the gateway strips, which a client never sets; - the upstream response and the client it came through are closed when the response ends, iterated or not. A backing server's env card names its endpoints relative to itself; :func:`rewrite_child_card` roots them under ``/svc/{key}`` so they resolve against the env's address, keeping the gateway's own routes where the gateway offers -the same operation. +the same operation. A website's card names its routes on ``urn:agentenv:website/v1``; :func:`website_routes` reads +them so :func:`route_name` can name the route a proxied request reached. """ from __future__ import annotations @@ -23,8 +28,9 @@ import json import logging import posixpath +import re from collections.abc import AsyncIterator, Awaitable, Callable -from typing import Any +from typing import Any, NamedTuple from urllib.parse import unquote import httpx @@ -35,9 +41,12 @@ # Relative import — gateway runs as a standalone container with no agent_env.* from .constants import ( AGENT_ENV_ROLE_HEADER, + AGENT_ENV_ROLE_KEY_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, CARD_FETCH_TIMEOUT_S, + DEFAULT_ROLE, + EXT_WEBSITE_URI, WELL_KNOWN_PATH, ) @@ -51,11 +60,16 @@ # 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() +ROLE_KEY_HEADER_BYTES = AGENT_ENV_ROLE_KEY_HEADER.lower().encode() + +FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX = "x-forwarded-host", "x-forwarded-proto", "x-forwarded-prefix" +_FORWARDED = frozenset(name.encode() for name in (FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX)) async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.Timeout, mcp_server: bool, - on_done: Callable[[int], Awaitable[None]] | None = None) -> Response: - """Forward `request` to the service at `base_url` as `role` and stream the response through. `mcp_server` says + role_key: str | None = None, on_done: Callable[[int], Awaitable[None]] | None = None) -> Response: + """Forward `request` to the service at `base_url` as `role`, vouched for by `role_key` when the deploy has one, + and stream the response through. `mcp_server` says the service is a backing MCP server, whose own endpoint takes the role in its JSON-RPC `_meta` as well as in the header. A 400 names a request-target the proxy cannot route; a 502 names an upstream that never answered. `on_done` is awaited with the status once the response has ended, after its last byte when streamed.""" @@ -64,7 +78,7 @@ async def forward(request: Request, base_url: str, *, role: str, timeout: httpx. except ValueError as e: # an escaped slash inside the route prefix; refuse rather than forward a wrong path return await _done(_json_response({"error": str(e)}, 400), on_done) content = await request.body() # /svc and /website request bodies are small (JSON-RPC, forms) - headers = upstream_headers(request, role) + headers = upstream_headers(request, role, role_key) if mcp_server and request.method == "POST" 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. @@ -101,18 +115,49 @@ def upstream_url(base_url: str, request: Request) -> str: return url -def upstream_headers(request: Request, role: str) -> httpx.Headers: +def upstream_headers(request: Request, role: str, role_key: str | None = None) -> 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.""" + and role key headers are the gateway's to set: the client's copies are dropped and the role the gateway + attributed the request to goes in their place, with `role_key` beside any role but `default`, so the upstream + never reads a value the gateway did not write. So are the forwarded-address headers (`forwarded_address`).""" 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]) + if k.lower() not in HOP_BY_HOP and k.lower() not in (ROLE_HEADER_BYTES, ROLE_KEY_HEADER_BYTES) + and k.lower() not in _FORWARDED]) headers[AGENT_ENV_ROLE_HEADER] = role + if role_key and role != DEFAULT_ROLE: + headers[AGENT_ENV_ROLE_KEY_HEADER] = role_key + headers.update(forwarded_address(request)) # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take. headers.setdefault("accept-encoding", "identity") return headers +def forwarded_address(request: Request) -> dict[str, str]: + """Where the client reached the proxied service: the host and scheme it addressed (as a front before the gateway + names them in its own ``X-Forwarded-Host``/``X-Forwarded-Proto``, else the request's ``Host`` and scheme; the + client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent. + A header with nothing to say is left out.""" + host = _first(request.headers.get(FORWARDED_HOST, "")) or request.headers.get("host", "") + proto = _first(request.headers.get(FORWARDED_PROTO, "")).lower() + address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else request.url.scheme, + FORWARDED_PREFIX: route_prefix(request)} + return {name: value for name, value in address.items() if value} + + +def _first(value: str) -> str: + return value.split(",")[0].strip() + + +def route_prefix(request: Request) -> str: + """The path the client reached the proxied service under, as it sent it: the request-target up to the `{path}` + remainder, without its trailing slash, e.g. ``/website/notes``, or under a root path ``/env/website/notes``.""" + raw_path = request.scope.get("raw_path") or request.scope["path"].encode() + decoded_rest = request.path_params.get("path", "") + if decoded_rest: + raw_path = raw_path[:-len(raw_remainder(raw_path, request.scope["path"], decoded_rest))] + return raw_path.decode("latin-1").rstrip("/") + + 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. @@ -326,5 +371,69 @@ def operations(params: dict) -> list[tuple[str | None, dict | None, dict, str, s if isinstance(holder.get("endpoint"), str) and holder["endpoint"]] +class WebsiteRoute(NamedTuple): + """One route a website's card names: its verb, the pattern its path matches (a Starlette converter's own, one + path segment for any other converter or none), and its name, None for a route the site builds in.""" + + method: str + pattern: re.Pattern + name: str | None + + +_CONVERTERS = {"path": ".*", "int": "[0-9]+", "float": r"[0-9]+(?:\.[0-9]+)?", + "uuid": "[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"} +_PATH_PARAM = re.compile(r"{\w+(?::(\w+))?}") + + +def website_routes(card: dict | None) -> list[WebsiteRoute]: + """The routes a website's card names on ``urn:agentenv:website/v1``, in its order: each one's method, its path + (``/api/notes/{note_id}``, with or without a Starlette converter such as ``{note_id:int}``) and its name. An entry + missing any of the three is skipped; a card naming none, or no card, has none. + + Before them come, unnamed, the built-in routes the card names, which answer before any of its routes would: a GET + to its ``health`` path, and a GET or DELETE to its ``session_path`` unless one of its routes is at that very path + (then the site serves its own there). A request one of them answers is named by no route. + + >>> [(r.method, r.pattern.pattern, r.name) for r in website_routes({"capabilities": {"extensions": [ + ... {"uri": "urn:agentenv:website/v1", "params": {"health": "/healthz", "routes": [ + ... {"method": "get", "path": "/api/notes/{note_id:int}", "name": "get_note"}, {"path": "/x"}]}}]}})] + [('GET', '/healthz', None), ('GET', '/api/notes/(?:[0-9]+)', 'get_note')] + """ + found: list[WebsiteRoute] = [] + for ext in ((card or {}).get("capabilities") or {}).get("extensions") or []: + params = ext.get("params") if isinstance(ext, dict) and ext.get("uri") == EXT_WEBSITE_URI else None + routes = params.get("routes") if isinstance(params, dict) else None + named = [route for route in (routes if isinstance(routes, list) else []) + if isinstance(route, dict) and all(isinstance(route.get(k), str) and route[k] + for k in ("method", "path", "name"))] + health, session = (params.get(k) if isinstance(params, dict) else None for k in ("health", "session_path")) + if isinstance(health, str) and health: + found.append(WebsiteRoute("GET", _path_pattern(health), None)) + if isinstance(session, str) and session and all(route["path"] != session for route in named): + found += [WebsiteRoute(verb, _path_pattern(session), None) for verb in ("GET", "DELETE")] + found += [WebsiteRoute(route["method"].upper(), _path_pattern(route["path"]), route["name"]) for route in named] + return found + + +def _path_pattern(path: str) -> re.Pattern: + parts, end = [], 0 + for param in _PATH_PARAM.finditer(path): + parts += [re.escape(path[end:param.start()]), f"(?:{_CONVERTERS.get(param.group(1) or '', '[^/]+')})"] + end = param.end() + return re.compile("".join(parts) + re.escape(path[end:])) + + +def route_name(routes: list[WebsiteRoute], method: str, path: str) -> str | None: + """The name of the first of `routes` a `method` request to `path` reaches (a HEAD reaches a GET route), else None, + as it is when that first one is a built-in route. + + >>> route_name(website_routes({"capabilities": {"extensions": [{"uri": "urn:agentenv:website/v1", "params": { + ... "routes": [{"method": "GET", "path": "/api/notes/{note_id}", "name": "get_note"}]}}]}}), "HEAD", "/api/notes/7") + 'get_note' + """ + verbs = {method.upper(), "GET"} if method.upper() == "HEAD" else {method.upper()} + return next((route.name for route in routes if route.method in verbs and route.pattern.fullmatch(path)), None) + + def _json_response(data: dict | list, status_code: int) -> Response: return Response(content=json.dumps(data), status_code=status_code, media_type="application/json") diff --git a/src/agent_env/env/gateway/triggers.py b/src/agent_env/env/gateway/triggers.py index 1017d3ef3..69dbf0d20 100644 --- a/src/agent_env/env/gateway/triggers.py +++ b/src/agent_env/env/gateway/triggers.py @@ -713,7 +713,8 @@ def _normalize_action_when(when: dict, tid: str) -> None: _normalize_provoked_when(when, tid) def _normalize_website_when(self, when: dict, tid: str) -> None: - """Validate a website `when`: a known app, upper-cased method, `path`/`status` predicates, `where`, `repeat`.""" + """Validate a website `when`: a known app, upper-cased method, the `op` (a route the site's card names), + `path`/`status` predicates, `where`, `repeat`.""" w = f"trigger '{tid}'" app = when.get("app") if not isinstance(app, str) or not app: @@ -726,6 +727,17 @@ def _normalize_website_when(self, when: dict, tid: str) -> None: raise TriggerError(f"{w}: when.method must be a non-empty HTTP method, e.g. 'POST'") else: when["method"] = when["method"].strip().upper() + if when.get("op") is None: + when.pop("op", None) + elif not isinstance(when["op"], str) or not when["op"]: + raise TriggerError(f"{w}: when.op must be a non-empty string naming one of the website's routes") + else: + names = self._gw.website_route_names(app) + if names is None: + raise TriggerError(f"{w}: when.op names a route of website {app!r}, whose card the gateway could not " + "read; match its path instead") + if when["op"] not in names: + raise TriggerError(f"{w}: when.op names no route of website {app!r} (routes: {sorted(names)})") for key in ("path", "status"): if when.get(key) is None: when.pop(key, None) @@ -921,17 +933,18 @@ def on_tool_call(self, role: str, name: str, arguments: dict, return pending, timeout_s def on_website_request(self, role: str, app: str, method: str, path: str, status: int, - args: dict) -> tuple[list[asyncio.Task], float | None]: - """Detection hook for an answered website request, before its response goes out; returns what on_tool_call does.""" + args: dict, op: str | None = None) -> tuple[list[asyncio.Task], float | None]: + """Detection hook for an answered website request, before its response goes out, `op` naming the route of + the site's card it reached; returns what on_tool_call does.""" provokable = [(tid, trig) for tid, trig in self._triggers.items() if self._provoked_by(trig, role)] if not provokable: return [], None if self._has_time_triggers: self._eval_time_triggers() - request = {"app": app, "method": method, "path": path, "status": status} + request = {"app": app, "method": method, "path": path, "status": status, **({"op": op} if op else {})} provoking = {"source": "website", **request, "role": role} matched = [(tid, trig) for tid, trig in provokable - if _website_matches(trig["spec"]["when"], app, method, path, status)] + if _website_matches(trig["spec"]["when"], app, method, path, status, op)] pending, timeout_s = self._detect(matched, provoking, {"args": args, "request": request}) if method not in _SAFE_METHODS and status < 400: self._reevaluate_states(provokable, provoking) @@ -963,10 +976,12 @@ def _reevaluate_states(self, provokable: list[tuple[str, dict]], provoking: dict trig["evaluating"] = True self._schedule(self._eval_state(tid, trig, provoking)) - def may_hold_website_request(self, role: str, app: str, method: str, path: str, status: int) -> bool: + def may_hold_website_request(self, role: str, app: str, method: str, path: str, status: int, + op: str | None = None) -> bool: """Whether a `provoking_call` barrier could hold this website request: a matcher check, `where` and status aside.""" return any((trig["spec"].get("barrier") or {}).get("at") == "provoking_call" and self._provoked_by(trig, role) - and _website_matches(trig["spec"]["when"], app, method, path, status) for trig in self._triggers.values()) + and _website_matches(trig["spec"]["when"], app, method, path, status, op) + for trig in self._triggers.values()) @property def has_website_triggers(self) -> bool: @@ -1608,10 +1623,13 @@ def _normalize_provoked_when(when: dict, tid: str) -> None: raise TriggerError(f"trigger '{tid}': when.repeat must be a boolean") -def _website_matches(when: dict, app: str, method: str, path: str, status: int) -> bool: - """Whether a request matches a website `when` (`where` aside); without `status`, only a response below 400 does.""" +def _website_matches(when: dict, app: str, method: str, path: str, status: int, op: str | None = None) -> bool: + """Whether a request matches a website `when` (`where` aside); with `op`, only one reaching that route does, and + without `status`, only a response below 400 does.""" if when["type"] != "website" or when["app"] != app or when.get("method", method) != method: return False + if "op" in when and when["op"] != op: + return False if "path" in when and not _eval_predicate(when["path"], path): return False return _eval_predicate(when["status"], status) if "status" in when else status < 400 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 03913adac..319329d45 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -7,6 +7,7 @@ import json import logging import re +import secrets import shlex import time from collections import Counter @@ -23,7 +24,9 @@ from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode from agent_env.env.gateway.constants import ( DEFAULT_MCP_SERVER_NAME, + GATEWAY_ROLE_KEY_ENV, STATE_TOOL_NAME_CLASHES_KEY, + WEBSITE_ROLE_KEY_ENV, WELL_KNOWN_PATH, random_mcp_server_name, tool_name_clashes_text, @@ -117,6 +120,7 @@ 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 + role_key: str | None = None @dataclass @@ -281,6 +285,7 @@ def create_docker_compose( host_ips: tuple[str, ...] = (), extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, + role_key: str | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. @@ -292,6 +297,8 @@ def create_docker_compose( gateway_image: Docker image for the gateway. gateway_port: Port to expose for the gateway. website_configs: Optional list of website configurations (frontend + backend pairs). + role_key: The deploy's role key, given to the gateway as AGENTENV_ROLE_KEY and to each website server as + WEBSITE_ROLE_KEY; without one, both believe the AgentEnv-Role header as before it. Returns: docker-compose.yml content as a string. @@ -489,6 +496,8 @@ 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)}") + if role_key: + gateway_env.append(f" - {GATEWAY_ROLE_KEY_ENV}={role_key}") publish = host_port or (lambda port: port) exposed_gateway_ports = [] if expose_gateway_port: @@ -534,6 +543,7 @@ def create_docker_compose( f" - ENVIRONMENT_NAME={wc.environment_name}", f" - PORT={AGENT_ENV_WEBSITE_BACKEND_PORT}", " - PYTHONUNBUFFERED=1", + *([f" - {WEBSITE_ROLE_KEY_ENV}={role_key}"] if role_key else []), " healthcheck:", f' test: ["CMD-SHELL", "curl -sf -o /dev/null http://localhost:{AGENT_ENV_WEBSITE_BACKEND_PORT}{WELL_KNOWN_PATH}' f' || curl -sf -o /dev/null http://localhost:{AGENT_ENV_WEBSITE_BACKEND_PORT}{wc.health_path}"]', @@ -601,10 +611,12 @@ async def deploy( attribution: Optional[Attribution] = None, artifact_id: str | None = None, artifact_version: int | None = None, + keyed_roles: bool = False, ) -> DeployedGatewayEnv: """A gateway in front of the env, with the state store it acquires (local Postgres by default); returns the record. A MultiEnv's gateway fronts its MCP servers and websites, a WebsiteEnv's its website, and any other env's its one MCP server. + With ``keyed_roles``, the deploy draws a role key (see ``create_gateway``). """ from agent_env.config import get_config, service_db_env_id from agent_env.env.env import Env @@ -623,6 +635,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, + keyed_roles=keyed_roles, ) return DeployedGatewayEnv( env_id=env.id, @@ -639,6 +652,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, + role_key=deployed.role_key, ) async def create_gateway( @@ -661,6 +675,7 @@ async def create_gateway( mcp_server_name: str | None = None, *, attribution: Optional[Attribution] = None, + keyed_roles: bool = False, ) -> DeployedGateway: """Create and start a gateway from explicit topology, for a caller that builds its own record rather than calling ``deploy()``. @@ -674,20 +689,26 @@ 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. + + ``keyed_roles`` draws a secret role key for the deploy, given to the gateway and each website server and kept on + the record: the gateway then believes a role other than ``default`` only beside it, so every caller acting as one + must send it. Off, as before the key, the role header alone is believed. """ 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, keyed_roles=keyed_roles, ) - async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, *, mcp_server_name: str | None = None, **gateway) -> DeployedGateway: + async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, *, mcp_server_name: str | None = None, + keyed_roles: bool = False, **gateway) -> DeployedGateway: """One gateway, deployed on the sandbox provider or each chained member in turn, then its tools probed.""" attempt = functools.partial( self._deploy_gateway, **gateway, mcp_server_name=mcp_server_name or random_mcp_server_name(), # drawn once, so every chained attempt presents the same name + role_key=secrets.token_urlsafe(32) if keyed_roles else None, ) earlier_sandbox = self._sandbox result = await self._run(sandbox_provider, env_id, attempt) @@ -739,6 +760,7 @@ async def _deploy_via_vm( existing_sandbox: Sandbox | None = None, sidecars: list[SidecarConfig] | None = None, mcp_server_name: str | None = None, + role_key: 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 @@ -830,6 +852,7 @@ async def _stand_up(_spec): host_ips=sandbox.host_ips, extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, + role_key=role_key, ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -923,7 +946,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, role_key=role_key, ) async def _deploy_via_containers( @@ -940,6 +963,7 @@ async def _deploy_via_containers( memory_mb: int | None = None, attribution: Optional[Attribution] = None, mcp_server_name: str | None = None, + role_key: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env from agent_env.config import gateway_env_id, get_config @@ -1036,6 +1060,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, + **({GATEWAY_ROLE_KEY_ENV: role_key} if role_key else {}), } gateway_sb = await sandbox_provider.create_container( image_name=gateway_env.docker_image_artifact.image_name, @@ -1069,6 +1094,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, + role_key=role_key, ) except Exception: await self.close() @@ -1370,6 +1396,7 @@ async def _deploy_gateway( sidecars: list[SidecarConfig] | None = None, mcp_server_name: str | None = None, attribution: Optional[Attribution] = None, + role_key: str | None = None, ) -> DeployedGateway: """One gateway deploy on one (unchained) sandbox provider: a VM on one that creates VMs, else a container per server on its private network.""" @@ -1394,7 +1421,7 @@ 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, role_key=role_key, **deploy_kwargs, ) earlier_sandbox = self._sandbox try: @@ -1403,7 +1430,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, role_key=role_key, **deploy_kwargs, ) except BaseException: await self._terminate_new_sandbox(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) @@ -1566,7 +1593,7 @@ def _compose_name_problems(mcp_servers: list[MCPServerConfig], website_configs: # Compose `environment:` keys whose values are secrets and must be masked before the # rendered compose is logged. Matches the key to the left of the first '=' on an # ` - KEY=VALUE` env line. -_SECRET_ENV_KEY_RE = re.compile(r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY)[A-Za-z0-9_]*)=.*$") +_SECRET_ENV_KEY_RE = re.compile(r"^(\s*-\s*)([A-Za-z0-9_]*(?:SECRET|PASSWORD|TOKEN|API_KEY|ROLE_KEY)[A-Za-z0-9_]*)=.*$") def _compose_ps_rows(text: str) -> list[dict]: diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index 2b082cd69..5ee2a6214 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -283,9 +283,11 @@ def _raise_unless_completed(self) -> None: if (error := self._run.exception()) is not None: raise error - async def _serve_in_process(self, agent: str, role: str | None, mcp_urls: dict[str, str]) -> DeployedAgent: + async def _serve_in_process(self, agent: str, role: str | None, mcp_urls: dict[str, str], + role_keys: dict[str, str | None] | None = None) -> DeployedAgent: """A proxy served in this process for ``agent``, which works in no sandbox, with each env server registered as - ``deploy_agent`` registers them with any agent, at the URL this process reaches: played as the others are.""" + ``deploy_agent`` registers them with any agent, at the URL this process reaches, with its deploy's role key + from ``role_keys``: played as the others are.""" server = InProcessProxy() self._in_process.append(server) url = await server.start() @@ -295,7 +297,8 @@ async def _serve_in_process(self, agent: str, role: str | None, mcp_urls: dict[s card = response.json() mcp_config = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) for env_id, mcp_url in mcp_urls.items(): - body = _mcp_add_body(mcp_config, mcp_url, sandbox_request_headers_for_url(mcp_url), env_id, role) + body = _mcp_add_body(mcp_config, mcp_url, sandbox_request_headers_for_url(mcp_url), env_id, role, + (role_keys or {}).get(env_id)) (await client.post(url + mcp_config["params"]["endpoint"], json=body)).raise_for_status() deployed = DeployedAgent(agent_name=agent, api_url=url, a2a_url=url, a2a_card=card, role=role) await self._attach(agent, deployed) diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py index 9bcc85816..7b8c12e51 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -45,11 +45,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: episode = _episode_of(context) agent = self.step.agent_name if not self.sandboxed: - mcp_urls = {} + mcp_urls, role_keys = {}, {} for env_id in getattr(self.step, "env_ids", None) or []: deployed = bound_deployment(context, env_id, env_step_id=self.step.env_step_id, step_id=self.id) mcp_urls[env_id] = deployed.mcp_url if deployed is not None else live_mcp_url(env_id) - deployed = await episode._serve_in_process(agent, getattr(self.step, "role", None), mcp_urls) + role_keys[env_id] = deployed.role_key if deployed is not None else None + deployed = await episode._serve_in_process(agent, getattr(self.step, "role", None), mcp_urls, role_keys) context.deployed_agents.append(deployed) if system_prompt := getattr(self.step, "system_prompt", None): config = A2AAgent.negotiate_agent_config(deployed.a2a_card, {"system_prompt": system_prompt}) diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index cbacdc804..1f37a3847 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -31,7 +31,7 @@ from agent_env.a2a_agent.staging import StagedObjectStore, transfer_store from agent_env.config import get_config from agent_env.env.env import DeployedEnv, DeployedSandboxEnv, Env -from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_KEY_HEADER from agent_env.providers.sandbox_providers.local_docker.provider import host_url_for from agent_env.a2a_agent.a2a_agent import DEFAULT_A2A_PORT, unwritten_local_agent from agent_env.providers.sandbox_providers.sandbox_provider import ( @@ -70,9 +70,11 @@ def _mcp_add_request(mcp_ext: dict) -> dict: return request if isinstance(request, dict) else {} -def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str | None, role: str | None = None) -> dict: +def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str | None, role: str | None = None, + role_key: str | None = None) -> dict: """Body for the agent's mcp-config `add`. The card name and the role header ride along only when the card's `add` - takes them, judged on the exact field set sent; what an env does with the role is the env's business.""" + takes them, judged on the exact field set sent, the deploy's `role_key` beside the role so its gateway believes + it; what an env does with the role is the env's business.""" request = _mcp_add_request(mcp_ext) body: dict = {"url": url} if headers: @@ -81,6 +83,8 @@ def _mcp_add_body(mcp_ext: dict, url: str, headers: dict | None, card_name: str body["name"] = card_name if role and card_request_accepts(request, {*body, "headers"}): body.setdefault("headers", {})[AGENT_ENV_ROLE_HEADER] = role + if role_key: + body["headers"][AGENT_ENV_ROLE_KEY_HEADER] = role_key return body @@ -433,7 +437,8 @@ async def _execute(self, context: TaskStepContext, deployed_agents: list) -> Tas url = host_url_for(live_mcp_url(env_id), deployed.sandbox_type) card_name = None if mcp_ext: - body = _mcp_add_body(mcp_ext, url, sandbox_request_headers_for_url(url), card_name, self.role) + body = _mcp_add_body(mcp_ext, url, sandbox_request_headers_for_url(url), card_name, self.role, + deployed_env.role_key if deployed_env is not None else None) registrations.append((env_id, body)) identity_payload: dict[str, str] = {"name": self.agent_name} diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 375a4c515..55d605b0a 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -48,6 +48,7 @@ def __init__( artifact_id: Optional[str] = None, artifact_version: Optional[int] = None, env_params: Optional[dict] = None, + keyed_roles: bool = False, ): super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) self.env_id = env_id @@ -66,6 +67,7 @@ def __init__( self.metadata = metadata or {} # The env type's own deploy options, which this step passes to its deploy() without reading them. self.env_params = env_params + self.keyed_roles = keyed_roles def run_state(self, user_overrides: Mapping[str, Any]) -> tuple[str | None, str | None]: """The env state type and attached state instance a run with ``user_overrides`` deploys on: each the run's, @@ -89,6 +91,7 @@ def to_dict(self) -> dict: base["env_state_instance_id"] = self.env_state_instance_id base["metadata"] = self.metadata base["env_params"] = self.env_params + base["keyed_roles"] = self.keyed_roles return base @classmethod @@ -109,6 +112,7 @@ def from_dict(cls, data: dict) -> DeployEnvTaskStep: env_state_instance_id=data.get("env_state_instance_id"), metadata=dict(data.get("metadata") or {}), env_params=data.get("env_params"), + keyed_roles=bool(data.get("keyed_roles", False)), ) def preflight(self) -> list[str]: @@ -133,7 +137,8 @@ def preflight(self) -> list[str]: return [f"deploy_env '{self.id}': {refusal}"] if refusal else [] def _env_params_option(self) -> dict: - return {"env_params": self.env_params} if self.env_params else {} + return {**({"env_params": self.env_params} if self.env_params else {}), + **({"keyed_roles": True} if self.keyed_roles else {})} async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.env.env import Env @@ -188,6 +193,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if "env_params" not in params and not accepts_kwargs: raise ValueError(f"deploy_env '{self.id}': {env.type} envs take no env_params") deploy_kwargs["env_params"] = self.env_params + if self.keyed_roles: + if "keyed_roles" not in params and not accepts_kwargs: + raise ValueError(f"deploy_env '{self.id}': {env.type} envs take no keyed_roles") + deploy_kwargs["keyed_roles"] = True deployed_env = await env.deploy(**deploy_kwargs) # create_instance runs inside env.deploy(), so an annotation set after it # returns has to be persisted explicitly to reach the instance record. diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index fdecb9e75..c79d6153c 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -645,7 +645,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if not (agent.on_host and agent.sandbox_type == LocalSandbox.type): gateway_url = reachable_url(gateway_url, from_sandbox_type=deployed.sandbox_type, to_sandbox_type=agent.sandbox_type) - install_path = await A2AAgent.install_cli(deployed_agent, artifact, gateway_url, **_agent_files(agent)) + install_path = await A2AAgent.install_cli(deployed_agent, artifact, gateway_url, role_key=deployed.role_key, + **_agent_files(agent)) agent_clis = context.metadata.setdefault("installed_clis", {}).setdefault(self.agent_name, {}) agent_clis[artifact.id] = { "install_path": install_path, diff --git a/src/agent_env/task_step/task_steps/mcp_cli_builder/codegen.py b/src/agent_env/task_step/task_steps/mcp_cli_builder/codegen.py index 4be0384b0..9853400f2 100644 --- a/src/agent_env/task_step/task_steps/mcp_cli_builder/codegen.py +++ b/src/agent_env/task_step/task_steps/mcp_cli_builder/codegen.py @@ -14,10 +14,11 @@ from agentenv_protocol import MANIFEST_VERSION, manifest_compatible -from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_KEY_HEADER GATEWAY_URL_ENV_VAR = "AGENT_ENV_GATEWAY_URL" ROLE_ENV_VAR = "AGENT_ENV_ROLE" +ROLE_KEY_ENV_VAR = "AGENT_ENV_ROLE_KEY" DEFAULT_CLI_ROLE = "cli" _COMMAND_NAME_RE = re.compile(r"^[a-zA-Z][a-zA-Z0-9_-]*$") _WORD_RE = re.compile(r"[A-Z]+(?=[A-Z][a-z]|$)|[A-Z]?[a-z]+|[0-9]+") @@ -166,6 +167,8 @@ def generate_cli_script( role_env_var=ROLE_ENV_VAR, default_role=DEFAULT_CLI_ROLE, role_header=AGENT_ENV_ROLE_HEADER, + role_key_env_var=ROLE_KEY_ENV_VAR, + role_key_header=AGENT_ENV_ROLE_KEY_HEADER, ) @@ -186,6 +189,8 @@ def generate_cli_script( ROLE_ENV = "{role_env_var}" DEFAULT_ROLE = "{default_role}" ROLE_HEADER = "{role_header}" +ROLE_KEY_ENV = "{role_key_env_var}" +ROLE_KEY_HEADER = "{role_key_header}" SERVICE_NAME = {environment_name} INTERFACE_MANIFEST = {interface_manifest_literal} _CONFIG_FILE_NAME = ".env" @@ -226,11 +231,17 @@ def _role(): return os.environ.get(ROLE_ENV, DEFAULT_ROLE) +def _role_headers(): + """The role header, with the deploy's role key beside it when one is set, so the gateway believes the role.""" + key = os.environ.get(ROLE_KEY_ENV) + return {{ROLE_HEADER: _role(), **({{ROLE_KEY_HEADER: key}} if key else {{}})}} + + def _post_step(payload): response = httpx.post( f"{{_gateway_url()}}/step", json=payload, - headers={{ROLE_HEADER: _role()}}, + headers=_role_headers(), timeout=120, ) response.raise_for_status() @@ -248,7 +259,7 @@ def _list_tools(): allowed_names = {{t.get("name") for t in (role_filtered.get("tools") or []) if t.get("name")}} state_resp = httpx.get( f"{{_gateway_url()}}/state", - headers={{ROLE_HEADER: _role()}}, + headers=_role_headers(), timeout=120, ) state_resp.raise_for_status() diff --git a/src/agent_env/task_step/task_steps/post_turn_hooks.py b/src/agent_env/task_step/task_steps/post_turn_hooks.py index 0afb54b71..cca67b1fd 100644 --- a/src/agent_env/task_step/task_steps/post_turn_hooks.py +++ b/src/agent_env/task_step/task_steps/post_turn_hooks.py @@ -22,7 +22,6 @@ from agent_env.artifact import FileArtifact from agent_env.env.env import DeployedEnv, EnvCapabilityUnsupported from agent_env.env.gateway.constants import ( - AGENT_ENV_ROLE_HEADER, EXT_CLOCK_URI, EXT_NOTIFICATIONS_URI, EXT_STEP_URI, @@ -152,7 +151,7 @@ async def call_tool(self, tool: str, args: Optional[dict], role: str) -> dict: "source": STEP_SOURCE_POST_TURN_HOOK, "hook_artifact_id": self._hook_artifact_id} async with httpx.AsyncClient() as client: resp = await client.request(op["method"], f"{self._deployed.environment_url}{op['endpoint']}", json=body, - headers={AGENT_ENV_ROLE_HEADER: role, **self._deployed.control_headers}, + headers={**self._deployed.role_headers(role), **self._deployed.control_headers}, timeout=_TOOL_CALL_TIMEOUT_S) if resp.status_code >= 400: raise RuntimeError(f"env '{self.env_id}': /step call_tool {tool!r} failed (HTTP {resp.status_code}): " diff --git a/src/agent_env/task_step/task_steps/register_env_triggers.py b/src/agent_env/task_step/task_steps/register_env_triggers.py index 156d17fb9..237f77259 100644 --- a/src/agent_env/task_step/task_steps/register_env_triggers.py +++ b/src/agent_env/task_step/task_steps/register_env_triggers.py @@ -21,6 +21,7 @@ TRIGGER_PROVOKING_ROLES, TRIGGER_UNLESS, TRIGGER_WEBSITE, + TRIGGER_WEBSITE_OP, ) from agent_env.task_step.context import TaskStepContext from agent_env.entity.refs import EntityRef @@ -146,12 +147,15 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: f"redeploy it on a gateway image that does.") # An older gateway ignores what it does not declare: named actions would go to the default executor, a # trigger's when.roles would leave it to the watched roles, delays and hours would not wait, a deadline - # would fire even after its `unless` call, and hours in a time zone would be read as UTC. + # would fire even after its `unless` call, hours in a time zone would be read as UTC, and a website trigger + # naming an `op` would fire on every request to the site. used = {TRIGGER_NAMED_EXECUTORS: bool(named), TRIGGER_PROVOKING_ROLES: _when_has(self.triggers, "roles"), TRIGGER_ACTION_DELAY: _delays(self.triggers), TRIGGER_EXECUTOR_HOURS: bool(self.executor_hours), TRIGGER_UNLESS: _when_has(self.triggers, "unless"), TRIGGER_WEBSITE: any(isinstance(t.get("when"), dict) and t["when"].get("type") == "website" for t in self.triggers), + TRIGGER_WEBSITE_OP: any(isinstance(t.get("when"), dict) and t["when"].get("type") == "website" + and t["when"].get("op") is not None for t in self.triggers), TRIGGER_EXECUTOR_HOURS_TIME_ZONE: _hours_in_a_time_zone(self.executor_hours), TRIGGER_NOTIFY: any(a.get("type") == "notify" for t in self.triggers for a in t.get("actions") or [] if isinstance(a, dict))} diff --git a/tst/integration/env/gateway/gateway_test.py b/tst/integration/env/gateway/gateway_test.py index ab4057715..eefe02c96 100644 --- a/tst/integration/env/gateway/gateway_test.py +++ b/tst/integration/env/gateway/gateway_test.py @@ -780,7 +780,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) cli_list = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "cli"}, + headers=result.role_headers("cli"), timeout=30, ) cli_list.raise_for_status() @@ -797,7 +797,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) cli_call = await client.post( f"{result.gateway_url}/step", json={"action": "call_tool", "tool_name": "list_emails", "arguments": {"folder_name": "INBOX"}}, - headers={AGENT_ENV_ROLE_HEADER: "cli"}, + headers=result.role_headers("cli"), timeout=30, ) assert cli_call.status_code == 200, f"cli call failed: {cli_call.status_code} {cli_call.text}" @@ -814,7 +814,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) async with streamable_http_client( result.mcp_url, - http_client=httpx.AsyncClient(headers={AGENT_ENV_ROLE_HEADER: "cli"}), + http_client=httpx.AsyncClient(headers=result.role_headers("cli")), ) as (read_stream, write_stream, _): async with ClientSession(read_stream, write_stream) as session: await session.initialize() @@ -826,7 +826,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) assert "list_emails" not in [t.name for t in await default_session.list_tools()] hidden = await default_session.call_tool("list_emails", {"folder_name": "INBOX"}) assert hidden.isError and "is disabled for role 'default'" in hidden.text, hidden - async with ToolSession(mcp_url=result.mcp_url, role="cli") as cli_session: + async with ToolSession(mcp_url=result.mcp_url, headers=result.role_headers("cli")) as cli_session: assert "list_emails" in [t.name for t in await cli_session.list_tools()] assert "alex.chen@techcorp.com" in (await cli_session.call_tool("list_emails", {"folder_name": "INBOX"})).text @@ -869,7 +869,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) assert state_resp.json()["roles"]["*"] == {"disabled": ["list_emails"], "allowed": []} for role in (None, "cli", "marketing"): - headers = {AGENT_ENV_ROLE_HEADER: role} if role else {} + headers = result.role_headers(role) if role else {} lst = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, headers=headers, timeout=30, ) @@ -879,14 +879,14 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) blocked = await client.post( f"{result.gateway_url}/step", json={"action": "call_tool", "tool_name": "list_emails", "arguments": {"folder_name": "INBOX"}}, - headers={AGENT_ENV_ROLE_HEADER: "marketing"}, + headers=result.role_headers("marketing"), timeout=30, ) assert blocked.status_code == 403, f"expected 403 for never-seen role under global disable, got {blocked.status_code} {blocked.text}" async with streamable_http_client( result.mcp_url, - http_client=httpx.AsyncClient(headers={AGENT_ENV_ROLE_HEADER: "cli"}), + http_client=httpx.AsyncClient(headers=result.role_headers("cli")), ) as (read_stream, write_stream, _): async with ClientSession(read_stream, write_stream) as session: await session.initialize() @@ -926,7 +926,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) mkt_lst = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "marketing"}, + headers=result.role_headers("marketing"), timeout=30, ) mkt_lst.raise_for_status() @@ -944,7 +944,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) cli_lst = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "cli"}, + headers=result.role_headers("cli"), timeout=30, ) cli_lst.raise_for_status() @@ -955,7 +955,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) mkt_lst = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "marketing"}, + headers=result.role_headers("marketing"), timeout=30, ) mkt_lst.raise_for_status() @@ -970,7 +970,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) cli_lst_after_reset = await client.post( f"{result.gateway_url}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "cli"}, + headers=result.role_headers("cli"), timeout=30, ) cli_lst_after_reset.raise_for_status() @@ -1002,7 +1002,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) blocked = await client.post( f"{result.gateway_url}/step", json={"action": "call_tool", "tool_name": "list_emails", "arguments": {"folder_name": "INBOX"}}, - headers={AGENT_ENV_ROLE_HEADER: "step-role"}, + headers=result.role_headers("step-role"), timeout=30, ) assert blocked.status_code == 403, f"expected 403 after step-driven disable, got {blocked.status_code} {blocked.text}" @@ -1020,7 +1020,7 @@ async def assert_generated_data_via_step(client, gateway_url, tools, tool_names) unblocked = await client.post( f"{result.gateway_url}/step", json={"action": "call_tool", "tool_name": "list_emails", "arguments": {"folder_name": "INBOX"}}, - headers={AGENT_ENV_ROLE_HEADER: "step-role"}, + headers=result.role_headers("step-role"), timeout=30, ) assert unblocked.status_code == 200, f"expected 200 after step-driven enable, got {unblocked.status_code} {unblocked.text}" @@ -1816,7 +1816,7 @@ async def test_gateway_changes_name_the_role_behind_each_write(mcp_server_envs, seen = (await client.get(f"{gw}/changes", params={"limit": 10000}, timeout=30)).json()["changes"] after_id = seen[-1]["id"] if seen else 0 for role, subject in (("default", "first"), ("alice@example.com", "second")): - resp = await client.post(f"{gw}/step", headers={AGENT_ENV_ROLE_HEADER: role}, timeout=60, json={ + resp = await client.post(f"{gw}/step", headers=result.role_headers(role), timeout=60, json={ "action": "call_tool", "tool_name": "send_email", "arguments": {"recipients": ["dana@example.com"], "subject": subject, "content": "body"}}) assert resp.status_code == 200, resp.text @@ -1850,7 +1850,7 @@ async def test_gateway_changes_name_the_request_behind_a_write_through_the_rest_ async with httpx.AsyncClient() as client: seen = (await client.get(f"{gw}/changes", params={"limit": 10000}, timeout=30)).json()["changes"] after_id = seen[-1]["id"] if seen else 0 - async with httpx.AsyncClient(headers={AGENT_ENV_ROLE_HEADER: "carol@example.com"}, timeout=60) as as_carol: + async with httpx.AsyncClient(headers=result.role_headers("carol@example.com"), timeout=60) as as_carol: async with streamable_http_client(f"{gw}/svc/mcp-email/mcp", http_client=as_carol) as (read, write, _): async with ClientSession(read, write) as session: await session.initialize() @@ -2170,7 +2170,7 @@ async def assert_tool_sees_items(session, tools_result, tool_names): await _verify_list_tools(mcp_url=result.mcp_url, tool_verifier_callback=assert_tool_sees_items) # The role a ToolSession sends reaches the child's tool as its caller, with the gateway's key for each session - async with (ToolSession(mcp_url=result.mcp_url, role="alice@example.com") as alice, + async with (ToolSession(mcp_url=result.mcp_url, headers=result.role_headers("alice@example.com")) as alice, ToolSession(mcp_url=result.mcp_url) as anonymous): seen = [json.loads((await session.call_tool("items_whoami")).text) for session in (alice, alice, anonymous)] assert [caller["role"] for caller in seen] == ["alice@example.com", "alice@example.com", None] @@ -2548,7 +2548,7 @@ async def _poll(client, gw, want_fired, timeout_s=120): hr = await client.post(f"{gw}/step", json={"action": "call_tool", "tool_name": "search_emails", "arguments": {"query": "alex.chen", "folder_name": "INBOX"}}, - headers={AGENT_ENV_ROLE_HEADER: "harness"}, timeout=60) + headers=result.role_headers("harness"), timeout=60) assert hr.status_code == 200, f"harness /step failed: {hr.text}" await asyncio.sleep(4) assert {t["id"]: t["status"] for t in (await _state(client, gw))["triggers"]}["decoy"] == "armed", ( @@ -3064,7 +3064,7 @@ async def test_gateway_triggers_nl_executor(sandbox_provider, mcp_server_envs, e # The one-tool constraint really held (determinism lever, verified live). lt = await client.post(f"{gw}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "executor"}, timeout=30) + headers=result.role_headers("executor"), timeout=30) exec_tools = [t["name"] for t in lt.json()["tools"]] assert exec_tools == ["send_email"], f"executor tool surface not locked to one tool: {exec_tools}" @@ -3105,7 +3105,7 @@ async def _state(): se = await client.post(f"{gw}/step", json={"action": "call_tool", "tool_name": "search_emails", "arguments": {"query": marker, "folder_name": "SENT"}}, - headers={AGENT_ENV_ROLE_HEADER: "harness"}, timeout=30) + headers=result.role_headers("harness"), timeout=30) assert marker in se.text, f"marker email not found in SENT: {se.text[:400]}" # Cascade-safety: nlact was detected exactly once (from the provoke), NOT from the @@ -3351,17 +3351,17 @@ async def _rate_check(client, gw, rate, sleep_s=5.0): await client.post(f"{gw}/tools/{TOOL_DISABLE_ACTION}", json={"role": "locked", "tools": ["*"]}, timeout=30) locked = await client.post(f"{gw}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "locked"}, timeout=30) + headers=result.role_headers("locked"), timeout=30) assert GET_TIME not in [t["name"] for t in locked.json()["tools"]], "wildcard deny did not hide get_time" denied = await client.post(f"{gw}/step", json={"action": "call_tool", "tool_name": GET_TIME, "arguments": {}}, - headers={AGENT_ENV_ROLE_HEADER: "locked"}, timeout=30) + headers=result.role_headers("locked"), timeout=30) assert denied.status_code == 403, denied.text assert GET_TIME in (await _get_time_surfaces(client, gw, result.mcp_url)).step, "default role lost it" await client.post(f"{gw}/tools/{TOOL_ENABLE_ACTION}", json={"role": "locked", "tools": [GET_TIME]}, timeout=30) relocked = await client.post(f"{gw}/step", json={"action": "list_tools"}, - headers={AGENT_ENV_ROLE_HEADER: "locked"}, timeout=30) + headers=result.role_headers("locked"), timeout=30) assert GET_TIME in [t["name"] for t in relocked.json()["tools"]], "re-enable did not restore it" finally: await env.close() diff --git a/tst/integration/task_step/test_next_env_event_hook_local.py b/tst/integration/task_step/test_next_env_event_hook_local.py index de63dc0d7..103127090 100644 --- a/tst/integration/task_step/test_next_env_event_hook_local.py +++ b/tst/integration/task_step/test_next_env_event_hook_local.py @@ -12,7 +12,6 @@ import agent_env from agent_env.a2a_agent import conversation_store from agent_env.artifact import FileArtifact -from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep @@ -65,7 +64,7 @@ async def _run(label: str, triggers: list[dict], params: dict) -> tuple: state = (await client.get("/triggers/state")).raise_for_status().json() clock = (await client.get("/clock/state")).raise_for_status().json() stored = await client.post("/step", json={"action": "call_tool", "tool_name": "list_items", "arguments": {}}, - headers={AGENT_ENV_ROLE_HEADER: "harness"}) + headers=deployed.role_headers("harness")) finally: await env.close() items = json.loads(stored.json()["content"][0]["text"])["items"] 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..7c6964cca 100644 --- a/tst/integration/task_step/test_register_env_triggers_local.py +++ b/tst/integration/task_step/test_register_env_triggers_local.py @@ -17,9 +17,9 @@ import pytest import pytest_asyncio +from agent_env.env.env import DeployedEnv from agent_env.env.envs.mcp_server import MCPServerEnv from agent_env.env.envs.multi_env import MultiEnv -from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER from agent_env.env.gateway.constants import METHOD_RESET, RPC_PATH, TRIGGER_IN_FLIGHT_STATUSES from agent_env.task_step.context import DeployedAgent, TaskStepContext from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep @@ -66,8 +66,8 @@ async def test_each_nl_action_reaches_the_executor_it_names(items_env): {"id": "on-add", "when": {"type": "action", "tool": "items_add_item", "repeat": True}, "actions": []}, ] try: - alice = _Executor("alice", "alice@example.com", deployed.gateway_url, host, started) - bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started) + alice = _Executor("alice", "alice@example.com", deployed, host, started) + bob = _Executor("bob", "bob@example.com", deployed, host, started) 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) @@ -76,7 +76,7 @@ async def test_each_nl_action_reaches_the_executor_it_names(items_env): assert provoke.status_code == 200, provoke.text state = await _quiet(client, {"to-alice": 1, "to-bob": 1, "on-add": 0}) stored = await client.post("/step", json={"action": "call_tool", "tool_name": "list_items", "arguments": {}}, - headers={AGENT_ENV_ROLE_HEADER: "harness"}) + headers=deployed.role_headers("harness")) finally: for executor in started: executor.stop() @@ -108,8 +108,8 @@ async def test_personas_react_to_each_other_until_one_stays_quiet(items_env): "actions": [{"type": "nl", "executor": "denise", "instruction": "DENISE-TO-BOB"}]}, ] try: - denise = _Executor("denise", "denise@example.com", deployed.gateway_url, host, started) - bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started, replies=1) + denise = _Executor("denise", "denise@example.com", deployed, host, started) + bob = _Executor("bob", "bob@example.com", deployed, host, started, replies=1) 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) @@ -118,7 +118,7 @@ async def test_personas_react_to_each_other_until_one_stays_quiet(items_env): assert provoke.status_code == 200, provoke.text state = await _quiet(client, {"denise-hears-solver": 1, "bob-hears-denise": 2, "denise-hears-bob": 1}) stored = await client.post("/step", json={"action": "call_tool", "tool_name": "list_items", "arguments": {}}, - headers={AGENT_ENV_ROLE_HEADER: "harness"}) + headers=deployed.role_headers("harness")) finally: for executor in started: executor.stop() @@ -140,7 +140,7 @@ async def test_a_reply_waits_its_delay_and_hours_and_its_writes_name_the_fire(it triggers = [{"id": "bob-hears-solver", "when": {"type": "action", "tool": "list_items"}, "actions": [{"type": "nl", "executor": "bob", "instruction": "BOB-ACTS", "delay": "PT2M"}]}] try: - bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started) + bob = _Executor("bob", "bob@example.com", deployed, 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: # The env is shared, so its trajectory holds the earlier tests' calls: read this test's, after them. @@ -185,7 +185,7 @@ async def test_a_deadline_fires_on_silence_and_a_reply_cancels_it(items_env): "where": {"item": {"regex": f"^REPLY-{n}"}}}}, "actions": [{"type": "nl", "executor": "bob", "instruction": f"ESCALATE-{n}"}]}] try: - bob = _Executor("bob", "bob@example.com", deployed.gateway_url, host, started) + bob = _Executor("bob", "bob@example.com", deployed, 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: await _arm_clock(client, "2026-10-14T10:00:00Z", 60) @@ -239,9 +239,10 @@ class _Executor: """An A2A executor on this machine that acts on each instruction through the gateway, under its own role; with ``replies``, it stops acting after that many.""" - def __init__(self, name: str, role: str, gateway_url: str, host: str, started: list, + def __init__(self, name: str, role: str, deployed: DeployedEnv, host: str, started: list, replies: int | None = None) -> None: - self.role, self.gateway_url, self.replies = role, gateway_url, replies + self.role, self.gateway_url, self.replies = role, deployed.gateway_url, replies + self.headers = deployed.role_headers(role) self.instructions: list[str] = [] self.acted_as: list[str] = [] self._server = ThreadingHTTPServer((host, 0), functools.partial(_ExecutorHandler, self)) @@ -265,7 +266,7 @@ def stop(self) -> None: def _call(self, tool: str, arguments: dict) -> dict: response = httpx.post(f"{self.gateway_url}/step", json={"action": "call_tool", "tool_name": tool, "arguments": arguments}, - headers={AGENT_ENV_ROLE_HEADER: self.role}, timeout=60) + headers=self.headers, timeout=60) response.raise_for_status() return response.json() diff --git a/tst/unit/env/envs/test_mcp_server_plugin_provider.py b/tst/unit/env/envs/test_mcp_server_plugin_provider.py index b37528070..a1d482d0d 100644 --- a/tst/unit/env/envs/test_mcp_server_plugin_provider.py +++ b/tst/unit/env/envs/test_mcp_server_plugin_provider.py @@ -230,7 +230,7 @@ async def test_a_multi_env_child_of_a_plugin_type_is_still_the_gateways(): multi = MultiEnv(id="multi", version=1, mcp_server_envs=[child]) record = SimpleNamespace(env_id="multi", env_version=1, sandbox_id="gw", sandbox_type="modal", sandbox_ids={}, env_state_instance_ids=[], instance_id="i", mcp_server_name="m", environment_card=None, - environment_card_url=None, gateway_url="https://gw.example") + environment_card_url=None, gateway_url="https://gw.example", role_key=None) gw = MagicMock(sandbox_id="gw", mode="vm", type="modal") with patch("agent_env.env.env.Env.get", return_value=multi), \ patch("agent_env.providers.build_sandbox_provider", return_value=MagicMock(get_sandbox=AsyncMock(return_value=gw))), \ diff --git a/tst/unit/env/envs/test_multi_env_name.py b/tst/unit/env/envs/test_multi_env_name.py index 0055bef90..70a9ea224 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", role_key=None) async def deploy_gateway(self, sandbox_provider, **kwargs): self._sandbox = SimpleNamespace(sandbox_id="gw", type="modal_vm") diff --git a/tst/unit/env/gateway/rest_proxy_stream_test.py b/tst/unit/env/gateway/rest_proxy_stream_test.py index f2378150f..79e8b5fb6 100644 --- a/tst/unit/env/gateway/rest_proxy_stream_test.py +++ b/tst/unit/env/gateway/rest_proxy_stream_test.py @@ -13,10 +13,13 @@ import httpx import pytest +from starlette.testclient import TestClient +from agentenv_protocol.website import ROLE_KEY_ENV, ApiContext, WebsiteEnvironment, api, website 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 import rest_proxy as rest_proxy_module +from agent_env.env.gateway.caller import Caller from agent_env.env.gateway.gateway import Gateway UPSTREAM = "http://upstream" @@ -479,3 +482,212 @@ async def test_bodies_that_are_not_a_child_mcp_request_are_forwarded_untouched(m 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] + + +ROLE_KEY = "deploy-secret" + + +def _role_and_key(request: httpx.Request) -> tuple[list[str], list[str]]: + return request.headers.get_list("agentenv-role"), request.headers.get_list("agentenv-role-key") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/svc/mcp-svc/x", "/website/web/x"]) +async def test_with_a_role_key_a_role_is_believed_only_beside_it_and_the_key_goes_upstream_only_with_it(monkeypatch, path): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}, role_key=ROLE_KEY) + for headers in ({"AgentEnv-Role": "admin"}, {"AgentEnv-Role": "admin", "AgentEnv-Role-Key": "guess"}, + {"AgentEnv-Role": "admin", "AgentEnv-Role-Key": ROLE_KEY}, {"AgentEnv-Role-Key": ROLE_KEY}): + status, _, _ = await _call(gw, path, headers=headers) + assert status == 200 + assert [_role_and_key(r) for r in up.forwarded] == [ + (["default"], []), (["default"], []), (["admin"], [ROLE_KEY]), (["default"], [])] + + +@pytest.mark.asyncio +async def test_without_a_role_key_the_role_header_alone_is_believed_and_a_client_s_key_never_goes_upstream(monkeypatch): + up = _Upstream([b"{}"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/x", headers={"AgentEnv-Role": "admin"}) + await _call(gw, "/svc/mcp-svc/x", headers={"AgentEnv-Role": "admin", "AgentEnv-Role-Key": "forged"}) + assert [_role_and_key(r) for r in up.forwarded] == [(["admin"], []), (["admin"], [])] + + +@pytest.mark.asyncio +async def test_with_a_role_key_step_attributes_a_role_sent_without_it_to_default(monkeypatch): + gw = _gateway(role_key=ROLE_KEY) + seen: list[str] = [] + monkeypatch.setattr(gw, "_step_list_tools", lambda role: (seen.append(role), gw._json_response({"tools": []}, 200))[1]) + try: + for headers in ({"AgentEnv-Role": "admin"}, {"AgentEnv-Role": "admin", "AgentEnv-Role-Key": ROLE_KEY}): + status, _, _ = await _call(gw, "/step", method="POST", body=b'{"action": "list_tools"}', headers=headers) + assert status == 200 + finally: + await gw._close_child_sessions() + assert seen == ["default", "admin"] + + +class _CallerRecorder: + """Stands in for FastMCP's app under the gateway's ASGI wrapper: records the caller bound for each request.""" + + def __init__(self) -> None: + self.router = type("Router", (), {"lifespan_context": None})() + self.roles: list[str] = [] + + async def __call__(self, scope, receive, send) -> None: + self.roles.append(Caller.current().role) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("role_key, expected", [ + (ROLE_KEY, ["default", "default", "admin", "default"]), + (None, ["admin", "admin", "admin", "default"]), +], ids=["keyed", "legacy"]) +async def test_the_mcp_path_binds_the_role_the_key_vouches_for(monkeypatch, role_key, expected): + gw = _gateway(role_key=role_key) + recorder = _CallerRecorder() + monkeypatch.setattr(gw._mcp, "streamable_http_app", lambda: recorder) + app = gw._asgi_app() + for headers in ([(b"agentenv-role", b"admin")], [(b"agentenv-role", b"admin"), (b"agentenv-role-key", b"guess")], + [(b"agentenv-role-key", ROLE_KEY.encode()), (b"agentenv-role", b"admin")], []): + await app({"type": "http", "method": "POST", "path": "/mcp", "headers": headers}, None, None) + assert recorder.roles == expected + + +def _forwarded(request: httpx.Request) -> tuple[list[str], list[str], list[str]]: + return tuple(request.headers.get_list(name) for name in ("x-forwarded-host", "x-forwarded-proto", "x-forwarded-prefix")) + + +@pytest.mark.asyncio +async def test_the_upstream_learns_the_host_scheme_and_route_prefix_the_client_reached_it_at(monkeypatch): + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/api/x/", headers={"host": "gw.example", "x-forwarded-prefix": "/elsewhere"}) + await _call(gw, "/website/web", headers={"host": "gw.example"}) + await _call(gw, "/website/web/", headers={"host": "gw.example"}) + await _call(gw, "/svc/mcp-svc/x%2Fy", headers={"host": "gw.example"}) + await _call(gw, "/gw%20x/svc/mcp-svc/x", headers={"host": "gw.example"}, root_path="/gw x") + assert [_forwarded(r) for r in up.forwarded] == [ + (["gw.example"], ["http"], ["/website/web"]), + (["gw.example"], ["http"], ["/website/web"]), + (["gw.example"], ["http"], ["/website/web"]), + (["gw.example"], ["http"], ["/svc/mcp-svc"]), + (["gw.example"], ["http"], ["/gw%20x/svc/mcp-svc"]), + ] + + +@pytest.mark.asyncio +async def test_a_front_s_own_forwarded_host_and_scheme_win_and_reach_the_upstream_once(monkeypatch): + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/x", raw_headers=[ + (b"host", b"gateway.internal"), (b"x-forwarded-host", b"env.example, proxy"), (b"x-forwarded-proto", b"https"), + (b"x-forwarded-proto", b"http"), (b"x-forwarded-for", b"10.0.0.1")]) + await _call(gw, "/website/web/x", headers={"host": "gw.example", "x-forwarded-proto": "gopher"}) + assert [_forwarded(r) for r in up.forwarded] == [ + (["env.example"], ["https"], ["/website/web"]), + (["gw.example"], ["http"], ["/website/web"]), + ] + assert up.forwarded[0].headers.get_list("x-forwarded-for") == ["10.0.0.1"] + + +@pytest.mark.asyncio +async def test_an_sdk_site_behind_the_gateway_believes_only_a_keyed_role_and_names_the_gateway_s_address(monkeypatch): + @website(domains=["notes.example"]) + class Notes(WebsiteEnvironment): + @api.get("/api/whoami") + def whoami(self, ctx: ApiContext) -> dict: + return {"role": ctx.role, "base_url": ctx.base_url} + + monkeypatch.setenv(ROLE_KEY_ENV, ROLE_KEY) + transport = httpx.ASGITransport(app=Notes().create_app()) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + gw = _gateway(website_urls={"notes": "http://notes"}, role_key=ROLE_KEY) + + async def whoami(path: str, **headers) -> tuple[int, dict, bytes]: + status, answered, chunks = await _call(gw, path, headers={"host": "gw.example", **headers}) + return status, answered, b"".join(chunks) + + _, _, spoofed = await whoami("/website/notes/api/whoami", **{"AgentEnv-Role": "admin"}) + _, _, keyed = await whoami("/website/notes/api/whoami", **{"AgentEnv-Role": "admin", "AgentEnv-Role-Key": ROLE_KEY}) + status, redirect, _ = await whoami("/website/notes/api/whoami/") + assert json.loads(spoofed) == {"role": "default", "base_url": "http://gw.example/website/notes"} + assert json.loads(keyed)["role"] == "admin" + assert (status, redirect["location"]) == (307, "http://gw.example/website/notes/api/whoami") + + +def test_the_route_an_sdk_card_names_reaches_its_nested_and_typed_paths(): + @website(domains=["files.example"]) + class Files(WebsiteEnvironment): + @api.get("/files/{rest:path}") + def files(self, rest: str) -> dict: + return {"rest": rest} + + @api.get("/api/notes/{note_id:int}") + def get_note(self, note_id: int) -> dict: + return {"id": note_id} + + card = TestClient(Files().create_app()).get("/.well-known/agent-env.json").json() + routes = rest_proxy_module.website_routes(card) + assert rest_proxy_module.route_name(routes, "GET", "/files/a/b.txt") == "files" + assert rest_proxy_module.route_name(routes, "GET", "/api/notes/7") == "get_note" + assert rest_proxy_module.route_name(routes, "GET", "/api/notes/seven") is None + + +def test_a_request_a_built_in_route_answers_is_named_by_no_card_route(): + @website(domains=["any.example"]) + class AnyApi(WebsiteEnvironment): + @api.get("/api/{rest:path}") + def any_api(self, rest: str) -> dict: + return {"rest": rest} + + @api.delete("/{rest:path}") + def any_delete(self, rest: str) -> dict: + return {"rest": rest} + + card = TestClient(AnyApi().create_app()).get("/.well-known/agent-env.json").json() + routes = rest_proxy_module.website_routes(card) + assert rest_proxy_module.route_name(routes, "GET", "/api/session") is None + assert rest_proxy_module.route_name(routes, "DELETE", "/api/session") is None + assert rest_proxy_module.route_name(routes, "GET", "/api/health") is None + assert rest_proxy_module.route_name(routes, "GET", "/api/notes") == "any_api" + assert rest_proxy_module.route_name(routes, "DELETE", "/api/health") == "any_delete" + + +def test_a_site_route_at_the_session_path_replaces_the_built_in_one(): + @website(domains=["own.example"]) + class OwnSession(WebsiteEnvironment): + @api.post("/api/session") + def sign_in(self) -> dict: + return {} + + @api.get("/api/{rest:path}") + def any_api(self, rest: str) -> dict: + return {"rest": rest} + + card = TestClient(OwnSession().create_app()).get("/.well-known/agent-env.json").json() + routes = rest_proxy_module.website_routes(card) + assert rest_proxy_module.route_name(routes, "POST", "/api/session") == "sign_in" + assert rest_proxy_module.route_name(routes, "GET", "/api/session") == "any_api" + + +@pytest.mark.parametrize("path, method, target, expected", [ + ("/api/notes/{note_id:int}", "GET", "/api/notes/7", "route"), + ("/api/notes/{note_id:int}", "GET", "/api/notes/seven", None), + ("/api/notes/{note_id}", "GET", "/api/notes/seven", "route"), + ("/api/notes/{note_id}", "GET", "/api/notes/7/extra", None), + ("/files/{rest:path}", "GET", "/files/a/b.txt", "route"), + ("/api/price/{p:float}", "GET", "/api/price/3.50", "route"), + ("/api/x/{id:uuid}", "GET", "/api/x/123e4567-e89b-12d3-a456-426614174000", "route"), + ("/api/a.b", "GET", "/api/aXb", None), + ("/api/notes", "HEAD", "/api/notes", "route"), + ("/api/notes", "POST", "/api/notes", None), +]) +def test_a_request_is_named_by_the_card_route_its_method_and_path_reach(path, method, target, expected): + card = {"capabilities": {"extensions": [{"uri": "urn:agentenv:website/v1", "params": { + "routes": [{"method": "GET", "path": path, "name": "route"}]}}]}} + assert rest_proxy_module.route_name(rest_proxy_module.website_routes(card), method, target) == expected diff --git a/tst/unit/env/gateway/triggers_website_test.py b/tst/unit/env/gateway/triggers_website_test.py index f885d541c..6571b0ed1 100644 --- a/tst/unit/env/gateway/triggers_website_test.py +++ b/tst/unit/env/gateway/triggers_website_test.py @@ -12,7 +12,7 @@ from agent_env.env.gateway import gateway as gateway_module from agent_env.env.gateway import rest_proxy as rest_proxy_module from agent_env.env.gateway.clock import Clock -from agent_env.env.gateway.constants import EXT_TRIGGERS_URI, GATEWAY_EXTENSIONS, TRIGGER_WEBSITE +from agent_env.env.gateway.constants import EXT_TRIGGERS_URI, GATEWAY_EXTENSIONS, TRIGGER_WEBSITE, TRIGGER_WEBSITE_OP from agent_env.env.gateway.gateway import GATEWAY_TRAJECTORY_FILE, Gateway from agent_env.env.gateway.triggers import TriggerEngine, TriggerError @@ -27,12 +27,16 @@ def __init__(self, responses: dict) -> None: self.events: list[dict] = [] self.rules: list[tuple[str, str, bool]] = [] self.website_urls = {"mail": "http://mail-frontend", "shop": "http://shop-frontend"} + self.route_names = {"mail": {"send_message", "list_messages"}} self._role_rules_lock = asyncio.Lock() self._server_tools: dict = {} self._tool_server_urls: dict = {} self._clock = Clock() self.TRIGGER_BARRIER_TIMEOUT_S = 30.0 + def website_route_names(self, app): + return self.route_names.get(app) + def _apply_rule(self, role, tool, value): self.rules.append((role, tool, value)) @@ -98,6 +102,10 @@ def _detected(engine) -> list[tuple[str, dict]]: "a website request carries args, not a result"), ({"type": "website", "app": "mail", "unless": {"tool": "t"}}, [], "when.unless is only valid on time triggers"), ({"type": "website", "app": "mail", "roles": []}, [], "when.roles must be a non-empty list"), + ({"type": "website", "app": "mail", "op": ""}, [], "when.op must be a non-empty string"), + ({"type": "website", "app": "mail", "op": "archive"}, [], + "when.op names no route of website 'mail' (routes: ['list_messages', 'send_message'])"), + ({"type": "website", "app": "shop", "op": "checkout"}, [], "whose card the gateway could not read"), ]) def test_registration_fails_loud(monkeypatch, when, actions, fragment): engine = _make_engine(monkeypatch) @@ -121,7 +129,7 @@ def test_registration_normalizes_the_method_and_state_echoes_the_request_matcher def test_the_card_declares_website_triggers(): triggers = next(e for e in GATEWAY_EXTENSIONS if e["uri"] == EXT_TRIGGERS_URI) - assert TRIGGER_WEBSITE in triggers["params"]["features"] + assert {TRIGGER_WEBSITE, TRIGGER_WEBSITE_OP} <= set(triggers["params"]["features"]) # ---------------------------------------------------------------- detection @@ -158,6 +166,21 @@ async def test_only_a_request_matching_app_method_path_status_and_where_provokes assert len(engine._test_gw.calls) == 1 +@pytest.mark.asyncio +async def test_a_trigger_naming_an_op_is_provoked_only_by_a_request_that_reached_that_route(monkeypatch): + engine = _make_engine(monkeypatch) + by_op = {**REPLY, "when": {"type": "website", "app": "mail", "op": "send_message", "repeat": True}} + engine.register({"triggers": [by_op]}) + assert engine._triggers["reply"]["spec"]["when"]["op"] == "send_message" + engine.on_website_request("default", "mail", "POST", "/api/messages", 201, dict(MESSAGE)) + engine.on_website_request("default", "mail", "POST", "/api/messages", 201, dict(MESSAGE), "list_messages") + engine.on_website_request("default", "mail", "POST", "/api/messages", 201, dict(MESSAGE), "send_message") + await _settle(engine) + request = {"app": "mail", "method": "POST", "path": "/api/messages", "status": 201, "op": "send_message"} + assert _detected(engine) == [("reply", {"source": "website", **request, "role": "default"})] + assert len(engine._test_gw.calls) == 1 + + @pytest.mark.asyncio async def test_a_status_predicate_can_match_a_failed_request(monkeypatch): engine = _make_engine(monkeypatch) @@ -335,12 +358,12 @@ async def test_the_proxy_hands_the_engine_the_role_path_status_and_args_of_each_ assert await _request(gw, "GET", "/website/calendar/") == 404 assert seen == [ - ("default", "mail", "GET", "/", 200, {}), - ("carol@example.com", "mail", "GET", "/inbox", 200, {"folder": "work", "tag": ["a", "b"]}), - ("default", "mail", "POST", "/api/messages", 200, {"draft": "0", "to": "dana@example.com"}), - ("default", "mail", "POST", "/compose", 200, {"subject": "Hi there", "to": "dana@example.com"}), - ("default", "mail", "POST", "/api/batch", 200, {"op": "send"}), - ("default", "mail", "PUT", "/api/drafts/a/b", 200, {}), + ("default", "mail", "GET", "/", 200, {}, None), + ("carol@example.com", "mail", "GET", "/inbox", 200, {"folder": "work", "tag": ["a", "b"]}, None), + ("default", "mail", "POST", "/api/messages", 200, {"draft": "0", "to": "dana@example.com"}, None), + ("default", "mail", "POST", "/compose", 200, {"subject": "Hi there", "to": "dana@example.com"}, None), + ("default", "mail", "POST", "/api/batch", 200, {"op": "send"}, None), + ("default", "mail", "PUT", "/api/drafts/a/b", 200, {}, None), ] assert [r.content for r in website.requests if r.url.path != "/.well-known/agent-env.json"][2:4] == [b'{"to": "dana@example.com", "draft": "0"}', b"subject=Hi+there&to=dana%40example.com"] @@ -419,3 +442,152 @@ def explode(*_): gw._trigger_engine.on_website_request = explode assert await _request(gw, "GET", "/website/mail/inbox") == 200 + + +@pytest.mark.asyncio +async def test_with_a_role_key_the_engine_sees_a_role_only_beside_it_so_a_page_cannot_pose_as_the_executor(monkeypatch): + _Website(status=200).install(monkeypatch) + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"mail": "http://mail-frontend"}, role_key="deploy-secret") + gw._trigger_engine.register({"triggers": [REPLY]}) + seen: list[str] = [] + gw._trigger_engine.on_website_request = lambda role, *a: seen.append(role) or ([], None) + + await _request(gw, "GET", "/website/mail/inbox", headers={"AgentEnv-Role": "executor"}) + await _request(gw, "GET", "/website/mail/inbox", headers={"AgentEnv-Role": "executor", "AgentEnv-Role-Key": "deploy-secret"}) + + assert seen == ["default", "executor"] + + +NOTES_CARD = {"name": "notes", "capabilities": {"extensions": [{"uri": "urn:agentenv:website/v1", "params": {"routes": [ + {"method": "GET", "path": "/api/notes", "name": "list_notes", "roles": None}, + {"method": "POST", "path": "/api/notes", "name": "create_note", "roles": None}, + {"method": "PATCH", "path": "/api/notes/{note_id}", "name": "edit_note", "roles": None}, +]}}]}} + + +class _NotesSite(_Website): + """A website whose server serves `NOTES_CARD` at the card's path, counting how often the gateway reads it.""" + + def install(self, monkeypatch) -> None: + self.card_reads = 0 + + async def body(): + yield b'{"ok": true}' + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/.well-known/agent-env.json": + self.card_reads += 1 + return httpx.Response(200, json=NOTES_CARD) + self.requests.append(request) + return httpx.Response(self.status, headers={"content-type": "application/json"}, content=body()) + + transport = httpx.MockTransport(handler) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + + +def _notes_gateway() -> Gateway: + return Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": "http://notes-website-backend:8000"}, + rest_proxy_urls={"notes": "http://notes-website-backend:8000"}) + + +async def _register(gw: Gateway, body: dict) -> tuple[int, dict]: + sent = json.dumps(body).encode() + answered: list[bytes] = [] + statuses: list[int] = [] + + async def receive(): + return {"type": "http.request", "body": sent, "more_body": False} + + async def send(message): + if message["type"] == "http.response.start": + statuses.append(message["status"]) + elif message["type"] == "http.response.body": + answered.append(message.get("body", b"")) + + scope = {"type": "http", "asgi": {"version": "3.0"}, "http_version": "1.1", "method": "POST", "scheme": "http", + "path": "/triggers/register", "raw_path": b"/triggers/register", "query_string": b"", "root_path": "", + "server": ("gateway", 80), "client": ("c", 1), "headers": [(b"content-type", b"application/json")]} + await gw._mcp.streamable_http_app()(scope, receive, send) + return statuses[0], json.loads(b"".join(answered)) + + +EDITED = {"id": "edited", "when": {"type": "website", "app": "notes", "op": "edit_note", "repeat": True}, "actions": []} + + +@pytest.mark.asyncio +async def test_registering_an_op_trigger_reads_the_site_s_routes_and_refuses_a_route_it_does_not_name(monkeypatch): + site = _NotesSite(status=200) + site.install(monkeypatch) + gw = _notes_gateway() + + status, answer = await _register(gw, {"triggers": [{**EDITED, "id": "typo", "when": {**EDITED["when"], "op": "edit"}}]}) + assert status == 400 and "routes: ['create_note', 'edit_note', 'list_notes']" in answer["error"] + status, _ = await _register(gw, {"triggers": [EDITED]}) + assert status == 200 and site.card_reads == 1 + + +@pytest.mark.asyncio +async def test_registering_an_op_trigger_whose_app_is_not_a_name_is_refused_as_bad_config(monkeypatch): + site = _NotesSite(status=200) + site.install(monkeypatch) + gw = _notes_gateway() + + status, answer = await _register(gw, {"triggers": [{**EDITED, "when": {**EDITED["when"], "app": []}}]}) + assert status == 400 and answer["ok"] is False and site.card_reads == 0 + + +@pytest.mark.asyncio +async def test_the_proxy_names_the_route_a_request_reached_to_the_engine_and_the_trajectory(monkeypatch): + site = _NotesSite(status=200) + site.install(monkeypatch) + gw = _notes_gateway() + assert (await _register(gw, {"triggers": [EDITED]}))[0] == 200 + changelog = iter(range(100)) + monkeypatch.setattr(gw, "_query_changelog_id", lambda: next(changelog)) + before = len(_trajectory()) + + await _request(gw, "PATCH", "/website/notes/api/notes/7", headers={"content-type": "application/json"}, body=b"{}") + await _request(gw, "PATCH", "/website/notes/api/notes/7/extra", headers={"content-type": "application/json"}, body=b"{}") + await _request(gw, "POST", "/website/notes/api/notes", headers={"content-type": "application/json"}, body=b"{}") + await asyncio.sleep(0.05) + + detected = [p for _, p in _detected(gw._trigger_engine)] + assert [(p["path"], p.get("op")) for p in detected] == [("/api/notes/7", "edit_note")] + writes = [e for e in _trajectory()[before:] if e["event_type"] == "http_request"] + assert [(w["method"], w["path"], w.get("op")) for w in writes] == [ + ("PATCH", "/api/notes/7", "edit_note"), ("PATCH", "/api/notes/7/extra", None), ("POST", "/api/notes", "create_note")] + + +@pytest.mark.asyncio +async def test_a_request_is_named_by_the_path_the_site_s_server_receives(monkeypatch): + """An encoded slash beside a dot segment routes here as one path and reaches the site's server as another; the + route, the trigger's match and the write's trajectory event all go by the one the server receives.""" + site = _NotesSite(status=200) + site.install(monkeypatch) + gw = _notes_gateway() + assert (await _register(gw, {"triggers": [EDITED]}))[0] == 200 + changelog = iter(range(100)) + monkeypatch.setattr(gw, "_query_changelog_id", lambda: next(changelog)) + before = len(_trajectory()) + + await _request(gw, "PATCH", "/website/notes/public%2Fx/../api/notes/7", headers={"content-type": "application/json"}, + body=b"{}") + await asyncio.sleep(0.05) + + assert [r.url.path for r in site.requests] == ["/api/notes/7"] + assert [(p["path"], p.get("op")) for _, p in _detected(gw._trigger_engine)] == [("/api/notes/7", "edit_note")] + writes = [e for e in _trajectory()[before:] if e["event_type"] == "http_request"] + assert [w.get("op") for w in writes] == ["edit_note"] + + +@pytest.mark.asyncio +async def test_composing_the_env_card_reads_each_site_s_routes(monkeypatch): + site = _NotesSite(status=200) + site.install(monkeypatch) + gw = _notes_gateway() + assert gw.website_route_names("notes") is None + await gw._serve_env_card() + assert gw.website_route_names("notes") == {"list_notes", "create_note", "edit_note"} and site.card_reads == 1 + assert gw._website_op("notes", "GET", "/api/notes") == "list_notes" diff --git a/tst/unit/env/test_env.py b/tst/unit/env/test_env.py index 02947560b..1498888bf 100644 --- a/tst/unit/env/test_env.py +++ b/tst/unit/env/test_env.py @@ -284,7 +284,7 @@ def test_a_gateway_record_is_filled_the_way_it_always_was(): sandbox_type="modal_vm", sandbox_ids={"gateway_server": "gw"}, db_web_url="https://pg/", db_mcp_url="https://dm/mcp", website_frontend_urls={"shop": "https://gw/shop/"}, vnc_url="https://vnc", env_state_instance_ids=["st-1"], environment_card_url=_CARD_URL, environment_card=_CARD, environment_card_read_at_utc="2026-09-28T00:00:00+00:00", - metadata={"k": "v"}, instance_id="i1", created_at_utc="c", expires_at_utc="x"), + metadata={"k": "v"}, instance_id="i1", created_at_utc="c", expires_at_utc="x", role_key="deploy-secret"), ], ids=["base", "sandbox", "gateway"]) def test_a_record_round_trips_through_its_stored_form(record): assert DeployedEnv.from_dict(dataclasses.asdict(record)) == record @@ -295,6 +295,14 @@ def test_each_record_class_loads_every_one_of_its_fields(cls): assert set(cls._fields_from({"env_id": "e", "env_version": 1})) == {f.name for f in dataclasses.fields(cls)} +def test_a_record_s_role_headers_carry_its_role_key_beside_the_role_and_an_older_record_s_only_the_role(): + keyed = DeployedEnv.from_dict({"env_id": "e", "env_version": 1, "gateway_url": "https://gw", "sandbox_id": "sb", + "role_key": "deploy-secret"}) + older = DeployedEnv.from_dict({"env_id": "e", "env_version": 1, "gateway_url": "https://gw", "sandbox_id": "sb"}) + assert keyed.role_headers("admin") == {"AgentEnv-Role": "admin", "AgentEnv-Role-Key": "deploy-secret"} + assert (older.role_key, older.role_headers("admin")) == (None, {"AgentEnv-Role": "admin"}) + + def test_the_gateway_only_fields_are_the_gateway_classs_own(): own = {f.name for f in dataclasses.fields(DeployedGatewayEnv)} - {f.name for f in dataclasses.fields(DeployedSandboxEnv)} assert set(_GATEWAY_ONLY_FIELDS) == own - {"gateway_mode", "env_provider_type"} 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 4345c8fde..7902db84a 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -375,6 +375,7 @@ async def _spy_acquire(self, ctx): ttl_seconds=60, disk_size_gb=10, mcp_server_name="crm", + role_key="deploy-secret", ) # Last create_container call is the gateway. Inspect its env var. @@ -395,6 +396,7 @@ async def _spy_acquire(self, ctx): assert re.search(r"^postgresql://agentenv:agentenv@\[fdaa::fake:[0-9a-f]+\]:5432/agentenv", env["SERVICE_DB_URL"]), env["SERVICE_DB_URL"] assert env["MCP_SERVER_NAME"] == "crm" assert result.mcp_server_name == "crm" + assert env["AGENTENV_ROLE_KEY"] == result.role_key == "deploy-secret" # Each MCP got DATABASE_URL with its scoped search_path mcp_calls = [c for c in create_calls if c["image_name"].startswith("mcp-")] @@ -1191,6 +1193,30 @@ def test_redact_compose_secrets_masks_only_secret_env_lines(): assert "s3cr3t-token-value" in compose +@pytest.mark.parametrize("role_key", ["deploy-secret", None], ids=["keyed", "unkeyed"]) +def test_compose_gives_the_role_key_to_the_gateway_and_every_website_server_and_masks_it_in_the_log(role_key): + """The gateway believes a role only beside the deploy's key, and a site holds the same key to believe the role the + gateway forwards; without a key neither is set, as before it.""" + from agent_env.env.envs.service_db import ServiceDBConfig + + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + gateway_image="agent-gateway", + website_configs=[WebsiteConfig(backend_image="notes-be", frontend_image=None, environment_name="notes")], + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + role_key=role_key, + ) + + services = yaml.safe_load(compose)["services"] + keyed = {name: [entry for entry in service.get("environment") or [] if "ROLE_KEY" in entry] + for name, service in services.items()} + expected = {GATEWAY_SERVICE_NAME: [f"AGENTENV_ROLE_KEY={role_key}"], "notes-website-backend": [f"WEBSITE_ROLE_KEY={role_key}"]} + assert {name: entries for name, entries in keyed.items() if entries} == (expected if role_key else {}) + if role_key: + assert role_key not in _redact_compose_secrets(compose) and " - AGENTENV_ROLE_KEY=***" in _redact_compose_secrets(compose) + + @pytest.mark.asyncio async def test_acquire_state_for_deploy_local_returns_none(): """Local Postgres cannot be pre-provisioned (its container's host doesn't exist yet), so the @@ -1953,12 +1979,28 @@ async def test_create_gateway_hands_the_deploy_path_every_argument(attribution, [call] = vm.await_args_list assert call.args == (provider, args.pop("mcp_servers"), args.pop("mcp_server_images")) - assert call.kwargs == {**args, "attribution": forwarded} + assert call.kwargs == {**args, "attribution": forwarded, "role_key": None} assert call.kwargs["attribution"] is not attribution @pytest.mark.asyncio -async def test_deploy_returns_the_gateway_record_for_the_env(caplog): +@pytest.mark.parametrize("keyed_roles", [True, False], ids=["keyed", "unkeyed"]) +async def test_create_gateway_draws_a_role_key_only_for_a_deploy_with_keyed_roles(keyed_roles): + """A keyed deploy draws a fresh secret for its gateway; any other deploy draws none, so its callers' role header + alone is believed, as before the key.""" + gp, provider = EnvironmentGatewayProvider(), _vm_provider() + 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, mcp_servers=[MagicMock()], mcp_server_images=[MagicMock()], + env_id="crm-env", keyed_roles=keyed_roles) + + role_key = vm.await_args.kwargs["role_key"] + assert (isinstance(role_key, str) and len(role_key) >= 32) if keyed_roles else role_key is None, role_key + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keyed_roles", [False, True], ids=["unkeyed", "keyed"]) +async def test_deploy_returns_the_gateway_record_for_the_env(caplog, keyed_roles): """deploy(env) on a scripted VM: the env's image behind the gateway, its state acquired, and main's record plus its provider type.""" from agent_env.env.gateway import GatewayMode from agent_env.providers.env_state import DatabaseStateProvider @@ -1984,7 +2026,7 @@ async def test_deploy_returns_the_gateway_record_for_the_env(caplog): patch.object(env_gateway_provider, "_tool_names", AsyncMock(return_value={"slack_send"})), \ patch("agent_env.providers.env_providers.env_gateway_provider.asyncio.sleep", AsyncMock()): record = await gp.deploy(env, _vm_provider(create_vm=AsyncMock(return_value=vm)), gateway_mode=GatewayMode.CONSISTENT, - ttl_seconds=60, env_state_type="remote") + ttl_seconds=60, env_state_type="remote", keyed_roles=keyed_roles) clashes_checked.assert_awaited_once_with("https://vm.example") @@ -1997,7 +2039,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"], + role_key=compose.call_args.kwargs["role_key"])) + assert (record.role_key and len(record.role_key) >= 32) if keyed_roles else record.role_key is None 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 "))] diff --git a/tst/unit/task_step/deploy_env_test.py b/tst/unit/task_step/deploy_env_test.py index 874832e4a..0fe426a6e 100644 --- a/tst/unit/task_step/deploy_env_test.py +++ b/tst/unit/task_step/deploy_env_test.py @@ -242,6 +242,36 @@ async def test_execute_merges_annotations_into_metadata_deploy_set(): +@pytest.mark.parametrize("keyed_roles", [True, False], ids=["keyed", "unkeyed"]) +def test_keyed_roles_roundtrips_and_an_older_step_is_unkeyed(keyed_roles): + step = DeployEnvTaskStep(id="t-1.deploy_env", version=None, env_id="art-1", keyed_roles=keyed_roles) + assert DeployEnvTaskStep.from_dict(step.to_dict()).keyed_roles is keyed_roles + older = {k: v for k, v in step.to_dict().items() if k != "keyed_roles"} + assert DeployEnvTaskStep.from_dict(older).keyed_roles is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keyed_roles", [True, False], ids=["keyed", "unkeyed"]) +async def test_execute_asks_for_a_role_key_only_for_a_keyed_step(keyed_roles): + """An unkeyed step passes no keyed_roles, so an env type or plugin provider that doesn't name it isn't broken.""" + env = _fake_env(_deployed_env(None)) + step = DeployEnvTaskStep(id="t-1.deploy_env", version=None, env_id="art-1", keyed_roles=keyed_roles) + with patch("agent_env.env.env.Env") as Env: + Env.get.return_value = env + await step.execute(TaskStepContext()) + assert env.deploy.await_args.kwargs.get("keyed_roles") is (True if keyed_roles else None) + + +@pytest.mark.asyncio +async def test_execute_refuses_keyed_roles_for_an_env_whose_deploy_does_not_take_them(): + env = _env_with_fixed_deploy(_deployed_env(None)) + env.type = "fixed" + step = DeployEnvTaskStep(id="d", version=None, env_id="art-1", keyed_roles=True) + with patch("agent_env.env.env.Env") as Env, pytest.raises(ValueError, match="fixed envs take no keyed_roles"): + Env.get.return_value = env + await step.execute(TaskStepContext()) + + # --- preflight: an option the step's env would refuse at deploy is a save-time problem --- @@ -276,6 +306,12 @@ def test_preflight_refuses_env_params_an_env_doesnt_take(kind, refusal): assert _preflight(env, env_params={"form": "container"}) == [f"deploy_env 'd': {refusal}"] +def test_preflight_passes_keyed_roles_for_a_gateway_env_and_refuses_them_for_a_server_env(): + assert _preflight(_mcp_env("gateway"), keyed_roles=True) == [] + assert _preflight(_mcp_env("server"), keyed_roles=True) == [ + "deploy_env 'd': env 'mcp-email' has env_provider_type 'server', which doesn't take keyed_roles"] + + @pytest.mark.parametrize("kind", ["server", "gateway", "other"]) @pytest.mark.parametrize("mode", ["bogus", "Performance", None]) def test_preflight_reports_an_invalid_gateway_mode_for_every_env(kind, mode): diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index 8320dfcd3..7af4f11e1 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -346,4 +346,34 @@ async def install_cli(deployed, cli_artifact, gateway_url, **placed): await LoadArtifactTaskStep(id="cli", version=None, env_id="env-x", artifact_id=cli.id, artifact_version=cli.version, agent_name="solver").execute(context) - assert installed == [(expected, {"on_host": True} if on_host else {})] + assert installed == [(expected, {"role_key": None, **({"on_host": True} if on_host else {})})] + + +@pytest.mark.asyncio +async def test_each_cli_is_installed_with_the_role_key_of_the_deploy_it_talks_to(monkeypatch, local_stores, tmp_path): + (tmp_path / "cli").mkdir() + (tmp_path / "cli" / "run.sh").write_text("#!/bin/sh\n") + clis = [CliArtifact.put(f"{name}-cli", command_name=name, entrypoint="run.sh", cli_dir=tmp_path / "cli") + for name in ("slack", "crm")] + installed = [] + + async def install_cli(deployed, cli_artifact, gateway_url, **placed): + installed.append((cli_artifact.command_name, gateway_url, placed["role_key"])) + return f"/opt/cli/{cli_artifact.command_name}/run.sh" + + monkeypatch.setattr(A2AAgent, "install_cli", staticmethod(install_cli)) + monkeypatch.setattr(agent_store, "get_a2a_agent_instance_store", + lambda: types.SimpleNamespace(get=lambda instance_id: object())) + context = TaskStepContext( + deployed_envs=[DeployedGatewayEnv(env_id=f"env-{name}", env_version=1, gateway_url=f"https://{name}.gw", + mcp_url="", sandbox_id=f"sb-{name}", sandbox_type="modal_vm", + role_key=f"key-{name}") for name in ("slack", "crm")], + deployed_agents=[DeployedAgent(agent_name="solver", api_url="http://agent", sandbox_id="sb-agent", + sandbox_type="modal_vm", instance_id="inst-1")], + ) + for cli in clis: + await LoadArtifactTaskStep(id=f"load-{cli.command_name}", version=None, env_id=f"env-{cli.command_name}", + artifact_id=cli.id, artifact_version=cli.version, agent_name="solver" + ).execute(context) + + assert installed == [("slack", "https://slack.gw", "key-slack"), ("crm", "https://crm.gw", "key-crm")] diff --git a/tst/unit/task_step/test_deploy_agent_role_delivery.py b/tst/unit/task_step/test_deploy_agent_role_delivery.py index e642fd0de..47bf80c46 100644 --- a/tst/unit/task_step/test_deploy_agent_role_delivery.py +++ b/tst/unit/task_step/test_deploy_agent_role_delivery.py @@ -14,7 +14,7 @@ from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.env.env import DeployedGatewayEnv -from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER +from agent_env.env.gateway import AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_KEY_HEADER from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps import deploy_agent from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep @@ -227,6 +227,30 @@ def test_a_selection_takes_exactly_one_of_disable_and_only(): assert DeployAgentTaskStep.from_dict({**step.to_dict(), "version": 1}).builtin_tools == {"disable": ["Bash"]} +@pytest.mark.asyncio +async def test_a_keyed_deploy_s_registration_carries_its_role_key_and_the_agent_s_environment_does_not(wiring): + wiring.context.deployed_envs[0].role_key = "deploy-secret" + + await wiring.deploy(_card(HEADERS_MCP), role="executor", env_ids=["crm", "ext"]) + + assert wiring.bodies("/ext/mcp-config") == [ + {"url": _GATEWAY_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "executor", AGENT_ENV_ROLE_KEY_HEADER: "deploy-secret"}, + "name": "crm"}, + {"url": _LIVE_MCP_URL, "headers": {AGENT_ENV_ROLE_HEADER: "executor"}}, + ] + assert wiring.env_vars == {ROLE_ENV_VAR: "executor"} + + +@pytest.mark.asyncio +async def test_a_keyed_deploy_gives_its_key_to_no_agent_without_a_role(wiring): + wiring.context.deployed_envs[0].role_key = "deploy-secret" + + await wiring.deploy(_card(HEADERS_MCP), role=None) + + assert wiring.bodies("/ext/mcp-config") == [{"url": _GATEWAY_MCP_URL, "name": "crm"}] + assert wiring.env_vars is None + + class _Wiring: """One deploy_agent run: the requests the agent received, the agent double, and the step's context.""" diff --git a/tst/unit/task_step/test_load_artifact_file.py b/tst/unit/task_step/test_load_artifact_file.py index cde1c353e..c736efa94 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -3,6 +3,7 @@ from __future__ import annotations +import base64 import re from types import SimpleNamespace @@ -275,3 +276,18 @@ async def test_a_cli_installs_on_the_vm_host_of_an_agent_installed_there(monkeyp "pull s3://bucket/check.py /opt/cli/tool/bin/tool", "chmod +x /opt/cli/tool/bin/tool"] assert [s.split("\n")[0] for s in sandbox.scripts if s.startswith("base64 -d")] == [ "base64 -d <<'ENDB64' > /opt/cli/tool/bin/.env"] + + +@pytest.mark.asyncio +async def test_each_installed_cli_carries_its_own_deploy_s_role_key_beside_its_gateway(monkeypatch): + async def _sidecar(role_key: str | None) -> str: + sandbox = _Vm() + monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) + await A2AAgent.install_cli(SimpleNamespace(sandbox_type=None, sandbox_id="sb-agent"), TOOL, "https://gateway", + on_host=True, role_key=role_key) + (written,) = [s for s in sandbox.scripts if s.startswith("base64 -d <<'ENDB64' > /opt/cli/tool/bin/.env")] + return base64.b64decode("".join(written.split("\n")[1:-1])).decode() + + assert await _sidecar("key-a") == "AGENT_ENV_GATEWAY_URL=https://gateway\nAGENT_ENV_ROLE_KEY=key-a\n" + assert await _sidecar("key-b") == "AGENT_ENV_GATEWAY_URL=https://gateway\nAGENT_ENV_ROLE_KEY=key-b\n" + assert await _sidecar(None) == "AGENT_ENV_GATEWAY_URL=https://gateway\n" diff --git a/tst/unit/task_step/test_mcp_cli_builder.py b/tst/unit/task_step/test_mcp_cli_builder.py index af7d6eb74..a2513e5b4 100644 --- a/tst/unit/task_step/test_mcp_cli_builder.py +++ b/tst/unit/task_step/test_mcp_cli_builder.py @@ -52,6 +52,16 @@ def test_codegen_compiles_and_renders_help(): assert "Auto-generated CLI for slack" in group_help.stdout +def test_the_generated_cli_sends_the_role_key_beside_its_role_only_when_one_is_set(monkeypatch, tmp_path): + namespace: dict = {"__name__": "generated_cli", "__file__": str(tmp_path / "svc")} + exec(compile(generate_cli_script("svc"), "svc", "exec"), namespace) + monkeypatch.setenv("AGENT_ENV_ROLE", "admin") + monkeypatch.delenv("AGENT_ENV_ROLE_KEY", raising=False) + assert namespace["_role_headers"]() == {"AgentEnv-Role": "admin"} + monkeypatch.setenv("AGENT_ENV_ROLE_KEY", "deploy-secret") + assert namespace["_role_headers"]() == {"AgentEnv-Role": "admin", "AgentEnv-Role-Key": "deploy-secret"} + + def test_codegen_loads_sibling_env_file(): src = generate_cli_script("svc") diff --git a/tst/unit/task_step/test_prompt_agent_post_turn_hooks.py b/tst/unit/task_step/test_prompt_agent_post_turn_hooks.py index f54d8a200..f485117b9 100644 --- a/tst/unit/task_step/test_prompt_agent_post_turn_hooks.py +++ b/tst/unit/task_step/test_prompt_agent_post_turn_hooks.py @@ -47,6 +47,7 @@ def __init__(self, monkeypatch, replies, *, clock: dict | None = None, engine: T self.clock = clock or {"armed": False} self.engine = engine self.steps: list[tuple[str, dict]] = [] + self.step_keys: list[str | None] = [] def _answer(self, request: httpx.Request) -> httpx.Response: origin = f"{request.url.scheme}://{request.url.host}" @@ -62,6 +63,7 @@ def _answer(self, request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"to": body.get("to", "default"), "items": [], "until": None}) if origin == ENV and request.url.path == "/step": self.steps.append((request.headers.get("AgentEnv-Role"), json.loads(request.content))) + self.step_keys.append(request.headers.get("AgentEnv-Role-Key")) return httpx.Response(200, json={"content": [{"type": "text", "text": "ok"}], "changelog_id": 7}) return super()._answer(request) @@ -475,6 +477,26 @@ async def post_turn(turn, params): "arguments": {"to": "a@example.com"}, "source": "post_turn_hook", "hook_artifact_id": "mailer"})] assert _recorded(ctx)[0]["result"]["changelog_id"] == 7 + assert peers.step_keys == [None] + + +@pytest.mark.asyncio +async def test_a_hook_s_tool_call_carries_the_deploy_s_role_key_beside_its_role(monkeypatch): + peers = EnvPeers(monkeypatch, {AGENT: [text("one")]}) + ctx = _context() + env = _env(EXT_STEP_URI) + env.role_key = "deploy-secret" + ctx.deployed_envs.append(env) + hook = _hook("keyed", """ + async def post_turn(turn, params): + [env] = turn.envs + await env.call_tool("send_email", {}, role="school") + return None + """) + + await _step(hook, turns=1).execute(ctx) + + assert [role for role, _ in peers.steps] == ["school"] and peers.step_keys == ["deploy-secret"] @pytest.mark.asyncio diff --git a/tst/unit/task_step/test_register_env_triggers.py b/tst/unit/task_step/test_register_env_triggers.py index 1abba6959..78f370aee 100644 --- a/tst/unit/task_step/test_register_env_triggers.py +++ b/tst/unit/task_step/test_register_env_triggers.py @@ -273,11 +273,13 @@ async def test_executor_hours_for_an_agent_no_trigger_uses_raise(monkeypatch): ([{"id": "x", "when": {"type": "action", "tool": "t"}, "actions": [{"type": "nl", "instruction": "x"}]}], {"exec-1": _HOURS}, "executor_hours"), ([{"id": "x", "when": {"type": "website", "app": "mail", "method": "POST"}, "actions": []}], None, "website"), -], ids=["delay", "unless", "hours", "website"]) + ([{"id": "x", "when": {"type": "website", "app": "mail", "op": "send"}, "actions": []}], None, "website_op"), +], ids=["delay", "unless", "hours", "website", "website-op"]) async def test_a_gateway_without_a_feature_refuses_the_triggers_that_use_it(monkeypatch, triggers, hours, missing): sent = _mock_gateway(monkeypatch, {}) context = _context(with_executor=True) - context.deployed_envs = [_record(_with_features("named_executors", "provoking_roles"))] + features = ("named_executors", "provoking_roles", "website") if missing == "website_op" else ("named_executors", "provoking_roles") + context.deployed_envs = [_record(_with_features(*features))] with pytest.raises(RuntimeError, match=rf"does not declare \['{missing}'\]"): await _step(triggers=triggers, executor_agent_name="exec-1", executor_hours=hours).execute(context) assert sent == []