diff --git a/AGENTS.md b/AGENTS.md index 1e1849646..f10da0cf3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -117,7 +117,7 @@ and deserialized through a registry. |---|---| | `env/` | `Env` base class: a subclass declares `type`, implements `from_dict` and `async deploy(...) -> DeployedEnv`, and may override `reset`, `close` and `load_file_artifact_universe`; teardown closes each env a run deployed before terminating its sandboxes by id. Built-ins in `env/envs/`: `MCPServerEnv` (`mcp_server`), `MultiEnv` (`multi`, MCP servers plus websites), `ServiceDBEnv` (`service_db`, PostgreSQL), `WebsiteEnv` (`website`), `GatewayEnv` (`gateway_server`), `ComputerDriverEnv` (`computer_driver`, `env/envs/computer/driver_env.py`: a computer driver's files and what its class declares for each platform it supports, read at put in a process of its own, `agentenv_protocol.computer._declare`, where the driver's own dependencies needn't be installed, with the `port` it serves on and its install by OS, `[install.]` (`env/envs/computer/utils/install.py`): files, apt `packages` with extra `sources`, a `python` requirements file for a venv, `downloads` by URL and sha256, `run` scripts, the services the base's init runs, and `ports`; it deploys only as a computer env's driver), `ComputerEnv` (`computer`, in `env/envs/computer/`: a base computer and a driver, matched at put by gate 1 before anything is stored, then composed by `composer.py` into an image for each form and arch on a builder of the put's sandbox provider, by default the env's (`utils/builder.py`; `--sandbox` on `env build`, `bundle register` and `run`): this machine's Docker for `local`, else a VM of the provider, which fetches the put's stored files by signed URL, takes the rest inline and pushes to the image store itself, so it builds only with stores a remote VM reaches; each arch goes to the first link that runs it natively, else one that emulates it (a VM registers QEMU; `--no-emulate` leaves it `skipped` instead, for a later put), and its pair records the provider that built it (`builder`); a run or register that reuses a computer env composes what it has left (`pending`); tagged `agentenv-computers:` in the image store by a key of the image's Dockerfile, build context and lock so any env with the same inputs reuses it; every layer's packages, requirements and downloads are resolved at put in a container of the base for each arch (`utils/packages.py`, `utils/resolve_in_base.py`), stored in the object store by sha256 under `agentenv-packages/`, and listed in a lock stored by a digest of its inputs, which any env with the same inputs reuses, and builds install only the lock's files, every install step `RUN --network=none`; each pair's status goes into its metadata, `composed`, the arches built are written as its computer, a `computer` artifact over one index, and `agent-env env build` resumes what a put left; gate 3 then boots each native build on the put's sandbox provider and probes it at `full`, making it `ready` or `failed`, and a build no provider runs natively stays `unverified`. Its deploy, through the `computer` env provider (`providers/env_providers/env_computer_provider.py`), boots its first `ready` build on a sandbox provider that runs its arch natively (`native_platforms()`), with no gateway in front and the driver's card as the env card, then checks its `display`; `deploy_env` `env_params` take a `form`, `computer_settings` (deploy-phase values of the base's settings; a secret one, here or on the env, only as a bare `secret:`/`env:` reference) and `driver_settings` (the driver's `ComputerSettings`, such as its default coordinate space); a `network` that allows nothing runs it on local Docker's internal network behind a relay (`providers/sandbox_providers/local_docker/relay/`); `reset` and `close` have the driver end its recordings first (`computer-recording/v1` `end_all`, each uploading what's left and its manifest under the grant it started with), and `reset` then boots a fresh container in its place (`reset_container`), keeping the computer's address but not an agent's MCP session to the old driver; its `load_file_artifact_universe` writes files under the session user's home through the driver's `data/add`, `Downloads` by default, each file's bytes through a read grant when the store's grants reach it and its network allows everything, else inline (`env/envs/computer/utils/inside.py`)), `ComputerAppEnv` (`computer_app`, `env/envs/computer/app_env.py`: an app on a computer, a `ComputerAppEnvironment` class its env.toml's `module` names, read at put by the same `_declare`, which declares each OS with its own decorator (`@linux(...)`, or `@install(os_type, spec)` for any OS type) holding its install, its launch and, optionally, a skill and settings and build settings that add to the class's there, and install steps and ready checks for every OS or those their `os=` names; everything that reads its settings, launch or steps asks for the computer's OS type; deployed alone it is a one-app computer env, a `ComputerEnv` of the computer and driver its env.toml names, the composer installing it under `/opt/agentenv/apps//` with its services as supervisord programs beside its own server (`agentenv_protocol.computer.serve_app`, in a venv with the protocol's `app` extra, on `.agentenv/server.sock` in its folder), which adds its launcher, launches a `boot` app once per boot and says through `computer-app/v1` whether it's ready, as the driver's `computer-ready` asks, its card named after it, and `env_params` `app_settings` handing it values of its `Settings` in `AGENTENV_APP_SETTINGS`, a secret one only as a bare reference; its env.toml's `build_settings`, values of its `BuildSettings` checked at put, are written to `.agentenv/build_settings.json` in its folder, so they're part of the image's key, its install steps read them, and a deploy can't change them). An env type whose services take environment data defines `load_environment_artifact` and `load_environment_universe_artifact`; `load_artifact` refuses an environment artifact aimed at a type that doesn't, at save and before it loads. An env built on the protocol SDK can serve agents a skill (`@agent_skill`, `skill/v1`), which `add_skills` `env_ids` installs under its environment name (behind a gateway, each server's under that server's name); a skill with files beside its SKILL.md goes as a bundle stored once per content under `agentenv-env-skills/` in the object store. `env/gateway/gateway.py` is the env-side gateway server, container-only code. While its clock is armed it serves agents `get_time` (`env/gateway/get_time.py`), which answers in UTC and, given an IANA `time_zone`, also in that zone; executors' working `hours` (`env/gateway/triggers.py`) take a fixed `utc_offset` or an IANA `time_zone`. | | `artifact/` | Immutable versioned artifacts; every `put()` creates a new version. The document goes to the document store, binary payloads to the object store. Every artifact but an environment universe, whose `metadata` key holds its legacy metadata file ids, carries a `metadata` dict outside its versioned fields, changed by compare-and-swap (`merge_metadata`). Built-ins: `CliArtifact` (`cli`), `ComputerArtifact` (`computer`, a base computer: its architectures, OS, session user, settings and container form, whose image `put` reads from its registry into `metadata.validated_images`), `FileArtifact` (`file`), `FileArtifactUniverse` (`file_artifact_universe`), `DockerImageArtifact` (`docker_image`), `EnvironmentArtifact` (`environment`: its environment's files by path, stored as the one wrapped `file_artifact_ref` every environment artifact had when it has one file at its filename, else as `file_artifact_refs`; `get_file_artifacts()` reads either, and a load hands each file to the env's `data/add` with its path, which a pre-v1 load, taking one file, refuses), `EnvironmentUniverseArtifact` (`environment_universe`), `SkillArtifact` (`skill`); `[artifacts] type_aliases` maps another stored spelling of a type onto one of them. A `DockerImageArtifact` holds a tar.gz of its image, or, written by `put_ref` (`--image-ref` on `env mcp-server put` and `a2a-agent put`), only a registry reference with its tag pinned to a digest, which sandboxes pull, or a published image's tag with its `image_digest` and its build context, which sandboxes pull by the digest and build from the context when they can't (the infra images, and a bundle agent built from a shipped agent's sources, `Plan.published`), or, written by `put_context`, only its build context, which a VM sandbox builds before running it (`VmSandbox.load_docker_images`) and Modal builds before creating its container (`ModalSandboxProvider.prepare_image`, outside a fallback chain's deadline); anything else that runs images by name (a container on the local provider, a lone server on any provider but Modal) refuses such an image. A build context (`utils/docker/context.py`) is what `docker build` would send, `.dockerignore` applied, written so equal files give equal bytes and an equal `source_digest`. | -| `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, `start_computer_recording` and `end_computer_recording`, and more). On an env whose card serves computer-exec/v1 (a computer), `run_code` with `env_id` and `run_container_unit_tests_verifier` with `env_id` run inside it as its session user through the driver: files staged with `data/add` and read with `data/get`, commands run with `run`; `env_outcome_verifier` hands a `verify` that takes an `env_url` keyword the env's address. `teardown_run` and the `teardown_sandboxes` step share one core (`task/teardown.py` `tear_down`): the recordings the run left open end with fresh grants, each env still up closes (`Env.close`; a computer env ends any recording still open), what stayed open is registered from its namespace, then every sandbox is terminated by id. `task/episode/` runs a task with some of its agents played from outside, as a training loop does (`EpisodeSpec`). | +| `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, `start_computer_recording` and `end_computer_recording`, and more). On an env whose card serves computer-exec/v1 (a computer), `run_code` with `env_id` and `run_container_unit_tests_verifier` with `env_id` run inside it as its session user through the driver: files staged with `data/add` and read with `data/get`, commands run with `run`; `env_outcome_verifier` hands a `verify` that takes an `env_url` keyword the env's address. A computer's framework-only routes (everything but MCP and the card) take the per-deploy control token its provider boots it with (`DeployedEnv.control_token`, `AGENTENV_CONTROL_TOKEN`); framework calls send `deployed.control_headers`, `Task.run` runs every step under `protocol_v1.authorized` with each deployed env's (a `verify` included), and a `run_code` script inside the computer gets `input["env"]["headers"]`. `teardown_run` and the `teardown_sandboxes` step share one core (`task/teardown.py` `tear_down`): the recordings the run left open end with fresh grants, each env still up closes (`Env.close`; a computer env ends any recording still open), what stayed open is registered from its namespace, then every sandbox is terminated by id. `task/episode/` runs a task with some of its agents played from outside, as a training loop does (`EpisodeSpec`). | | `store/` | Four store ABCs with local and cloud implementations: `DocumentStore` (SQLite, MongoDB, DynamoDB, Firestore with MongoDB compatibility), `ObjectStore` (filesystem, S3, Cloud Storage), `ImageStore` (local OCI registry, ECR), `SecretStore` (env vars or file, AWS Secrets Manager, Google Cloud Secret Manager). `VersionedEntityStore` implements the shared versioned get/put logic, `QueryBuilder` is the immutable chained query API, `store/base.py` holds the error types. `store/routing.py` is the `@local` namespace: while the CLI or a `build_app` server runs, an `@local/…` id and what an `@local` task's run records go to local stores beside the configured ones, and that run's writes to the configured stores are refused. A new backend must pass the conformance kits in `tst/store/`. | | `config/` | The `Config` singleton (`get_config`, `configure`, `reset_config`) in `config/runtime.py`, file discovery in `config/loader.py`, and `load_impl`, which resolves `module:Class` pointers. `agent_env.store` re-exports the config names for compatibility. | | `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`, `sail_vm` (the `sail` extra), `smol_vm` (the `smol` extra), `tensorlake` (the `tensorlake` extra), `tenki` (the `tenki` extra), `vercel` (the `vercel` extra), `boxd` (the `boxd` extra), `coreweave` (the `coreweave` extra), `daytona` (the `daytona` extra), `boat` (the `boat` extra); `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, whose deploys skip a provider that can't run them. A sandbox provider declares what it runs (`ON_THIS_MACHINE`, and for VMs on this machine the address their requests to it arrive at, `REACHES_HOST_AT`; `CONTAINER_ACCEPTS`, `PRIVATE_NETWORK`; `SUPPORTED_SANDBOX_MODES`, derived from whether it implements `create_sandbox` and `create_vm`; and `SANDBOX_ACCEPTS`, the images an agent's sandbox runs, which defaults to those a VM loads when it implements `create_vm` and to `CONTAINER_ACCEPTS` otherwise, so only a provider that implements `create_vm` yet gives an agent a container, like the local one, declares it), each sandbox carries the `platform` of the images it runs (linux/amd64 unless its provider sets another when it creates it; None, any, unchecked, for a local sandbox), and code outside `providers/sandbox_providers/` reads those, never which provider it has. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store), `EnvironmentGatewayProvider` (`gateway`, an env's default), which renders a docker-compose for the gateway and its MCP servers inside the sandbox, and `EnvironmentServerProvider` (`server`), one MCP server in its own container with no gateway in front and no state store; `providers/env_state/` holds env-state providers (`local_postgres` built in). | @@ -238,6 +238,9 @@ orchestrators call it the same way. What callers rely on: through `record_task_failure()`, retries through `undo_steps()`. - The runner writes its run identifier to `context.metadata["workflow_id"]` before calling `run()`; core stores it opaquely with the instance. +- Each step runs under `protocol_v1.authorized` for the deployed envs in its context: a protocol-client + call the step makes to a deployed env's address, naming no headers of its own, sends that env's + `DeployedEnv.control_headers`. - A runner records every run it accepts in the run index (`agent_env.runner.store`): `insert_run` at submit, `set_running` when it starts, `mark_terminal` when it ends. The explorer lists, counts and groups runs from that collection and never calls the runner to ask, so a run missing there is diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 8cc9fb639..5462c89e6 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -227,6 +227,23 @@ async def display_get(self) -> dict: ... async def display_set(self, width: int, height: int) -> dict: ... ``` +### Framework-only routes and the control token + +Everything an environment serves but its MCP endpoint and its card is agent-env's own: the data +plane, every `@extension` route and the tool-access routes. When agent-env deploys an environment +that takes one, it boots it with a per-deploy token in `AGENTENV_CONTROL_TOKEN` +(`CONTROL_TOKEN_ENV`), and the SDK then refuses a call to those routes that doesn't carry it as +`Authorization: Bearer `, answering 401 with the code `unauthorized` before the route runs. +The MCP endpoint and the card stay open, so an agent reaches the environment's tools as before and +nothing else. Without the variable, as when the environment runs on its own, every route is open. +`control_guard()` wraps a route an environment adds itself the same way, as a computer driver does +its `/svc//` routes. + +On the client side each call takes `headers`, and `client.control_headers(token)` builds them. +Code that names none runs under `client.authorized({env_url: client.control_headers(token)})`: +within it, a call to an address under `env_url` sends those headers. agent-env runs every task +step, a verifier's `verify` included, under such a grant for each environment it deployed. + ### Tool access without a gateway An environment agents reach directly, with no gateway in front, can enforce per-role tool access @@ -588,7 +605,10 @@ report = await run_probe("http://localhost:18765", "full") # a ProbeReport failed = [c for c in report.checks if c.status == "failed"] ``` -The CLI exits 0 when no check failed, 1 when one did and 2 when it can't reach the driver. +The CLI exits 0 when no check failed, 1 when one did and 2 when it can't reach the driver. A +driver booted with a control token is probed with it: `--token`, by default +`$AGENTENV_CONTROL_TOKEN`, or `run_probe(..., token=...)`. Given one, the probe also checks that +`computer-ready`'s status refuses a call without it and with a wrong one. | Level | Checks | Needs on the computer | |---|---|---| diff --git a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py index b081ab73c..61f9cbab9 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py @@ -53,6 +53,7 @@ manifest_compatible, ) from .types import ( + CONTROL_TOKEN_ENV, DATA_OBJECTS_EXTENSION_URI, INTAKE_EXTENSION_URI, MCP_PATH, @@ -161,6 +162,7 @@ "ExtensionCallError", "PROTOCOL_VERSION", "RPC_PATH", + "CONTROL_TOKEN_ENV", "WELL_KNOWN_PATH", "METHOD_RESET", "METHOD_ADD", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index c88634844..de6450ae1 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -5,6 +5,7 @@ import decimal import enum import functools +import hmac import inspect import json import logging @@ -28,7 +29,7 @@ from .tool_access import extensions as tool_access_extensions from .tool_context import EMPTY, ToolContext, bound, context_from_http, context_from_mcp from .transfers import WriteNamespaceGrant -from .types import DATA_OBJECTS_EXTENSION_URI, MCP_TRANSPORT, SKILL_EXTENSION_URI, METHOD_ADD, METHOD_GET, METHOD_RESET, RPC_PATH, WELL_KNOWN_PATH, AddDataRequest, AddDataResponse, EnvironmentCapabilities, EnvironmentCard, EnvironmentExtension, EnvironmentInterface, EnvironmentTool, GetDataResponse, ResetDataResponse, error_body +from .types import CONTROL_TOKEN_ENV, DATA_OBJECTS_EXTENSION_URI, MCP_TRANSPORT, SKILL_EXTENSION_URI, METHOD_ADD, METHOD_GET, METHOD_RESET, RPC_PATH, WELL_KNOWN_PATH, AddDataRequest, AddDataResponse, EnvironmentCapabilities, EnvironmentCard, EnvironmentExtension, EnvironmentInterface, EnvironmentTool, GetDataResponse, ResetDataResponse, error_body logger = logging.getLogger(__name__) @@ -343,10 +344,11 @@ def add_routes_to_app(self, app: Any, *, rpc_url: str = RPC_PATH, card_url: str if isinstance(self.handler, AgentEnvEnvironment): self.handler._application = self self.environment_card = self._card_for(app) - self._add_route(app, rpc_url, ["POST"], self._dispatch) + guard = control_guard() + self._add_route(app, rpc_url, ["POST"], guard(self._dispatch)) self._add_route(app, card_url, ["GET"], self._serve_card) for path, method, handler in self._ext_routes: - self._add_route(app, path, [method], handler) + self._add_route(app, path, [method], guard(handler)) for descriptor, fn in self._tools: self._register_tool(app, descriptor, fn) for resource in self._resources: @@ -1544,6 +1546,25 @@ def _error_response(error: AgentEnvError) -> Response: return _json_response(error_body(**error.fields()), error.status) +def control_guard() -> Callable[[Callable], Callable]: + """What wraps a framework-only route, everything an env serves but its MCP endpoint and its card: with + ``CONTROL_TOKEN_ENV`` set, a call that doesn't carry that token as its bearer is refused before the route runs; + without it, nothing.""" + token = os.environ.get(CONTROL_TOKEN_ENV) + if not token: + return lambda handler: handler + expected = f"Bearer {token}".encode() + + def guard(handler: Callable) -> Callable: + async def guarded(request: Request) -> Response: + if not hmac.compare_digest(request.headers.get("authorization", "").encode(), expected): + return _error_response(AgentEnvError( + "unauthorized", "this route is agent-env's own: it takes the env's control token", status=401)) + return await handler(request) + return guarded + return guard + + def _tool_access_handler(access: ToolAccess, action: str) -> Callable: async def handler(request: Request) -> Response: raw = await request.body() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/client.py b/packages/agentenv-protocol/src/agentenv_protocol/client.py index 33f6f2375..dcddaf541 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/client.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/client.py @@ -1,7 +1,12 @@ -"""Data-plane wire protocol client (JSON-RPC) for servers implementing the AgentEnv protocol.""" +"""Data-plane wire protocol client (JSON-RPC) for servers implementing the AgentEnv protocol. A call to an env's +framework-only routes carries its control token: as the call's ``headers`` (``control_headers``), or, for code that +names none, from the grant it runs under (``authorized``).""" from __future__ import annotations +import contextlib import logging +from collections.abc import Iterator +from contextvars import ContextVar import httpx @@ -11,6 +16,9 @@ logger = logging.getLogger(__name__) +#: The headers that carry the control token of each env address the running code is granted, by address. +_GRANTS: ContextVar[dict[str, dict[str, str]]] = ContextVar("agentenv_control_grants", default={}) + class ExtensionCallError(httpx.HTTPStatusError): """A non-2xx answer from an extension method, with the error fields the server sent. A body without them @@ -35,10 +43,45 @@ def _error_fields(response: httpx.Response) -> dict: return error if isinstance(error, dict) else {} -async def _rpc(base_url: str, method: str, params: dict, timeout: int, verify: bool) -> dict: +def control_headers(token: str | None) -> dict[str, str]: + """The headers that carry an env's control token to its framework-only routes; none without one.""" + return {"Authorization": f"Bearer {token}"} if token else {} + + +@contextlib.contextmanager +def authorized(grants: dict[str | None, dict[str, str]]) -> Iterator[None]: + """Within it, a call to an address under one of ``grants``, an env's address and the headers that carry its + control token (``control_headers``), sends those headers when the call names none of its own: how the steps + agent-env runs, and the verifiers they call, reach an env's framework-only routes as written.""" + granted = _GRANTS.set({**_GRANTS.get(), + **{url.rstrip("/"): sent for url, sent in grants.items() if url and sent}}) + try: + yield + finally: + _GRANTS.reset(granted) + + +def request_headers(url: str, headers: dict[str, str] | None = None) -> dict[str, str]: + """What a call to ``url`` sends: ``headers`` when it names them, else the headers its grant gives + (``authorized``), if any.""" + if headers is not None: + return headers + return granted_headers(url, _GRANTS.get()) + + +def granted_headers(url: str, grants: dict[str | None, dict[str, str]]) -> dict[str, str]: + """What ``grants``, env addresses and the headers that carry their control tokens, give a call to ``url``: those of + the address it is or is under, and none anywhere else.""" + return next((sent for base, sent in grants.items() + if base and (url == base.rstrip("/") or url.startswith(f"{base.rstrip('/')}/"))), {}) + + +async def _rpc(base_url: str, method: str, params: dict, timeout: int, verify: bool, + headers: dict[str, str] | None) -> dict: request = {"jsonrpc": "2.0", "id": 1, "method": method, "params": params} + url = f"{base_url}{RPC_PATH}" async with httpx.AsyncClient(verify=verify) as client: - response = await client.post(f"{base_url}{RPC_PATH}", json=request, timeout=timeout) + response = await client.post(url, json=request, timeout=timeout, headers=request_headers(url, headers)) response.raise_for_status() body = response.json() if "error" in body: @@ -47,18 +90,20 @@ async def _rpc(base_url: str, method: str, params: dict, timeout: int, verify: b return body["result"] -async def reset_data(base_url: str, timeout: int = 30, verify: bool = True) -> ResetDataResponse: - return ResetDataResponse.model_validate(await _rpc(base_url, METHOD_RESET, {}, timeout, verify)) +async def reset_data(base_url: str, timeout: int = 30, verify: bool = True, *, + headers: dict[str, str] | None = None) -> ResetDataResponse: + return ResetDataResponse.model_validate(await _rpc(base_url, METHOD_RESET, {}, timeout, verify, headers)) -async def add_data(base_url: str, parts: list[Part], timeout: int = 120, verify: bool = True) -> AddDataResponse: +async def add_data(base_url: str, parts: list[Part], timeout: int = 120, verify: bool = True, *, + headers: dict[str, str] | None = None) -> AddDataResponse: params = {"parts": [p.model_dump(mode="json", exclude_none=True) for p in parts]} - return AddDataResponse.model_validate(await _rpc(base_url, METHOD_ADD, params, timeout, verify)) + return AddDataResponse.model_validate(await _rpc(base_url, METHOD_ADD, params, timeout, verify, headers)) async def get_data( base_url: str, timeout: int = 30, verify: bool = True, *, write_namespace: WriteNamespaceGrant | None = None, - paths: list[str] | None = None, + paths: list[str] | None = None, headers: dict[str, str] | None = None, ) -> GetDataResponse: """`data/get`. A server advertising `DATA_OBJECTS_EXTENSION_URI` may be handed `write_namespace` to upload its export under; `uploaded_object_path` reads from the answer where it did. `paths` names what a computer @@ -66,7 +111,7 @@ async def get_data( params = {} if write_namespace is None else {"write_namespace": write_namespace.model_dump(mode="json", exclude_none=True)} if paths is not None: params["paths"] = paths - return GetDataResponse.model_validate(await _rpc(base_url, METHOD_GET, params, timeout, verify)) + return GetDataResponse.model_validate(await _rpc(base_url, METHOD_GET, params, timeout, verify, headers)) async def get_card(base_url: str, timeout: int = 10, verify: bool = True) -> dict: @@ -180,7 +225,8 @@ def find_child(card: dict, name: str) -> dict | None: return None -async def get_interface_manifest(base_url: str, interface: str, endpoint: str = INTERFACE_MANIFEST_PATH, timeout: int = 30, verify: bool = True): +async def get_interface_manifest(base_url: str, interface: str, endpoint: str = INTERFACE_MANIFEST_PATH, timeout: int = 30, verify: bool = True, *, + headers: dict[str, str] | None = None): """Fetch one interface's manifest from a server's manifest index endpoint. ``endpoint`` is the index path — pass the one advertised on the card under @@ -191,12 +237,13 @@ async def get_interface_manifest(base_url: str, interface: str, endpoint: str = """ url = f"{base_url.rstrip('/')}{endpoint}/{interface}" async with httpx.AsyncClient(verify=verify) as client: - response = await client.get(url, timeout=timeout) + response = await client.get(url, timeout=timeout, headers=request_headers(url, headers)) response.raise_for_status() return response.json() -async def invoke_extension(base_url: str, card: dict, uri: str, params: dict | None = None, timeout: int = 30, verify: bool = True, *, method: str | None = None): +async def invoke_extension(base_url: str, card: dict, uri: str, params: dict | None = None, timeout: int = 30, verify: bool = True, *, method: str | None = None, + headers: dict[str, str] | None = None): """Invoke a card extension via its advertised REST endpoint. Reads the endpoint and HTTP verb of the extension method named `method` from the card, so the @@ -215,11 +262,12 @@ async def invoke_extension(base_url: str, card: dict, uri: str, params: dict | N if not op["endpoint"]: raise ValueError(f"extension {uri} has no endpoint in params") url = f"{base_url}{op['endpoint']}" + sent = request_headers(url, headers) async with httpx.AsyncClient(verify=verify) as client: if op["method"] == "GET": - response = await client.get(url, params=params or {}, timeout=timeout) + response = await client.get(url, params=params or {}, timeout=timeout, headers=sent) else: - response = await client.request(op["method"], url, json=params or {}, timeout=timeout) + response = await client.request(op["method"], url, json=params or {}, timeout=timeout, headers=sent) if response.status_code >= 400: raise ExtensionCallError(response) return response.json() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__init__.py index 58cb6c6db..b102b1be4 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__init__.py @@ -11,9 +11,11 @@ __all__ = ["run_probe", "ProbeReport", "CheckResult", "Level", "Status", "Verdict"] -async def run_probe(url: str, level: Level = "full", *, timeout_s: float = 60, ready_timeout_s: float = 60) -> ProbeReport: - """Probe the driver serving its card at ``url``. ``ready_timeout_s`` is how long it may take to report ready.""" - driver = await Driver.connect(url, timeout_s) +async def run_probe(url: str, level: Level = "full", *, timeout_s: float = 60, ready_timeout_s: float = 60, + token: str | None = None) -> ProbeReport: + """Probe the driver serving its card at ``url``, sending ``token``, its control token, to its framework-only routes. + ``ready_timeout_s`` is how long it may take to report ready.""" + driver = await Driver.connect(url, timeout_s, token) async with driver.session() as session: probe = Probe(driver, session, ready_timeout_s) await probe.basic() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__main__.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__main__.py index 06fd2ad51..bf502a101 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__main__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/__main__.py @@ -1,12 +1,15 @@ """``python -m agentenv_protocol.computer.conformance ``: probe a running driver, print each check, and -exit 1 when one failed or 2 when the driver couldn't be reached.""" +exit 1 when one failed or 2 when the driver couldn't be reached. ``--token`` is its control token, by default +``CONTROL_TOKEN_ENV``'s.""" from __future__ import annotations import argparse import asyncio +import os import sys from pathlib import Path +from ...types import CONTROL_TOKEN_ENV from . import run_probe _MARK = {"passed": "PASS", "failed": "FAIL", "unverified": "----", "not_applicable": "n/a "} @@ -19,9 +22,12 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--level", choices=("basic", "full"), default="full") parser.add_argument("--json", type=Path, help="also write the report to this file") parser.add_argument("--timeout", type=float, default=60, help="seconds for each call, and to report ready") + parser.add_argument("--token", default=os.environ.get(CONTROL_TOKEN_ENV), + help=f"the driver's control token, if it was given one (default: ${CONTROL_TOKEN_ENV})") args = parser.parse_args(argv) try: - report = asyncio.run(run_probe(args.url, args.level, timeout_s=args.timeout, ready_timeout_s=args.timeout)) + report = asyncio.run(run_probe(args.url, args.level, timeout_s=args.timeout, ready_timeout_s=args.timeout, + token=args.token)) except Exception as e: print(f"couldn't probe {args.url}: {type(e).__name__}: {e}", file=sys.stderr) return 2 diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/checks.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/checks.py index 16f139d89..ebcf64843 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/checks.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/checks.py @@ -88,7 +88,8 @@ async def basic(self) -> None: for check, section in (("P0", self._ready), ("P1", self._card), ("P2", self._frame), ("P5", self._refusals), ("X", self._display), ("X", self._spaces), ("X", self._keyboard_settings), ("X", self._exec), ("X", self._data_get), ("X", self._app_exports), - ("X", self._tool_access), ("X", self._pointer), ("X", self._skill)): + ("X", self._tool_access), ("X", self._pointer), ("X", self._skill), + ("X", self._control)): await self._guarded(check, section) self._grants() @@ -422,6 +423,24 @@ async def export(app: dict = app) -> str | None: await self.check("X", f"app {app.get('name')}'s data/get answers its parts through the driver", export, *DATA_PLANE) + async def _control(self) -> None: + if not self.driver.headers: + return + + async def refused() -> str | None: + for sent, without in (({}, "without"), (protocol.control_headers("not-its-token"), "with a wrong")): + try: + await protocol.invoke_extension(self.driver.url, self.driver.card, COMPUTER_READY, method="status", + timeout=int(self.driver.timeout_s), headers=sent) + except ExtensionCallError as e: + if e.response.status_code != 401: + return f"computer-ready status answered a call {without} control token {e.response.status_code}" + continue + return f"computer-ready status answered a call {without} control token" + return None + + await self.check("X", "a framework-only route refuses a call without the control token", refused) + async def _tool_access(self) -> None: if not (self.driver.declares(DISABLE_TOOL_EXTENSION_URI) and self.driver.declares(ENABLE_TOOL_EXTENSION_URI)): return diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/client.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/client.py index a941e4f52..6750f84d4 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/client.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/conformance/client.py @@ -24,14 +24,15 @@ class Driver: - def __init__(self, url: str, card: dict, timeout_s: float) -> None: + def __init__(self, url: str, card: dict, timeout_s: float, token: str | None = None) -> None: self.url = url.rstrip("/") self.card = card self.timeout_s = timeout_s + self.headers = protocol.control_headers(token) if token else None @classmethod - async def connect(cls, url: str, timeout_s: float) -> Driver: - return cls(url, await protocol.get_card(url.rstrip("/"), timeout=int(timeout_s)), timeout_s) + async def connect(cls, url: str, timeout_s: float, token: str | None = None) -> Driver: + return cls(url, await protocol.get_card(url.rstrip("/"), timeout=int(timeout_s)), timeout_s, token) @property def name(self) -> str: @@ -59,12 +60,13 @@ async def session(self, role: str | None = None) -> AsyncIterator[ClientSession] async def invoke(self, uri: str, method: str, params: dict | None = None) -> Any: return await protocol.invoke_extension(self.url, self.card, uri, params or {}, method=method, - timeout=int(self.timeout_s)) + timeout=int(self.timeout_s), headers=self.headers) async def rpc(self, method: str, params: dict | None = None, card: dict | None = None) -> dict: """``method``'s answer from the data plane of ``card``, the driver's own or a child's it fronts.""" async with httpx.AsyncClient(timeout=self.timeout_s) as http: - response = await http.post(self.url + ((card or self.card).get("url") or "/agentenv"), + url = self.url + ((card or self.card).get("url") or "/agentenv") + response = await http.post(url, headers=protocol.request_headers(url, self.headers), json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params or {}}) return response.json() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py index 0b660db85..6177a4aac 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py @@ -12,7 +12,7 @@ from starlette.responses import Response from ..agent_env_environment import (_CARD_CONFIG_ATTR, AgentEnvEnvironment, AgentEnvFastMCPApplication, - _error_response, environment_card) + _error_response, control_guard, environment_card) from ..tool_context import ToolContext from ..types import EnvironmentCard, EnvironmentExtension, EnvironmentTool from . import _linux, apps, card @@ -135,7 +135,7 @@ async def children_environments(self) -> list[EnvironmentCard] | None: return await apps.children(self._computer.apps) if self._computer.apps else None def _computer_front_apps(self, app: Any) -> Any: - app.custom_route("/svc/{name}/{path:path}", methods=["GET", "POST"])(self._computer_forward) + app.custom_route("/svc/{name}/{path:path}", methods=["GET", "POST"])(control_guard()(self._computer_forward)) return app async def _computer_forward(self, request: Request) -> Response: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/serve_app.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/serve_app.py index 4ef274fa7..c3b2f52ab 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/serve_app.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/serve_app.py @@ -1,19 +1,22 @@ """Serve a computer app's class from its file, as its computer's program starts it: ``python -m agentenv_protocol.computer.serve_app :``, run in the app's folder, serves the class's card and -``computer-app/v1`` on the Unix socket ``.agentenv/server.sock`` there. Before it serves, it keeps a copy of the app's -data paths as installed, which a reset puts back; once it serves, it adds the app's launcher and launches a boot app -(``AppRuntime.start``). A server that can't start, or whose start fails once it serves, -writes why to ``.agentenv/server.failure``, which its computer's driver reports, and exits. Needs the ``app`` extra.""" +``computer-app/v1`` on the Unix socket ``.agentenv/server.sock`` there, which only the computer reaches, so it takes no +control token (``CONTROL_TOKEN_ENV``): the driver checks it in front of the app. Before it serves, it keeps a copy of +the app's data paths as installed, which a reset puts back; once it serves, it adds the app's launcher and launches a +boot app (``AppRuntime.start``). A server that can't start, or whose start fails once it serves, writes why to +``.agentenv/server.failure``, which its computer's driver reports, and exits. Needs the ``app`` extra.""" from __future__ import annotations import asyncio import contextlib +import os import sys import traceback from pathlib import Path import uvicorn +from ..types import CONTROL_TOKEN_ENV from ._load import load_class from .app import FAILURE, SOCKET, ComputerAppEnvironment @@ -35,6 +38,7 @@ def _serve(argv: list[str]) -> None: if len(argv) != 1 or ":" not in argv[0]: raise SystemExit(USAGE) file, _, name = argv[0].partition(":") + os.environ.pop(CONTROL_TOKEN_ENV, None) app = load_class(Path(file), name, ComputerAppEnvironment)() runtime = app._computer_app_runtime() app._computer_app_data_plane().keep_initial() diff --git a/packages/agentenv-protocol/src/agentenv_protocol/types.py b/packages/agentenv-protocol/src/agentenv_protocol/types.py index 448b430b6..0c91725df 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/types.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/types.py @@ -16,6 +16,9 @@ def error_body(code: str, message: str, **fields: Any) -> dict: WELL_KNOWN_PATH = "/.well-known/agent-env.json" RPC_PATH = "/agentenv" +#: The variable that hands an env the token its framework-only routes take as their bearer, everything it serves but +#: its MCP endpoint and its card; agent-env sets it when it deploys the env, and without it those routes are open. +CONTROL_TOKEN_ENV = "AGENTENV_CONTROL_TOKEN" # A card declares its MCP endpoint as an `additionalInterfaces` entry with this transport; a card # without one is reached at MCP_PATH by convention. diff --git a/packages/agentenv-protocol/tests/test_computer_app_server.py b/packages/agentenv-protocol/tests/test_computer_app_server.py index 231c281e4..8c535345a 100644 --- a/packages/agentenv-protocol/tests/test_computer_app_server.py +++ b/packages/agentenv-protocol/tests/test_computer_app_server.py @@ -1,6 +1,7 @@ """A computer app's own server: it adds the app's launcher, launches a boot app once per boot and says whether the app is ready (``computer-app/v1``); its computer's driver finds it on the socket in the app's folder and asks it, lists it -as a child environment and forwards its requests to it, and leaves the paths it owns out of its own data/get.""" +as a child environment and forwards its requests to it, only with the computer's control token when it has one, and +leaves the paths it owns out of its own data/get.""" from __future__ import annotations import asyncio @@ -31,7 +32,7 @@ from agentenv_protocol.computer import app, app_runtime, apps, session # noqa: E402 from agentenv_protocol.computer.app import COMPUTER_APP, FAILURE, SOCKET # noqa: E402 from agentenv_protocol.computer.platform import Platform # noqa: E402 -from agentenv_protocol.types import DATA_OBJECTS_EXTENSION_URI, WELL_KNOWN_PATH # noqa: E402 +from agentenv_protocol.types import CONTROL_TOKEN_ENV, DATA_OBJECTS_EXTENSION_URI, WELL_KNOWN_PATH # noqa: E402 @@ -311,6 +312,30 @@ async def test_the_driver_lists_its_apps_as_children_forwards_each_ones_requests "Documents/note.txt"] +@pytest.mark.asyncio +async def test_with_a_control_token_the_driver_lists_its_apps_and_forwards_to_one_only_with_it(tmp_path, monkeypatch): + from computer_fake import FakeDriver, http_client + + monkeypatch.setenv("HOME", str(tmp_path)) + (tmp_path / "Mail").mkdir() + (tmp_path / "Mail/inbox.txt").write_text("hi") + get = {"jsonrpc": "2.0", "id": 1, "method": "data/get", "params": {}} + with tempfile.TemporaryDirectory(dir="/tmp", prefix="apps-") as root: + monkeypatch.setattr(apps, "APPS", Path(root)) + async with _serving(Inbox(), Path(root) / "mail", monkeypatch): + monkeypatch.setenv(CONTROL_TOKEN_ENV, "s3cret") + async with http_client(FakeDriver()) as http: + card = (await http.get(WELL_KNOWN_PATH)).json() + refused = [await http.post("/svc/mail/agentenv", json=get, headers=headers) + for headers in ({}, client.control_headers("guess"))] + forwarded = await http.post("/svc/mail/agentenv", json=get, headers=client.control_headers("s3cret")) + + assert [child["name"] for child in card["children_environments"]] == ["mail"] + assert [answer.status_code for answer in refused] == [401, 401] + assert [part["metadata"][DATA_OBJECTS_EXTENSION_URI]["path"] for part in forwarded.json()["result"]["parts"]] == [ + "Mail/inbox.txt"] + + @pytest.mark.asyncio async def test_the_session_launches_an_app_in_a_session_of_its_own_from_its_folder(tmp_path, monkeypatch): monkeypatch.setenv("DISPLAY", ":99") diff --git a/packages/agentenv-protocol/tests/test_computer_conformance.py b/packages/agentenv-protocol/tests/test_computer_conformance.py index 154738a95..3f2f8ddd6 100644 --- a/packages/agentenv-protocol/tests/test_computer_conformance.py +++ b/packages/agentenv-protocol/tests/test_computer_conformance.py @@ -1,6 +1,7 @@ """The conformance probe against fake drivers on loopback: the fake driver with the fake fixture passes every check; a -driver without what full needs gets those checks unverified; a driver broken one way fails the check for that way; each -app a driver fronts that exports is asked for its export once; the report's verdicts and the CLI.""" +driver given a control token passes with it, refusing calls without it, and fails without it; a driver without what +full needs gets those checks unverified; a driver broken one way fails the check for that way; each app a driver +fronts that exports is asked for its export once; the report's verdicts and the CLI.""" from __future__ import annotations import json @@ -33,6 +34,7 @@ from agentenv_protocol.computer.conformance import checks # noqa: E402 from agentenv_protocol.computer.conformance.client import Driver # noqa: E402 from agentenv_protocol.computer.conformance.__main__ import main # noqa: E402 +from agentenv_protocol import CONTROL_TOKEN_ENV # noqa: E402 from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI # noqa: E402 from agentenv_protocol.types import EnvironmentCard # noqa: E402 @@ -71,6 +73,18 @@ async def test_the_fake_driver_passes_every_check_at_full_and_cleans_up(home): assert ["kill", "4242"] in driver.commands +@pytest.mark.asyncio +async def test_a_driver_given_a_control_token_passes_with_it_and_refuses_calls_without_it(monkeypatch): + monkeypatch.setenv(CONTROL_TOKEN_ENV, "s3cret") + + report = await _probe(FixtureDriver(), "basic", token="s3cret") + unauthorized = await _probe(FixtureDriver(), "basic") + + assert report.passed, _failed(report) + assert _status(report, "a framework-only route refuses a call without the control token") == "passed" + assert not unauthorized.passed and "401 unauthorized" in _failed(unauthorized)[0] + + @pytest.mark.asyncio async def test_a_driver_whose_default_space_is_norm_passes_at_full(): report = await _probe(FixtureDriver(settings=ComputerSettings(default_space="norm"))) diff --git a/packages/agentenv-protocol/tests/test_computer_serve.py b/packages/agentenv-protocol/tests/test_computer_serve.py index 84aa77113..775d48d26 100644 --- a/packages/agentenv-protocol/tests/test_computer_serve.py +++ b/packages/agentenv-protocol/tests/test_computer_serve.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import os import sys import pytest @@ -11,6 +12,7 @@ pytest.importorskip("PIL") pytest.importorskip("uvicorn") +from agentenv_protocol import CONTROL_TOKEN_ENV # noqa: E402 from agentenv_protocol.computer import ComputerDriverEnvironment, run_install_step, serve_app # noqa: E402 from agentenv_protocol.computer.app import BUILD_SETTINGS, FAILURE, SOCKET # noqa: E402 from agentenv_protocol.computer.app_runtime import AppRuntime # noqa: E402 @@ -164,10 +166,11 @@ async def lifespan(): return tmp_path, ran -def test_an_apps_server_serves_on_the_socket_in_its_folder_and_starts_the_app_once_it_serves(app_folder, - monkeypatch): +def test_an_apps_server_serves_on_the_socket_in_its_folder_with_no_control_token_and_starts_the_app_once_it_serves( + app_folder, monkeypatch): folder, ran = app_folder started = [] + monkeypatch.setenv(CONTROL_TOKEN_ENV, "s3cret") async def start(runtime): started.append(runtime) @@ -177,6 +180,7 @@ async def start(runtime): serve_app.main(["app.py:Notes"]) (server,) = ran + assert CONTROL_TOKEN_ENV not in os.environ assert server.config.uds == str(SOCKET) and not server.should_exit assert not (folder / SOCKET).exists() and not (folder / FAILURE).exists() assert len(started) == 1 and isinstance(started[0], AppRuntime) diff --git a/packages/agentenv-protocol/tests/test_control_token.py b/packages/agentenv-protocol/tests/test_control_token.py new file mode 100644 index 000000000..a083ea4e1 --- /dev/null +++ b/packages/agentenv-protocol/tests/test_control_token.py @@ -0,0 +1,148 @@ +"""The control token: with CONTROL_TOKEN_ENV set when an env mounts, everything it serves but its MCP endpoint and its +card refuses a call without that token, or with another, before the route runs, and answers one with it; the MCP +endpoint and the card stay open, and without the variable every route is open. The client sends a call's own headers, +or else the headers a grant gives the address it calls.""" +from __future__ import annotations + +import contextlib + +import anyio +import httpx +import pytest + +pytest.importorskip("mcp.server.fastmcp") + +from mcp import ClientSession # noqa: E402 +from mcp.client import streamable_http as _streamable_http # noqa: E402 +from mcp.server.fastmcp import FastMCP # noqa: E402 + +from agentenv_protocol import ( # noqa: E402 + CONTROL_TOKEN_ENV, + RPC_PATH, + WELL_KNOWN_PATH, + AgentEnvEnvironment, + add_data, + client, + environment_card, + extension, + tool, +) + +pytestmark = pytest.mark.skipif(not hasattr(_streamable_http, "streamable_http_client"), + reason="mcp without streamable_http_client") + +_BASE = "http://127.0.0.1:8000" +RAN: list[str] = [] + + +@environment_card(name="desk") +class _Desk(AgentEnvEnvironment): + tool_access = True + + @tool() + def read(self) -> str: + return "read" + + @add_data + async def load(self, parts: list) -> None: + RAN.append("data/add") + + @extension("urn:agentenv:test-exec/v1", op="run") + async def run(self, argv: list) -> dict: + RAN.append("run") + return {"ran": argv} + + +@contextlib.asynccontextmanager +async def _served(): + """An ``_Desk`` mounted on mcp's FastMCP as the environment's variables stand, and its REST and MCP clients.""" + app = FastMCP("desk", json_response=True) + _Desk().mount(app) + asgi = app.streamable_http_app() + async with app.session_manager.run(): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=asgi), base_url=_BASE) as rest: + async with _streamable_http.streamable_http_client(f"{_BASE}/mcp", http_client=rest) as (read, write, _): + async with ClientSession(read, write) as mcp: + await mcp.initialize() + yield rest, mcp + + +async def _calls(rest: httpx.AsyncClient, headers: dict) -> list[int]: + """The status of a call to each framework-only route with ``headers``: the data plane, an extension, tool access.""" + load = {"jsonrpc": "2.0", "id": 1, "method": "data/add", "params": {"parts": [{"kind": "data", "data": {}}]}} + return [(await rest.post(RPC_PATH, json=load, headers=headers)).status_code, + (await rest.post("/agentenv/ext/test-exec/run", json={"argv": ["ls"]}, headers=headers)).status_code, + (await rest.post("/tools/disable", json={"role": "agent", "tools": ["read"]}, headers=headers)).status_code] + + +@pytest.mark.asyncio +async def test_with_a_token_every_framework_only_route_refuses_a_call_without_it_before_it_runs(monkeypatch): + monkeypatch.setenv(CONTROL_TOKEN_ENV, "s3cret") + RAN.clear() + + with anyio.fail_after(30): + async with _served() as (rest, mcp): + without = await _calls(rest, {}) + wrong = await _calls(rest, client.control_headers("guess")) + refusal = (await rest.post("/agentenv/ext/test-exec/run", json={"argv": ["ls"]})).json() + ran_before = list(RAN) + right = await _calls(rest, client.control_headers("s3cret")) + card = await rest.get(WELL_KNOWN_PATH) + tools = [t.name for t in (await mcp.list_tools()).tools] + + assert without == wrong == [401, 401, 401] and ran_before == [] + assert refusal["error"]["code"] == "unauthorized" and refusal["error"]["applied"] == "no" + assert right == [200, 200, 200] and RAN == ["data/add", "run"] + assert card.status_code == 200 and tools == ["read"] + + +@pytest.mark.asyncio +async def test_without_a_token_every_route_is_open(monkeypatch): + monkeypatch.delenv(CONTROL_TOKEN_ENV, raising=False) + + with anyio.fail_after(30): + async with _served() as (rest, _): + assert await _calls(rest, {}) == [200, 200, 200] + + +def test_a_call_sends_its_own_headers_or_else_the_ones_its_grant_gives_an_address_under_the_envs(): + env = "http://desk:18765" + + with client.authorized({env: client.control_headers("s3cret"), "http://other:1": {}, + None: client.control_headers("orphan")}): + granted = [client.request_headers(url) for url in (env, f"{env}/svc/mail/agentenv", f"{env}0/x", + "http://other:1/x")] + own = client.request_headers(f"{env}/agentenv", {"X-Mine": "1"}) + with client.authorized({env: client.control_headers("rotated")}): + inner = client.request_headers(f"{env}/agentenv") + outer = client.request_headers(f"{env}/agentenv") + after = client.request_headers(f"{env}/agentenv") + + bearer = {"Authorization": "Bearer s3cret"} + assert granted == [bearer, bearer, {}, {}] and own == {"X-Mine": "1"} + assert inner == {"Authorization": "Bearer rotated"} and outer == bearer and after == {} + assert client.control_headers(None) == {} and client.control_headers("") == {} + + +@pytest.mark.asyncio +async def test_invoke_extension_and_the_data_plane_send_the_token(monkeypatch): + seen = [] + + def handle(request: httpx.Request) -> httpx.Response: + seen.append((request.url.path, request.headers.get("authorization"))) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {"parts": []}}) + + real = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + card = {"capabilities": {"extensions": [{"uri": "urn:agentenv:test-exec/v1", "params": { + "methods": {"run": {"endpoint": "/agentenv/ext/test-exec/run"}}}}]}} + + await client.invoke_extension("http://desk", card, "urn:agentenv:test-exec/v1", {}, method="run", + headers=client.control_headers("own")) + await client.get_data("http://desk", headers=client.control_headers("own")) + with client.authorized({"http://desk": client.control_headers("granted")}): + await client.add_data("http://desk", []) + await client.reset_data("http://desk", headers={}) + + assert seen == [("/agentenv/ext/test-exec/run", "Bearer own"), ("/agentenv", "Bearer own"), + ("/agentenv", "Bearer granted"), ("/agentenv", None)] diff --git a/src/agent_env/env/env.py b/src/agent_env/env/env.py index 3bdc55afc..411cceb91 100644 --- a/src/agent_env/env/env.py +++ b/src/agent_env/env/env.py @@ -93,6 +93,8 @@ class DeployedEnv: instance_id: Optional[str] = None created_at_utc: Optional[str] = None expires_at_utc: Optional[str] = None + # The token the env's framework-only routes take, from the provider that booted it with it; None when they're open. + control_token: Optional[str] = None def __post_init__(self) -> None: # Only older records lack their card: they keep the stored values. @@ -110,6 +112,16 @@ def environment_url(self) -> Optional[str]: """The env's address: its card URL without the well-known path; card paths join onto it.""" return _environment_url(self.environment_card_url) if self.environment_card_url else None + @property + def control_headers(self) -> dict[str, str]: + """What a call to the env's framework-only routes sends: its control token, when it has one.""" + return protocol_v1.control_headers(self.control_token) + + 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 supports(self, uri: str, method: str) -> bool: """Whether the stored env card offers `method` on the extension at `uri`; no network.""" return protocol_v1.find_extension_method(self.environment_card or {}, uri, method) is not None @@ -135,7 +147,8 @@ async def invoke(self, uri: str, method: str, params: Optional[dict] = None, *, card = self.environment_card if environment_name is None else self.get_child_env_card(environment_name) if card is None or not self.environment_url or protocol_v1.find_extension_method(card, uri, method) is None: raise EnvCapabilityUnsupported(uri, method, self.env_id, self.instance_id, environment_name) - return await protocol_v1.invoke_extension(self.environment_url, card, uri, params, timeout, method=method) + return await protocol_v1.invoke_extension(self.environment_url, card, uri, params, timeout, method=method, + headers=self.control_headers) @classmethod def _fields_from(cls, data: dict) -> dict: @@ -152,6 +165,7 @@ def _fields_from(cls, data: dict) -> dict: instance_id=data.get("instance_id"), created_at_utc=data.get("created_at_utc"), expires_at_utc=data.get("expires_at_utc"), + control_token=data.get("control_token"), ) diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index 70b4ea55c..daa8fd269 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -11,6 +11,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Optional +from agentenv_protocol import CONTROL_TOKEN_ENV from agentenv_protocol import client as protocol_v1 from agentenv_protocol.computer import APP_SETTINGS_ENV, COMPUTER_READY, SETTINGS_ENV from agentenv_protocol.computer.apps import ANSWER_TIMEOUT_S @@ -35,6 +36,8 @@ _POLL_S = 0.5 #: How a setting a record keeps shows a secret's or a generated one's value, which it never holds. WITHHELD = "" +#: How many random bytes a computer's control token carries: 256 bits, past guessing. +CONTROL_TOKEN_BYTES = 32 @dataclass @@ -87,19 +90,31 @@ def recorded_settings(computer: ComputerArtifact, values: dict[str, Any]) -> dic for name, setting in computer.config.items() if setting.phase == "deploy"} -async def boot(sandbox_provider: SandboxProvider, *, image: str, port: int, environment: dict[str, str], +def control_token() -> str: + """A token for one computer's framework-only routes (``CONTROL_TOKEN_ENV``), which agent-env holds and gives no + agent.""" + return secrets.token_urlsafe(CONTROL_TOKEN_BYTES) + + +def control_environment(token: Optional[str]) -> dict[str, str]: + """What hands a computer its control ``token`` in its boot environment; nothing without one, as for a computer an + older agent-env booted.""" + return {CONTROL_TOKEN_ENV: token} if token else {} + + +async def boot(sandbox_provider: SandboxProvider, *, image: str, port: int, environment: dict[str, str], token: str, network_policy: Optional[NetworkPolicy], timeout: int, attribution: Optional[Attribution], disk_size_gb: float = 10, cpu: Optional[float] = None, memory_mb: Optional[int] = None, ready_s: float = READY_TIMEOUT_S) -> Booted: - """``image`` booted with ``environment`` on ``sandbox_provider``, which the caller chose to run its arch natively, - its driver's ``port`` reached on the sandbox, once the driver says the computer is ready; a boot that fails is torn - down.""" + """``image`` booted with ``environment`` and its control ``token`` on ``sandbox_provider``, which the caller chose + to run its arch natively, its driver's ``port`` reached on the sandbox, once the driver, asked with that token, says + the computer is ready; a boot that fails is torn down.""" sandbox = await sandbox_provider.create_container( - image_name=image, port=port, env=environment, network_policy=network_policy, + image_name=image, port=port, env={**environment, **control_environment(token)}, network_policy=network_policy, disk_size_gb=disk_size_gb, timeout=timeout, attribution=dict(attribution or {}), **_size_kwargs(cpu, memory_mb)) try: url = sandbox.tunnel_urls[port] - return Booted(sandbox, url, await ready(url, ready_s)) + return Booted(sandbox, url, await ready(url, ready_s, token)) except BaseException: await discard(sandbox) raise @@ -124,10 +139,10 @@ def ready_timeout(apps: list[ComputerAppEnv], os_type: str) -> float: if (launch := app.launch(os_type)) and launch.at == "boot" and launch.ready)]) -async def ready(url: str, timeout: float) -> dict: - """The card of the driver at ``url``, read once its ``computer-ready`` status says it's ready, so each app it fronts - is on it as the app's own server serves it; raises TimeoutError, naming what it last saw, when it isn't within - ``timeout`` seconds.""" +async def ready(url: str, timeout: float, token: Optional[str] = None) -> dict: + """The card of the driver at ``url``, read once its ``computer-ready`` status, asked with its control ``token``, + says it's ready, so each app it fronts is on it as the app's own server serves it; raises TimeoutError, naming + what it last saw, when it isn't within ``timeout`` seconds.""" last = "nothing answered" try: async with asyncio.timeout(timeout): @@ -135,7 +150,8 @@ async def ready(url: str, timeout: float) -> dict: try: card = await protocol_v1.get_card(url, timeout=5) status = await protocol_v1.invoke_extension(url, card, COMPUTER_READY, None, _STATUS_TIMEOUT_S, - method="status") + method="status", + headers=protocol_v1.control_headers(token)) if status.get("ready"): return await protocol_v1.get_card(url, timeout=5) last = "not ready: " + "; ".join(f"{check['name']}: {check.get('detail')}" diff --git a/src/agent_env/env/envs/computer/utils/composer.py b/src/agent_env/env/envs/computer/utils/composer.py index dd3ce211f..195e2c9d6 100644 --- a/src/agent_env/env/envs/computer/utils/composer.py +++ b/src/agent_env/env/envs/computer/utils/composer.py @@ -182,13 +182,15 @@ def _probed(env: ComputerEnv, driver: ComputerDriverEnv, form: str, arch: str, r async def _probe(links: list[SandboxProvider], env: ComputerEnv, driver: ComputerDriverEnv, image: str) -> tuple[ProbeReport, dict, str]: - """``image`` booted at the env's settings, as a deploy boots it, on the first of ``links`` that boots it, probed at - ``full`` within ``_PROBE_TIMEOUT_S`` and torn down; the report, the driver's card and the provider's name.""" + """``image`` booted at the env's settings with a control token of its own, as a deploy boots it, on the first of + ``links`` that boots it, probed at ``full`` with that token within ``_PROBE_TIMEOUT_S`` and torn down; the report, + the driver's card and the provider's name.""" + token = boot.control_token() environment = boot.settings_environment(env.computer, _values(env), environment_name=env.environment_name) problems = [] for link in links: try: - booted = await boot.boot(link, image=image, port=driver.port, environment=environment, + booted = await boot.boot(link, image=image, port=driver.port, environment=environment, token=token, network_policy=None, timeout=_PROBE_TTL_S, attribution=None, ready_s=boot.ready_timeout(env.composed_apps(), env.computer.os.type)) except Exception as e: @@ -196,7 +198,7 @@ async def _probe(links: list[SandboxProvider], env: ComputerEnv, driver: Compute continue try: async with asyncio.timeout(_PROBE_TIMEOUT_S): - return await run_probe(booted.url, "full"), booted.card, type(link).__name__ + return await run_probe(booted.url, "full", token=token), booted.card, type(link).__name__ finally: await boot.discard(booted.sandbox) raise RuntimeError("; ".join(problems)) diff --git a/src/agent_env/env/envs/computer/utils/inside.py b/src/agent_env/env/envs/computer/utils/inside.py index 7be9f394f..c03ccecaa 100644 --- a/src/agent_env/env/envs/computer/utils/inside.py +++ b/src/agent_env/env/envs/computer/utils/inside.py @@ -47,9 +47,11 @@ def max_timeout_s(deployed: DeployedSandboxEnv) -> float: def loopback(deployed: DeployedSandboxEnv) -> dict: - """The env's id, address and MCP endpoint as code inside the computer reaches them: the driver on loopback.""" + """The env's id, address and MCP endpoint as code inside the computer reaches them, the driver on loopback, and + the ``headers`` its framework-only routes take: code agent-env runs in the computer gets the computer's token.""" url = f"http://127.0.0.1:{deployed.metadata['computer']['port']}" - return {"id": deployed.env_id, "url": url, "mcp_url": url + deployed.mcp_url.removeprefix(deployed.environment_url)} + return {"id": deployed.env_id, "url": url, "mcp_url": url + deployed.mcp_url.removeprefix(deployed.environment_url), + "headers": deployed.control_headers} async def add_artifacts(deployed: DeployedSandboxEnv, files: dict[str, FileArtifact], @@ -79,7 +81,8 @@ async def add_files(deployed: DeployedSandboxEnv, files: dict[str, bytes]) -> No async def read_file(deployed: DeployedSandboxEnv, path: str) -> bytes: """The bytes of the file at ``path``; FileNotFoundError when there is none.""" try: - answer = await protocol_v1.get_data(deployed.environment_url, timeout=_READ_TIMEOUT_S, paths=[path]) + answer = await protocol_v1.get_data(deployed.environment_url, timeout=_READ_TIMEOUT_S, paths=[path], + headers=deployed.control_headers) except httpx.TransportError as e: raise _broken(deployed, "data/get", e) from e for part in answer.parts: @@ -102,7 +105,8 @@ async def run(deployed: DeployedSandboxEnv, argv: list[str], *, timeout_s: float async def _add(deployed: DeployedSandboxEnv, parts: list[FilePart], size: int, url: Optional[str] = None) -> None: try: - await protocol_v1.add_data(url or deployed.environment_url, parts, timeout=data_plane_load_timeout_s(size)) + await protocol_v1.add_data(url or deployed.environment_url, parts, timeout=data_plane_load_timeout_s(size), + headers=deployed.control_headers) except httpx.TransportError as e: raise _broken(deployed, "data/add, so some of the files it carried may be written, each whole", e) from e diff --git a/src/agent_env/env/legacy_protocol.py b/src/agent_env/env/legacy_protocol.py index e9533a844..16899025b 100644 --- a/src/agent_env/env/legacy_protocol.py +++ b/src/agent_env/env/legacy_protocol.py @@ -96,7 +96,8 @@ async def _export_state_at(base_url: str, timeout: int, verify: bool) -> dict: async def service_state(deployed: Optional[DeployedEnv], gateway_url: Optional[str], environment_name: str) -> dict: - """An MCP service's state as JSON: its v1 ``data/get`` answer when that is data, else ``GET /export-state``. + """An MCP service's state as JSON: its v1 ``data/get`` answer when that is data, asked with the deployment's + control token, else ``GET /export-state``. A service that exports a file from ``data/get`` (a bundle of its database, to reload from) still serves its state as JSON at ``/export-state``. @@ -104,7 +105,8 @@ async def service_state(deployed: Optional[DeployedEnv], gateway_url: Optional[s base_url = await v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is None: return await export_state(gateway_url, environment_name) - response = await protocol_v1.get_data(base_url) + response = await protocol_v1.get_data(base_url, + headers=deployed.control_headers_for(base_url) if deployed is not None else None) if not response.parts: return {} if isinstance(response.parts[0], DataPart): diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index e4903e743..3db560b1b 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -17,6 +17,8 @@ from agent_env.env.envs.computer.utils.boot import ( app_settings_environment, boot, + control_environment, + control_token, discard, ready, ready_timeout, @@ -90,11 +92,11 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par raise ValueError(f"{type(sandbox_provider).__name__} runs none of {env.id}'s ready builds natively " f"({', '.join(f'{form}/{arch}' for form, arch, _ in builds)})") form, arch, build = native[0] - computer, name = env.computer_artifact(), env.environment_name + computer, name, token = env.computer_artifact(), env.environment_name, control_token() values = _values(computer, values) policy = NetworkPolicy.from_dict(env.network) if env.network else None booted = await boot( - sandbox_provider, image=build["image"], port=env.driver.port, network_policy=policy, + sandbox_provider, image=build["image"], port=env.driver.port, network_policy=policy, token=token, environment=_environment(env, computer, values, params), timeout=ttl_seconds, attribution=attribution, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, ready_s=ready_timeout(env.composed_apps(), env.computer.os.type)) @@ -107,7 +109,7 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par environment_card_url=f"{booted.url}{WELL_KNOWN_PATH}", environment_card=booted.card, environment_card_read_at_utc=datetime.now(timezone.utc).isoformat(), sandbox_id=booted.sandbox.sandbox_id, sandbox_type=booted.sandbox.type, - sandbox_ids={ComputerEnv.type: {name: booted.sandbox.sandbox_id}}, + sandbox_ids={ComputerEnv.type: {name: booted.sandbox.sandbox_id}}, control_token=token, metadata={"computer": { "form": form, "arch": arch, "provider": type(sandbox_provider).__name__, "image": build["image"], "computer": {"id": env.metadata["computer"]["id"], "version": env.metadata["computer"]["version"]}, @@ -123,8 +125,8 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: """Replace the computer ``deployed`` records with a fresh container of the image it booted, once its recordings - ended, its settings applied as that deploy applied them (a generated one minted afresh), on the same sandbox and - port, readied as a deploy readies it.""" + ended, its settings applied as that deploy applied them (a generated one minted afresh) and its control token + the record's, on the same sandbox and port, readied as a deploy readies it.""" record = deployed.metadata["computer"] params, values, _ = computer_params.checked(env, record["env_params"]) computer = ComputerArtifact.get(record["computer"]["id"], record["computer"]["version"]) @@ -134,10 +136,12 @@ async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: sandbox = await sandbox_provider.get_sandbox(deployed.sandbox_id) try: await sandbox_provider.reset_container(sandbox, image_name=record["image"], port=env.driver.port, - env=_environment(env, computer, values, params)) + env={**_environment(env, computer, values, params), + **control_environment(deployed.control_token)}) except NotImplementedError as e: # a reset the env supports, which this provider can't carry out raise RuntimeError(f"'{env.id}' can't be reset on {type(sandbox_provider).__name__}: {e}") from None - card = await ready(deployed.environment_url, ready_timeout(env.composed_apps(), env.computer.os.type)) + card = await ready(deployed.environment_url, ready_timeout(env.composed_apps(), env.computer.os.type), + deployed.control_token) _require_name(env, card) await _require_display(deployed, values.get("resolution")) diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 4c2beb166..6840a58d8 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -17,6 +17,8 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Callable, ClassVar, Optional +from agentenv_protocol import client as protocol_v1 + from agent_env.config import run_env_ids from agent_env.plugins import _registration from agent_env.plugins._requirements import unmet_problems @@ -629,7 +631,8 @@ async def _run_one(step: TaskStep, idx: int, dispatch_epoch: int) -> None: started_mono = time.monotonic() started_utc = datetime.now(timezone.utc).isoformat() try: - await step.execute(context) + with protocol_v1.authorized({d.environment_url: d.control_headers for d in context.deployed_envs}): + await step.execute(context) except asyncio.CancelledError: raise except BaseException as exc: diff --git a/src/agent_env/task_step/task_steps/apply_server_config.py b/src/agent_env/task_step/task_steps/apply_server_config.py index d1495fb9b..41f82ffd0 100644 --- a/src/agent_env/task_step/task_steps/apply_server_config.py +++ b/src/agent_env/task_step/task_steps/apply_server_config.py @@ -255,7 +255,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: ) result = await protocol_v1.invoke_extension( base_url, card, directive.uri, params=directive.args, - timeout=self.timeout_seconds, + timeout=self.timeout_seconds, headers=deployed.control_headers, ) except ApplyServerConfigError: raise diff --git a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py index 16adf5090..20d25ea43 100644 --- a/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py +++ b/src/agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py @@ -65,7 +65,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: protocol = {} for method, probe_params in _CORE_OPERATIONS: if probe_params is not None: - supported = await self._probe_operation(deployed.environment_url, method, probe_params) + supported = await self._probe_operation(deployed.environment_url, method, probe_params, + deployed.control_headers) else: supported = self._advertised_support(card, method) protocol[method] = {"supported": supported} @@ -90,19 +91,20 @@ def _advertised_support(card: Optional[dict], method: str) -> Optional[bool]: advertised = protocol_v1.operations(backing) return True if advertised is None else method in advertised - async def _probe_operation(self, environment_url: str, method: str, params: dict) -> Optional[bool]: - """Mutation-free: -32601 = unregistered; -32000 = gateway couldn't forward (None); + async def _probe_operation(self, environment_url: str, method: str, params: dict, + headers: dict[str, str]) -> Optional[bool]: + """Mutation-free: -32601 = unregistered; -32000 = gateway couldn't forward, or a 401 = not let in (None); anything else — a result, or -32602 from the empty-parts add probe — proves registration.""" body = {"jsonrpc": "2.0", "id": f"verify-{method}", "method": method, "params": params} try: async with httpx.AsyncClient(timeout=15) as client: - response = await client.post(f"{environment_url}{RPC_PATH}", json=body) + response = await client.post(f"{environment_url}{RPC_PATH}", json=body, headers=headers) error = (response.json() or {}).get("error") or {} except Exception as e: logger.warning(f"operation probe {method} for '{self.env_id}' failed: {e}") return None if error.get("code") == -32601: return False - if error.get("code") == -32000: + if response.status_code == 401 or error.get("code") == -32000: return None return True 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 4a04ab10d..0afb54b71 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 @@ -152,7 +152,8 @@ 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}, timeout=_TOOL_CALL_TIMEOUT_S) + headers={AGENT_ENV_ROLE_HEADER: 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}): " f"{resp.text[:300]}") diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 98b773526..e1ae345bc 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -5,7 +5,9 @@ A script run for an env also gets ``input["env"] = {"id", "url", "mcp_url"}``: the env's address, where its card is at ``/.well-known/agent-env.json``, and its MCP endpoint, as the script reaches them: on the env's host as that host does (``Sandbox.url_on_host``), since a host can't -always call its own public URL; inside a computer, the driver on loopback. +always call its own public URL; inside a computer, the driver on loopback, with the ``headers`` that +carry the computer's control token to its framework-only routes. A script run in an agent's sandbox +never gets the token. ``script_artifact_id`` is either a single-file ``FileArtifact`` or a ``FileArtifactUniverse`` whose ``.py`` members are staged side by side so the entry diff --git a/src/agent_env/task_step/task_steps/snapshot_env.py b/src/agent_env/task_step/task_steps/snapshot_env.py index 4b082f3ef..48fc9d6dc 100644 --- a/src/agent_env/task_step/task_steps/snapshot_env.py +++ b/src/agent_env/task_step/task_steps/snapshot_env.py @@ -99,7 +99,8 @@ def _snapshot_upload(card: dict, timeout_seconds: float, sandbox_type: Optional[ return _SnapshotUpload(store, grant) -async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float) -> None: +async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float, + headers: Optional[dict[str, str]]) -> None: """Best-effort: push the object store's shared AWS creds + bucket to a service advertising add_s3_credentials, if the store shares any. Never raises (get_data falls back). An @local run's snapshots land in the local stores, so it hands a service no credentials to @@ -132,6 +133,7 @@ async def _push_s3_credentials(base_url: str, card: dict, timeout_seconds: float _S3_CREDENTIALS_EXTENSION_URI, params, timeout=int(timeout_seconds), + headers=headers, ) logger.info("snapshot_env: pushed S3 creds to %s", base_url) except Exception: @@ -377,13 +379,14 @@ def _resolve_deployed_env(self, context: TaskStepContext): async def _capture_env_trajectory(self, entry: dict, context: TaskStepContext, env_id: str, deployed) -> None: """Stream the trajectory to a temp file, upload it verbatim (no envelope), and fill ``entry``.""" method, url = _trajectory_route(deployed, self.gateway_url) + headers = deployed.control_headers_for(url) if deployed is not None else None store = get_config().get_object_store() event_count = 0 timeout = httpx.Timeout(_TRAJECTORY_CAPTURE_BUDGET_SECONDS, connect=_CONNECT_TIMEOUT_SECONDS) tmp = tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) try: async with httpx.AsyncClient(timeout=timeout) as client: - async with client.stream(method, url) as response: + async with client.stream(method, url, headers=headers) as response: response.raise_for_status() async for chunk in response.aiter_bytes(): tmp.write(chunk) @@ -404,6 +407,7 @@ async def _capture_env_trajectory(self, entry: dict, context: TaskStepContext, e async def _capture_env_changes(self, entry: dict, context: TaskStepContext, env_id: str, deployed) -> None: """Page the gateway's changes into a JSONL upload; fill ``entry`` with its url and attribution counts.""" method, url = _changes_route(deployed, self.gateway_url) + headers = deployed.control_headers_for(url) if deployed is not None else None counts = {"change_count": 0, "attributed_count": 0, "ambiguous_count": 0} timeout = httpx.Timeout(_CHANGES_CAPTURE_BUDGET_SECONDS, connect=_CONNECT_TIMEOUT_SECONDS) tmp = tempfile.NamedTemporaryFile(suffix=".jsonl", delete=False) @@ -411,7 +415,8 @@ async def _capture_env_changes(self, entry: dict, context: TaskStepContext, env_ after_id, has_more = 0, True async with httpx.AsyncClient(timeout=timeout) as client: while has_more: - response = await client.request(method, url, params={"after_id": after_id, "limit": CHANGES_PAGE_MAX}) + response = await client.request(method, url, params={"after_id": after_id, "limit": CHANGES_PAGE_MAX}, + headers=headers) response.raise_for_status() page = response.json() for change in page["changes"]: @@ -496,18 +501,21 @@ async def _export_environment_to_file( base_url = await legacy_protocol.v1_base_url(deployed, gateway_url, environment_name, mcp=True) if base_url is not None: + headers = deployed.control_headers_for(base_url) if deployed is not None else None card_base, card = await legacy_protocol.child_env_card(deployed, gateway_url, environment_name) if returns_file_tree(card or {}): upload = await asyncio.to_thread(_snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None), FILE_TREE_LIMITS) - return await get_data(base_url, environment_name, timeout_seconds, upload, **(data_get_params or {})) + return await get_data(base_url, environment_name, timeout_seconds, upload, headers=headers, + **(data_get_params or {})) upload = await asyncio.to_thread( _snapshot_upload, card or {}, timeout_seconds, getattr(deployed, "sandbox_type", None) ) # Push S3 creds (no-op unless the service advertises the extension), with a grant too: a # bundle the grant cannot hold can still be uploaded with them. - await _push_s3_credentials(card_base, card or {}, timeout_seconds) - resp, upload = await get_data(base_url, environment_name, timeout_seconds, upload, **(data_get_params or {})) + await _push_s3_credentials(card_base, card or {}, timeout_seconds, headers) + resp, upload = await get_data(base_url, environment_name, timeout_seconds, upload, headers=headers, + **(data_get_params or {})) part = resp.parts[0] if resp.parts else None if upload is not None and (path := uploaded_object_path(part)) is not None: return upload.object_url(path) diff --git a/src/agent_env/task_step/task_steps/sync_env_clock.py b/src/agent_env/task_step/task_steps/sync_env_clock.py index 1dc48a55d..ede3a409b 100644 --- a/src/agent_env/task_step/task_steps/sync_env_clock.py +++ b/src/agent_env/task_step/task_steps/sync_env_clock.py @@ -128,5 +128,5 @@ async def _sync_one(self, deployed, gateway_url: str, service: str, env_get_time raise RuntimeError(f"gateway did not expose env_get_time_url; cannot sync {service}") result = await protocol_v1.invoke_extension(base_url, card, EXT_CLOCK_URI, params={"env_get_time_url": env_get_time_url}, - timeout=self.timeout_seconds) + timeout=self.timeout_seconds, headers=deployed.control_headers) return {"synced": True, "result": result} diff --git a/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py b/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py index 1e8c89d1b..e9dfa989f 100644 --- a/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py @@ -25,7 +25,9 @@ class EnvOutcomeVerifierTaskStep(TaskStep): """Runs the artifact's ``verify(mcp_url)``; a ``run`` parameter also gets the ``Run`` of ``prompt_id``, and an - ``env_url`` parameter the env's address, which the protocol client takes as it is.""" + ``env_url`` parameter the env's address, which the protocol client takes as it is: like every step, ``verify`` + runs under ``Task.run``'s grant of each deployed env's control token, so its calls reach the env's framework-only + routes as written.""" type: ClassVar[str] = "env_outcome_verifier" entity_refs = ( diff --git a/tst/integration/computer/app/artifacts/check_desk/run.py b/tst/integration/computer/app/artifacts/check_desk/run.py index 38db8dd25..20c7c0cfd 100644 --- a/tst/integration/computer/app/artifacts/check_desk/run.py +++ b/tst/integration/computer/app/artifacts/check_desk/run.py @@ -18,7 +18,9 @@ def run(input): time.sleep(0.5) with urllib.request.urlopen(f"{input['env']['url']}/.well-known/agent-env.json", timeout=10) as response: card = json.load(response) - with urllib.request.urlopen(f"{input['env']['url']}/agentenv/ext/computer-ready/status", timeout=30) as response: + status = urllib.request.Request(f"{input['env']['url']}/agentenv/ext/computer-ready/status", + headers=input["env"]["headers"]) + with urllib.request.urlopen(status, timeout=30) as response: ready = json.load(response) with urllib.request.urlopen("http://127.0.0.1:8137/", timeout=10) as response: clock = json.load(response) diff --git a/tst/integration/computer/app/artifacts/check_notes/run.py b/tst/integration/computer/app/artifacts/check_notes/run.py index e91675f11..50e145ae3 100644 --- a/tst/integration/computer/app/artifacts/check_notes/run.py +++ b/tst/integration/computer/app/artifacts/check_notes/run.py @@ -18,7 +18,9 @@ def run(input): time.sleep(0.5) with urllib.request.urlopen(f"{input['env']['url']}/.well-known/agent-env.json", timeout=10) as response: card = json.load(response) - with urllib.request.urlopen(f"{input['env']['url']}/agentenv/ext/computer-ready/status", timeout=30) as response: + status = urllib.request.Request(f"{input['env']['url']}/agentenv/ext/computer-ready/status", + headers=input["env"]["headers"]) + with urllib.request.urlopen(status, timeout=30) as response: ready = json.load(response) return {**json.loads(written.read_text()), "card": card["name"], "app_ready": [check for check in ready["checks"] if check["name"].startswith("app ")], diff --git a/tst/integration/computer/app/artifacts/reset_apps/run.py b/tst/integration/computer/app/artifacts/reset_apps/run.py index fedff288f..72f7e9904 100644 --- a/tst/integration/computer/app/artifacts/reset_apps/run.py +++ b/tst/integration/computer/app/artifacts/reset_apps/run.py @@ -1,5 +1,5 @@ -"""Each app's data/reset, sent from inside the computer through the driver's route to it, and what the apps keep -after: the notes in ~/Notes, and whether the clock's alarms are there.""" +"""Each app's data/reset, sent from inside the computer through the driver's route to it with the computer's control +token, and what the apps keep after: the notes in ~/Notes, and whether the clock's alarms are there.""" import json import urllib.request from pathlib import Path @@ -10,7 +10,7 @@ def run(input): for app in ("notes", "clock"): request = urllib.request.Request(f"{input['env']['url']}/svc/{app}/agentenv", data=json.dumps( {"jsonrpc": "2.0", "id": 1, "method": "data/reset", "params": {}}).encode(), - headers={"Content-Type": "application/json"}) + headers={"Content-Type": "application/json", **input["env"]["headers"]}) with urllib.request.urlopen(request, timeout=60) as response: answers[app] = "error" not in json.load(response) return {"reset": answers, "notes": sorted(p.name for p in (Path.home() / "Notes").iterdir()), diff --git a/tst/integration/computer/computer_env_test.py b/tst/integration/computer/computer_env_test.py index 0f0467849..d258d04ea 100644 --- a/tst/integration/computer/computer_env_test.py +++ b/tst/integration/computer/computer_env_test.py @@ -7,17 +7,26 @@ kept by their paths in its environment artifact after its upload under the grant fails, comes back from a reset fresh, takes the snapshot back through load_artifact, and leaves nothing once torn down; an open one uploads each file under the snapshot's grant, kept where it was put; and a recording open when the computer resets or is torn down is ended by -the computer and registered whole. Needs Docker and the local registry on localhost:5000.""" +the computer and registered whole. A deployed computer answers its framework-only routes only with the control token +agent-env booted it with, from this machine or from another container as an agent would call, and keeps it through a +reset, while its MCP endpoint and card stay open to anyone. Needs Docker and the local registry on localhost:5000.""" import asyncio +import json import platform import subprocess import time import uuid +import httpx import pytest +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client -from agentenv_protocol.computer import COMPUTER_READY +from agentenv_protocol import RPC_PATH, WELL_KNOWN_PATH +from agentenv_protocol import client as protocol_v1 +from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_EXEC, COMPUTER_READY, COMPUTER_RECORDING +from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI from agentenv_protocol.computer.conformance import run_probe from agentenv_protocol.computer.conformance.client import Driver @@ -120,7 +129,7 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an artifact_id=note.id).execute(context) context = await RunCodeTaskStep(id="read", version=None, env_id=closed.id, script_artifact_id=script.id).execute(context) - driver = await Driver.connect(deployed.environment_url, 30) + driver = await Driver.connect(deployed.environment_url, 30, deployed.control_token) egress = await driver.run("python3", "-c", "import urllib.request; urllib.request.urlopen('http://1.1.1.1', timeout=5)") await driver.run("sh", "-c", "echo x > Desktop/left.txt") context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=closed.id, @@ -128,10 +137,10 @@ async def test_with_no_network_it_is_reached_reaches_nothing_comes_back_fresh_an file_paths=["Desktop", "Downloads/note.txt"]).execute(context) snapshot = {path: file.load() for path, file in _snapshotted(context, closed).items()} await ResetEnvTaskStep(id="reset", version=None, env_id=closed.id, fail_task_on_error=True).execute(context) - left = await (await Driver.connect(deployed.environment_url, 30)).run("test", "-e", "Desktop/left.txt") + left = await (await Driver.connect(deployed.environment_url, 30, deployed.control_token)).run("test", "-e", "Desktop/left.txt") context = await LoadArtifactTaskStep(id="restore", version=None, env_id=closed.id, artifact_from_step_id="snapshot").execute(context) - restored = await (await Driver.connect(deployed.environment_url, 30)).run("cat", "Desktop/left.txt") + restored = await (await Driver.connect(deployed.environment_url, 30, deployed.control_token)).run("cat", "Desktop/left.txt") finally: report = await teardown_run(context) sandbox = deployed.sandbox_id @@ -154,7 +163,7 @@ async def test_an_open_computer_uploads_each_file_under_the_snapshots_grant_and_ deployed = await Env.get(desktop.id).deploy() context = TaskStepContext(deployed_envs=[deployed]) try: - driver = await Driver.connect(deployed.environment_url, 30) + driver = await Driver.connect(deployed.environment_url, 30, deployed.control_token) await driver.run("sh", "-c", "mkdir -p Documents/reports && printf one > Documents/reports/a.txt " "&& printf two > Downloads/b.txt") context = await SnapshotComputerEnvTaskStep(id="snapshot", version=None, env_id=desktop.id, @@ -194,6 +203,74 @@ async def test_a_recording_open_at_a_reset_or_at_teardown_is_ended_by_its_comput assert entry["segments"] >= 1 +@pytest.mark.asyncio +async def test_only_agent_env_reaches_a_deployed_computers_framework_only_routes_while_its_mcp_and_card_stay_open( + desktop): + with namespace_routing(): + deployed = await Env.get(desktop.id).deploy() + context = TaskStepContext(deployed_envs=[deployed]) + try: + url, routes = deployed.environment_url, _framework_only(deployed.environment_card) + async with httpx.AsyncClient(timeout=30) as http: + refused = {name: [(await http.request(verb, url + path, json=body, headers=headers)).status_code + for headers in ({}, protocol_v1.control_headers("guess"))] + for name, (verb, path, body) in routes.items()} + card = await http.get(url + WELL_KNOWN_PATH) + from_a_container = _post_from_a_container(url + routes["computer-exec run"][1], + routes["computer-exec run"][2]) + async with streamablehttp_client(deployed.mcp_url) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + tools = {tool.name for tool in (await session.list_tools()).tools} + waited = await session.call_tool("computer_wait", {"ms": 0}) + driver = await Driver.connect(url, 30, deployed.control_token) + untouched = await driver.run("test", "-e", "Desktop/pwned") + await ResetEnvTaskStep(id="reset", version=None, env_id=desktop.id, + fail_task_on_error=True).execute(context) + after_reset = await (await Driver.connect(url, 30, deployed.control_token)).run("echo", "back") + verb, path, body = routes["computer-exec run"] + async with httpx.AsyncClient(timeout=30) as http: + still_refused = (await http.request(verb, url + path, json=body)).status_code + finally: + report = await teardown_run(context) + + assert not report.failed + assert refused == {name: [401, 401] for name in routes}, refused + assert from_a_container == 401 + assert card.status_code == 200 and card.json()["name"] == "desktop" + assert "computer_screenshot" in tools and not waited.isError + assert untouched["exit_code"] != 0, "a refused computer-exec call ran" + assert (after_reset["exit_code"], after_reset["stdout"]) == (0, "back\n") and still_refused == 401 + + +def _framework_only(card: dict) -> dict[str, tuple[str, str, dict]]: + """A call to each kind of framework-only route the card declares, by name: verb, path and body. The command would + leave Desktop/pwned behind, were it run.""" + calls = {"computer-exec run": (COMPUTER_EXEC, "run", {"argv": ["touch", "Desktop/pwned"], "request_id": "pwned"}), + "computer-ready status": (COMPUTER_READY, "status", None), + "computer-display get": (COMPUTER_DISPLAY, "get", None), + "computer-recording end_all": (COMPUTER_RECORDING, "end_all", {"reason": "pwned"}), + "disable-tool": (DISABLE_TOOL_EXTENSION_URI, "disable", {"role": "agent", "tools": ["computer_wait"]})} + routes = {name: (op["method"], op["endpoint"], body) + for name, (uri, method, body) in calls.items() + if (op := protocol_v1.find_extension_method(card, uri, method)) is not None} + assert len(routes) == len(calls), routes + return {**routes, "data/get": ("POST", RPC_PATH, {"jsonrpc": "2.0", "id": 1, "method": "data/get", "params": {}})} + + +def _post_from_a_container(url: str, body: dict) -> int: + """The status a POST of ``body`` to ``url`` gets from another container on this machine's network, as an agent + with a shell would send it.""" + script = ("import json, sys, urllib.request, urllib.error\n" + "request = urllib.request.Request(sys.argv[1], data=sys.argv[2].encode(), " + "headers={'Content-Type': 'application/json'})\n" + "try:\n print(urllib.request.urlopen(request, timeout=30).status)\n" + "except urllib.error.HTTPError as e:\n print(e.code)\n") + out = docker("run", "--rm", "--network", "host", "public.ecr.aws/docker/library/python:3.12-slim", + "python", "-c", script, url, json.dumps(body)) + return int(out.stdout.decode().strip()) + + def _record(name: str, env_id: str) -> StartComputerRecordingTaskStep: return StartComputerRecordingTaskStep(id=f"start-{name}", version=None, env_id=env_id, name=name, fail_task_on_error=True) diff --git a/tst/integration/computer/recording_chaos_test.py b/tst/integration/computer/recording_chaos_test.py index b619def33..a2ffb8583 100644 --- a/tst/integration/computer/recording_chaos_test.py +++ b/tst/integration/computer/recording_chaos_test.py @@ -7,6 +7,7 @@ import json import platform import subprocess +import time import uuid from contextlib import asynccontextmanager from datetime import datetime @@ -20,6 +21,7 @@ from agentenv_protocol.computer import COMPUTER_READY, COMPUTER_RECORDING, Recording, RecordingManifest from agentenv_protocol.computer.conformance.client import Driver from agentenv_protocol.computer.recording import DISK_CHECK_S, RECORDINGS_DIR +from agent_env.env.envs.computer.utils.boot import READY_TIMEOUT_S from tst.integration.computer.conftest import docker pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] @@ -123,8 +125,10 @@ async def _desk(image: str, *docker_args: str): async def _ready(name: str) -> Driver: - """A client for the container's driver once its computer is ready; a restart may move its port.""" - for _ in range(240): + """A client for the container's driver once its computer is ready, within the time a deploy gives one; a restart may + move its port.""" + deadline = time.monotonic() + READY_TIMEOUT_S + while time.monotonic() < deadline: port = subprocess.run(["docker", "port", name, str(PORT)], capture_output=True, text=True).stdout if port: try: @@ -134,7 +138,7 @@ async def _ready(name: str) -> Driver: except Exception: # not serving yet pass await asyncio.sleep(0.25) - raise TimeoutError(f"{name}'s computer never became ready") + raise TimeoutError(f"{name}'s computer wasn't ready within {READY_TIMEOUT_S:.0f} s") async def _layout(driver: Driver) -> dict: diff --git a/tst/unit/computer_home.py b/tst/unit/computer_home.py index 1d9d8e3dc..9564ca2fd 100644 --- a/tst/unit/computer_home.py +++ b/tst/unit/computer_home.py @@ -1,6 +1,6 @@ """A computer at its driver's boundary, for steps that reach inside one: its home a folder on this machine, ``data/add`` writing there, ``data/get`` reading there, and computer-exec/v1 ``run`` running commands there; each call recorded, -and the address each ``data/add`` is sent to. +with the headers it sent, and the address each ``data/add`` is sent to. ``pkill`` matches nothing, so one test never ends another's processes on this machine.""" from __future__ import annotations @@ -19,6 +19,8 @@ from agent_env.env.env import DeployedSandboxEnv PORT = 18765 +#: The control token the computer was booted with, which each call to its driver carries. +TOKEN = "t0ken" CARD = {"name": "desk", "capabilities": {"operations": ["data/reset", "data/add", "data/get"], "extensions": [ {"uri": COMPUTER_EXEC, "params": {"endpoint": "/computer/exec", "methods": {"run": {}}, "max_output_bytes": 1 << 20, "max_timeout_s": 3600}}]}} @@ -29,13 +31,14 @@ def __init__(self, home: Path): self.home = home self.calls: list[tuple] = [] self.added_at: list[str] = [] + self.headers: list[dict | None] = [] self.add_fails: Exception | None = None def deployed(self, env_id: str = "desk", card: dict | None = None, **computer) -> DeployedSandboxEnv: return DeployedSandboxEnv(env_id=env_id, env_version=1, env_provider_type="computer", environment_card_url="http://127.0.0.1:20001/.well-known/agent-env.json", environment_card=card or CARD, sandbox_id="local-1", sandbox_type="local", - metadata={"computer": {"port": PORT, **computer}}) + metadata={"computer": {"port": PORT, **computer}}, control_token=TOKEN) def install(self, monkeypatch) -> Computer: monkeypatch.setattr(protocol_v1, "add_data", self.add_data) @@ -43,9 +46,10 @@ def install(self, monkeypatch) -> Computer: monkeypatch.setattr(protocol_v1, "invoke_extension", self.invoke_extension) return self - async def add_data(self, base_url, parts, timeout=120, verify=True): + async def add_data(self, base_url, parts, timeout=120, verify=True, *, headers=None): self.calls.append(("data/add", [p.metadata[DATA_OBJECTS_EXTENSION_URI] for p in parts], timeout)) self.added_at.append(base_url) + self.headers.append(headers) if self.add_fails is not None: raise self.add_fails for part in parts: @@ -54,15 +58,18 @@ async def add_data(self, base_url, parts, timeout=120, verify=True): target.parent.mkdir(parents=True, exist_ok=True) target.write_bytes(base64.b64decode(part.file.bytes)) - async def get_data(self, base_url, timeout=30, verify=True, *, write_namespace=None, paths=None): + async def get_data(self, base_url, timeout=30, verify=True, *, write_namespace=None, paths=None, headers=None): self.calls.append(("data/get", paths)) + self.headers.append(headers) return GetDataResponse(parts=[ FilePart(file=FileWithBytes(bytes=base64.b64encode((self.home / path).read_bytes()).decode()), metadata={DATA_OBJECTS_EXTENSION_URI: {"path": path}}) for path in paths if (self.home / path).is_file()]) - async def invoke_extension(self, base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def invoke_extension(self, base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, + headers=None): assert (uri, method) == (COMPUTER_EXEC, "run") + self.headers.append(headers) self.calls.append(("run", params, timeout)) return await asyncio.to_thread(self._run, params) diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index d1095a050..9aaf728a2 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -6,6 +6,7 @@ import asyncio import contextlib +import dataclasses import hashlib import importlib.metadata import json @@ -19,6 +20,7 @@ import httpx import pytest +from agentenv_protocol import CONTROL_TOKEN_ENV from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI from click.testing import CliRunner @@ -34,6 +36,7 @@ from agent_env.config import get_config from agent_env.env import Env from agent_env.env import env as env_module +from agent_env.env.env import DeployedEnv from agent_env.env.envs.computer.utils import boot, builder, composer, gate, inside, packages, resolve_in_base from agent_env.env.envs.computer.utils.composer import Composing, _probe from agent_env.env.envs.computer.env import ComputerEnv @@ -251,23 +254,27 @@ def test_the_pairs_left_to_compose_are_those_failed_interrupted_or_unprobed_on_a assert composer.pending(Env.get(env.id)) == [] -def test_the_probe_boots_a_build_at_the_envs_settings_as_a_deploy_does(base, driver, docker, monkeypatch): +def test_the_probe_boots_a_build_at_the_envs_settings_with_a_token_of_its_own_as_a_deploy_does(base, driver, docker, + monkeypatch): env = _put(base, driver, compose=False, computer_settings={"timezone": "Europe/Paris"}, display={"width": 1600, "height": 900}) booted = [] async def fake_boot(link, **kwargs): - booted.append(kwargs["environment"]) + booted.append(kwargs) return boot.Booted(sandbox=None, url="http://computer", card=CARD) monkeypatch.setattr(boot, "boot", fake_boot) monkeypatch.setattr(boot, "discard", AsyncMock()) - monkeypatch.setattr(composer, "run_probe", AsyncMock(return_value="report")) + probed = AsyncMock(return_value="report") + monkeypatch.setattr(composer, "run_probe", probed) assert asyncio.run(_probe([SimpleNamespace()], env, driver, "img")) == ( "report", CARD, "SimpleNamespace") - assert {k: v for k, v in booted[0].items() if k.startswith("AGENTENV_SETTING_")} == { + assert {k: v for k, v in booted[0]["environment"].items() if k.startswith("AGENTENV_SETTING_")} == { "AGENTENV_SETTING_RESOLUTION": "1600x900", "AGENTENV_SETTING_TIMEZONE": "Europe/Paris"} + assert len(booted[0]["token"]) >= 32 + probed.assert_awaited_once_with("http://computer", "full", token=booted[0]["token"]) def test_a_build_whose_probe_fails_or_that_never_gets_ready_fails_at_the_probe_and_isnt_deployed(base, driver, docker, @@ -503,8 +510,8 @@ def test_builds_of_one_env_at_once_keep_each_others_pairs_and_a_failure_doesnt_r @pytest.fixture def boots(monkeypatch): """The boots a deploy asks for, each a sandbox serving the desktop driver's card at 127.0.0.1:20001; the driver - shows a 1600x900 screen, and each extension call it gets is recorded.""" - made = SimpleNamespace(boots=[], calls=[], sandboxes=[], display=(1600, 900)) + shows a 1600x900 screen, and each extension call it gets is recorded, with the headers it carried.""" + made = SimpleNamespace(boots=[], calls=[], headers=[], sandboxes=[], display=(1600, 900)) async def fake_boot(sandbox_provider, **kwargs): made.boots.append({**kwargs, "provider": type(sandbox_provider).__name__}) @@ -519,8 +526,9 @@ async def terminate(): return boot.Booted(sandbox, "http://127.0.0.1:20001", {**DRIVER_CARD, "name": kwargs["environment"]["ENVIRONMENT_NAME"]}) - async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, headers=None): made.calls.append((uri, method, params)) + made.headers.append(headers) width, height = made.display return {"width": width, "height": height, "scale": 1, "modes": []} if uri == COMPUTER_DISPLAY else {} @@ -530,7 +538,7 @@ async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, m @pytest.mark.asyncio -async def test_a_computer_env_deploys_its_first_native_ready_build_with_its_settings_and_display( +async def test_a_computer_env_deploys_its_first_native_ready_build_with_its_settings_display_and_a_control_token( base, driver, docker, boots): env = _put(base, driver, display={"width": 1600, "height": 900}) @@ -545,7 +553,12 @@ async def test_a_computer_env_deploys_its_first_native_ready_build_with_its_sett assert (environment["AGENTENV_SETTING_RESOLUTION"], environment["AGENTENV_SETTING_TIMEZONE"]) == ("1600x900", "Asia/Tokyo") assert json.loads(environment["AGENTENV_COMPUTER_SETTINGS"]) == {"default_space": "norm"} assert environment["ENVIRONMENT_NAME"] == "desktop" and "AGENTENV_SETTING_THEME" not in environment + assert len(deployed.control_token) >= 32 and booted["token"] == deployed.control_token assert boots.calls == [(COMPUTER_DISPLAY, "get", None)] + assert boots.headers == [{"Authorization": f"Bearer {deployed.control_token}"}] + assert DeployedEnv.from_dict(dataclasses.asdict(deployed)).control_token == deployed.control_token + second = await Env.get(env.id).deploy() + assert second.control_token != deployed.control_token assert (deployed.env_provider_type, deployed.sandbox_ids) == ("computer", {"computer": {"desktop": "local-1"}}) assert deployed.environment_card == DRIVER_CARD and deployed.mcp_url == "http://127.0.0.1:20001/mcp" record = deployed.metadata["computer"] @@ -580,7 +593,9 @@ async def reset_container(self, sandbox, **kwargs): assert replaced is sandbox and not sandbox.terminated assert (kwargs["image_name"], kwargs["port"]) == (deployed.metadata["computer"]["image"], 18765) assert (kwargs["env"]["AGENTENV_SETTING_TIMEZONE"], kwargs["env"]["ENVIRONMENT_NAME"]) == ("Asia/Tokyo", "desktop") - env_computer_provider.ready.assert_awaited_once_with(deployed.environment_url, boot.READY_TIMEOUT_S) + assert kwargs["env"][CONTROL_TOKEN_ENV] == deployed.control_token + env_computer_provider.ready.assert_awaited_once_with(deployed.environment_url, boot.READY_TIMEOUT_S, + deployed.control_token) assert boots.calls == [(COMPUTER_RECORDING, "end_all", {"reason": "reset"}), "reset_container", (COMPUTER_DISPLAY, "get", None)] @@ -659,37 +674,46 @@ async def test_a_computer_that_boots_at_another_size_than_its_display_fails_the_ @pytest.mark.asyncio -async def test_a_boot_that_never_gets_ready_takes_its_sandbox_and_work_folder_down(monkeypatch, tmp_path): +async def test_a_boot_hands_the_computer_its_token_and_one_never_ready_takes_its_sandbox_and_work_folder_down( + monkeypatch, tmp_path): monkeypatch.setenv("AGENT_ENV_LOCAL_SANDBOX_DIR", str(tmp_path)) sandbox = local_sandbox.LocalSandbox(exposed_ports=[18765], sandbox_id="local-unready") sandbox.terminate = AsyncMock() monkeypatch.setattr(boot, "ready", AsyncMock(side_effect=TimeoutError("not ready"))) + create = AsyncMock(return_value=sandbox) with pytest.raises(TimeoutError): - await boot.boot(SimpleNamespace(create_container=AsyncMock(return_value=sandbox)), image="img", port=18765, - environment={}, network_policy=None, timeout=60, attribution=None) + await boot.boot(SimpleNamespace(create_container=create), image="img", port=18765, + environment={"ENVIRONMENT_NAME": "desk"}, token="t0ken", network_policy=None, timeout=60, + attribution=None) + assert create.await_args.kwargs["env"] == {"ENVIRONMENT_NAME": "desk", CONTROL_TOKEN_ENV: "t0ken"} + assert boot.ready.await_args.args[2] == "t0ken" sandbox.terminate.assert_awaited_once() assert local_sandbox.LocalSandbox.find_work_dir("local-unready") is None @pytest.mark.asyncio -async def test_a_boot_is_ready_with_the_card_read_once_its_status_says_so_each_app_on_it(monkeypatch): +async def test_a_boot_is_ready_with_the_card_read_once_its_status_asked_with_its_token_says_so_each_app_on_it( + monkeypatch): cards = iter([{"name": "desk"}, {"name": "desk"}, {"name": "desk", "children_environments": [{"name": "mail"}]}]) statuses = iter([{"ready": False, "checks": [{"name": "app mail", "ok": False, "detail": "starting"}]}, {"ready": True}]) + sent = [] async def card(url, timeout): return next(cards) - async def status(url, card, uri, params, timeout, method): + async def status(url, card, uri, params, timeout, method, headers): + sent.append(headers) return next(statuses) monkeypatch.setattr(boot.protocol_v1, "get_card", card) monkeypatch.setattr(boot.protocol_v1, "invoke_extension", status) monkeypatch.setattr(boot, "_POLL_S", 0) - assert await boot.ready("http://desk", 5) == {"name": "desk", "children_environments": [{"name": "mail"}]} + assert await boot.ready("http://desk", 5, "t0ken") == {"name": "desk", "children_environments": [{"name": "mail"}]} + assert sent == [{"Authorization": "Bearer t0ken"}] * 2 def test_a_boot_mints_a_generated_setting_resolves_a_secret_its_record_withholds_both_and_a_secret_is_a_reference( diff --git a/tst/unit/env/task_step/verify_env_core_protocol_test.py b/tst/unit/env/task_step/verify_env_core_protocol_test.py index f1c75be26..857c516c6 100644 --- a/tst/unit/env/task_step/verify_env_core_protocol_test.py +++ b/tst/unit/env/task_step/verify_env_core_protocol_test.py @@ -108,8 +108,10 @@ def test_missing_deployed_env_raises(self): assert "not found in context.deployed_envs" in str(e) -def _probe(payload): +def _probe(payload, status=200): class _Resp: + status_code = status + def json(self): return payload @@ -127,7 +129,7 @@ async def post(self, *args, **kwargs): return _Resp() with patch("agent_env.task_step.task_steps.env_card_validator.verify_env_core_protocol.httpx.AsyncClient", _Client): - return asyncio.run(_make_step()._probe_operation("https://gw", "data/add", {"parts": []})) + return asyncio.run(_make_step()._probe_operation("https://gw", "data/add", {"parts": []}, {})) class TestProbeOperation: @@ -143,6 +145,9 @@ def test_gateway_forward_error_is_indeterminate(self): def test_result_means_registered(self): assert _probe({"jsonrpc": "2.0", "id": 1, "result": {"parts": []}}) is True + def test_a_refusal_for_the_control_token_is_indeterminate(self): + assert _probe({"error": {"code": "unauthorized", "message": "no", "applied": "no"}}, status=401) is None + class TestWire: @pytest.mark.parametrize("gateway_url", ["https://gw", "https://gw/", "https://sandbox.example.com/sandbox/sandbox-vm-1-18765"]) @@ -150,7 +155,7 @@ def test_reads_the_card_and_probes_the_data_plane_at_the_cards_address(self, gat base = gateway_url.rstrip("/") ctx = TaskStepContext(deployed_envs=[DeployedGatewayEnv( env_id="e1", env_version=3, gateway_url=gateway_url, mcp_url=f"{base}/mcp", db_web_url=None, sandbox_id="sb", - environment_card_url=f"{gateway_url}{WELL_KNOWN_PATH}", + environment_card_url=f"{gateway_url}{WELL_KNOWN_PATH}", control_token="s3cret", )]) sent = [] @@ -169,6 +174,7 @@ def handle(request): ("GET", f"{base}{WELL_KNOWN_PATH}"), ("POST", f"{base}/agentenv"), ("POST", f"{base}/agentenv"), ] assert [(b["method"], b["params"]) for b in (json.loads(r.content) for r in sent[1:])] == [("data/add", {"parts": []}), ("data/get", {})] + assert [r.headers.get("authorization") for r in sent] == [None, "Bearer s3cret", "Bearer s3cret"] assert fake_env.update_metadata.call_args[0][0]["validated_environment_protocol"] == { "data/reset": {"supported": True}, "data/add": {"supported": True}, "data/get": {"supported": True}, } diff --git a/tst/unit/env/test_env.py b/tst/unit/env/test_env.py index 0f40fe177..03bda1b95 100644 --- a/tst/unit/env/test_env.py +++ b/tst/unit/env/test_env.py @@ -113,6 +113,23 @@ async def test_a_record_without_a_card_offers_nothing(sent): assert sent == [] +@pytest.mark.asyncio +async def test_a_record_with_a_control_token_sends_it_on_each_extension_call_only_to_its_address_and_keeps_it( + composed, sent): + deployed = _record(composed) + deployed.control_token = "t0ken" + env = deployed.environment_url + + await deployed.invoke(EXT_STEP_URI, "step", {"action": "list_tools"}) + + assert [request.headers.get("authorization") for request in sent] == ["Bearer t0ken"] + assert [deployed.control_headers_for(url) for url in (env, f"{env}/svc/clock/agentenv", f"{env}0/x", + "https://override/trajectory")] == [ + {"Authorization": "Bearer t0ken"}] * 2 + [{}] * 2 + assert DeployedEnv.from_dict(dataclasses.asdict(deployed)).control_token == "t0ken" + assert _record(composed).control_headers == {} and _record(composed).control_headers_for(env) == {} + + def test_child_env_cards_come_from_the_card_by_name(composed): deployed = _record(composed) assert deployed.get_child_env_card("slack")["url"] == "/svc/mcp-slack/agentenv" diff --git a/tst/unit/env/test_legacy_protocol.py b/tst/unit/env/test_legacy_protocol.py index 6f5a60a19..9508e8d2e 100644 --- a/tst/unit/env/test_legacy_protocol.py +++ b/tst/unit/env/test_legacy_protocol.py @@ -10,6 +10,7 @@ from agentenv_protocol.client import GetDataResponse from agent_env.env import legacy_protocol +from agent_env.env.env import DeployedSandboxEnv @pytest.mark.asyncio @@ -145,6 +146,26 @@ def handle(request: httpx.Request) -> httpx.Response: return asked +@pytest.mark.asyncio +async def test_a_childs_state_is_read_with_its_deployments_control_token(monkeypatch): + seen, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + seen.append((request.url.path, request.headers.get("authorization"))) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {"parts": [ + {"kind": "data", "data": {"alarms": []}}]}}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + deployed = DeployedSandboxEnv( + env_id="desk", env_version=1, sandbox_id="sb-1", environment_card_url=f"http://desk:18765{WELL_KNOWN_PATH}", + environment_card={"name": "desk", "capabilities": {"operations": ["data/get"]}, + "children_environments": [{"name": "clock", "url": "/svc/clock/agentenv"}]}, + control_token="t0ken") + + assert await legacy_protocol.service_state(deployed, None, "clock") == {"alarms": []} + assert seen == [("/svc/clock/agentenv", "Bearer t0ken")] + + @pytest.mark.asyncio @pytest.mark.parametrize("answer, state, reads_export_state", [ ([DataPart(data={"messages": 2})], {"messages": 2}, False), diff --git a/tst/unit/store/aws/test_shared_credentials.py b/tst/unit/store/aws/test_shared_credentials.py index 2b50bcab5..049f4a01c 100644 --- a/tst/unit/store/aws/test_shared_credentials.py +++ b/tst/unit/store/aws/test_shared_credentials.py @@ -19,7 +19,7 @@ def pushed(monkeypatch): calls: list[dict] = [] - async def fake_invoke(base_url, card, uri, params=None, timeout=30): + async def fake_invoke(base_url, card, uri, params=None, timeout=30, headers=None): calls.append(params) monkeypatch.setattr(protocol_v1, "invoke_extension", fake_invoke) @@ -43,7 +43,7 @@ def _merged_env(caller_env: dict[str, str] | None = None, metadata: dict | None async def test_nothing_is_shared_by_default(s3, host_credentials, pushed): if s3: set_object_store(S3ObjectStore(object(), "bucket", region="us-west-2")) - await _push_s3_credentials("https://svc", _CARD, 30) + await _push_s3_credentials("https://svc", _CARD, 30, None) assert not any(k.startswith("AWS_") for k in _merged_env()) assert pushed == [] @@ -51,7 +51,7 @@ async def test_nothing_is_shared_by_default(s3, host_credentials, pushed): @pytest.mark.asyncio async def test_an_opted_in_s3_store_shares_the_frozen_chain(host_credentials, pushed): set_object_store(S3ObjectStore(object(), "bucket", region="us-west-2", share_credentials=True)) - await _push_s3_credentials("https://svc", _CARD, 30) + await _push_s3_credentials("https://svc", _CARD, 30, None) assert {k: v for k, v in _merged_env().items() if k.startswith("AWS_")} == { "AWS_ACCESS_KEY_ID": "AKIA-HOST", "AWS_SECRET_ACCESS_KEY": "host-secret", @@ -72,7 +72,7 @@ async def test_an_opted_in_s3_store_shares_the_frozen_chain(host_credentials, pu async def test_an_local_run_hands_a_service_no_credentials_to_upload_a_snapshot_with(owner, shared, host_credentials, pushed, cli_routing): set_object_store(S3ObjectStore(object(), "bucket", region="us-west-2", share_credentials=True)) with routing.run_scope(owner): - await _push_s3_credentials("https://svc", _CARD, 30) + await _push_s3_credentials("https://svc", _CARD, 30, None) assert [call["bucket"] for call in pushed] == (["bucket"] if shared else []) diff --git a/tst/unit/task/control_grant_test.py b/tst/unit/task/control_grant_test.py new file mode 100644 index 000000000..8e76ffbcc --- /dev/null +++ b/tst/unit/task/control_grant_test.py @@ -0,0 +1,74 @@ +"""Every step runs under a grant of each deployed env's control token: a protocol-client call the step makes to an env's +address carries that env's token unless the call names its own headers, an env without one is sent none, an env an +earlier step deployed is granted to the steps after it, and the grant ends with the step.""" + +import pytest +from agentenv_protocol import WELL_KNOWN_PATH +from agentenv_protocol import client as protocol_v1 + +from agent_env.config.runtime import Config +from agent_env.env.env import DeployedEnv +from agent_env.task import Task +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep + +DESK = "http://desk:18765" +MAIL = "http://mail:8000" + + +def _deployed(env_id: str, url: str, token: str | None = None) -> DeployedEnv: + return DeployedEnv(env_id=env_id, env_version=1, mcp_url=f"{url}/mcp", environment_card_url=f"{url}{WELL_KNOWN_PATH}", + control_token=token) + + +class _Calls(TaskStep): + """Records what its calls to the desk, to the mail env and elsewhere send, and one to the desk with headers of its + own.""" + + type = "control_grant_test_calls" + entity_refs = () + seen: list = [] + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data)) + + async def execute(self, context): + urls = (f"{DESK}/agentenv/ext/computer-exec/run", f"{MAIL}/agentenv", "http://elsewhere/agentenv") + _Calls.seen.append([protocol_v1.request_headers(url) for url in urls] + + [protocol_v1.request_headers(f"{DESK}/agentenv", {"X-Own": "1"})]) + return context + + +class _DeployDesk(TaskStep): + """Deploys the desk, with its token.""" + + type = "control_grant_test_deploy" + entity_refs = () + + @classmethod + def from_dict(cls, data): + return cls(**cls._base_from_dict(data)) + + async def execute(self, context): + context.deployed_envs.append(_deployed("desk", DESK, "s3cret")) + return context + + +@pytest.fixture(autouse=True) +def registries(monkeypatch): + steps = {**Config().task_step_registry(), _Calls.type: _Calls, _DeployDesk.type: _DeployDesk} + monkeypatch.setattr(Config, "task_step_registry", lambda self: steps) + monkeypatch.setattr(_Calls, "seen", []) + + +@pytest.mark.asyncio +async def test_a_step_calls_each_deployed_env_with_its_own_token_and_the_grant_ends_with_the_step(local_stores): + task = Task(id="t", version=None, steps=[_Calls(id="before", version=None), _DeployDesk(id="deploy", version=None), + _Calls(id="after", version=None)]) + + await task.run(context=TaskStepContext(deployed_envs=[_deployed("mail", MAIL)])) + + assert _Calls.seen == [[{}, {}, {}, {"X-Own": "1"}], + [{"Authorization": "Bearer s3cret"}, {}, {}, {"X-Own": "1"}]] + assert protocol_v1.request_headers(f"{DESK}/agentenv") == {} diff --git a/tst/unit/task_step/add_skills_env_test.py b/tst/unit/task_step/add_skills_env_test.py index 3b8001ea5..576536f04 100644 --- a/tst/unit/task_step/add_skills_env_test.py +++ b/tst/unit/task_step/add_skills_env_test.py @@ -35,7 +35,7 @@ def agent(monkeypatch): """The agent's skill registrations and the env calls the step made.""" made = SimpleNamespace(calls=[], register=AsyncMock(return_value={"ok": True})) - async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, headers=None): made.calls.append((base_url, card["name"], uri, method, params)) return {"name": params["name"], "description": "How to use this env", "skill_md": f"---\nname: {params['name']}\n"} @@ -69,7 +69,7 @@ async def test_each_envs_skill_is_installed_under_its_environment_name_and_recor async def test_a_skill_with_files_is_stored_once_and_installed_as_a_bundle(agent, local_stores, monkeypatch): files = {"LINUX.md": "# Linux\n", "WORKFLOW.md": "# Workflow\n"} - async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, headers=None): return {"name": params["name"], "description": "Drive the desktop", "skill_md": "# Skill\n", "files": files} monkeypatch.setattr(env_module.protocol_v1, "invoke_extension", invoke) @@ -86,7 +86,7 @@ async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, m @pytest.mark.asyncio async def test_a_skill_file_whose_path_leaves_its_folder_is_refused_before_anything_is_stored(agent, local_stores, monkeypatch): - async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def invoke(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, headers=None): return {"name": params["name"], "description": "d", "skill_md": "# Skill\n", "files": {"../../victim": "x"}} monkeypatch.setattr(env_module.protocol_v1, "invoke_extension", invoke) diff --git a/tst/unit/task_step/snapshot_env_upload_grant_test.py b/tst/unit/task_step/snapshot_env_upload_grant_test.py index 9554f333c..b9c82a9d4 100644 --- a/tst/unit/task_step/snapshot_env_upload_grant_test.py +++ b/tst/unit/task_step/snapshot_env_upload_grant_test.py @@ -43,7 +43,7 @@ def store(local_stores, tmp_path): def pushes(monkeypatch): pushed: list[str] = [] - async def push(base_url, card, timeout_seconds): + async def push(base_url, card, timeout_seconds, headers): pushed.append(base_url) monkeypatch.setattr(snapshot_mod, "_push_s3_credentials", push) diff --git a/tst/unit/task_step/snapshot_env_without_gateway_test.py b/tst/unit/task_step/snapshot_env_without_gateway_test.py index a262d6407..ad213fe01 100644 --- a/tst/unit/task_step/snapshot_env_without_gateway_test.py +++ b/tst/unit/task_step/snapshot_env_without_gateway_test.py @@ -1,6 +1,6 @@ -"""snapshot_env on an env deployed without a gateway: reached through the card its deployment recorded, its -environments the card's children its names reach that export, and the env itself when it exports beside them, or the -env alone, each kept as one file as through a gateway and sent its own data/get params; an environment +"""snapshot_env on an env deployed without a gateway: reached through the card its deployment recorded, with its control +token, its environments the card's children its names reach that export, and the env itself when it exports beside +them, or the env alone, each kept as one file as through a gateway and sent its own data/get params; an environment whose data/get returns a file tree, a computer's, kept file by file by its paths under a grant sized for a file tree; and an env that fails under the upload grant asked again without it.""" @@ -124,6 +124,24 @@ async def test_the_env_itself_is_reached_when_its_card_has_no_children(local_sto assert sent == [("/agentenv", {})] +@pytest.mark.asyncio +async def test_a_record_with_a_control_token_sends_it_on_its_data_get(local_stores, monkeypatch): + seen, real = [], httpx.AsyncClient + + def handle(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("authorization")) + return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {"parts": [ + {"kind": "data", "data": {"rows": 1}}]}}) + + monkeypatch.setattr(httpx, "AsyncClient", lambda *a, **k: real(transport=httpx.MockTransport(handle))) + record = _record(_server()) + record.control_token = "t0ken" + + await SnapshotEnvTaskStep.snapshot_env_state(env=_Desk(), gateway_url=None, snapshot_id="snap-t", deployed=record) + + assert seen == ["Bearer t0ken"] + + @pytest.mark.asyncio async def test_a_file_tree_is_kept_file_by_file_by_its_paths_under_a_grant_sized_for_it(local_stores, tmp_path, monkeypatch): diff --git a/tst/unit/task_step/sync_env_clock_test.py b/tst/unit/task_step/sync_env_clock_test.py index 31ff6735f..471ca3a3a 100644 --- a/tst/unit/task_step/sync_env_clock_test.py +++ b/tst/unit/task_step/sync_env_clock_test.py @@ -35,7 +35,7 @@ async def test_sync_one_skips_unadvertised(monkeypatch): async def fake_card(base_url, **kw): return {"capabilities": {"extensions": []}} monkeypatch.setattr(mod.protocol_v1, "get_card", fake_card) - out = await _step()._sync_one(None, "http://gw", "email", CLOCK_URL) + out = await _step()._sync_one(_cardless(), "http://gw", "email", CLOCK_URL) assert out == {"synced": False, "reason": "not_advertised"} @@ -55,7 +55,7 @@ async def fake_invoke(base_url, c, uri, params=None, **kw): monkeypatch.setattr(mod.protocol_v1, "get_card", fake_card) monkeypatch.setattr(mod.protocol_v1, "invoke_extension", fake_invoke) - out = await _step()._sync_one(None, "http://gw", "email", CLOCK_URL) + out = await _step()._sync_one(_cardless(), "http://gw", "email", CLOCK_URL) assert out["synced"] is True assert calls["uri"] == EXT_CLOCK_URI and calls["params"] == {"env_get_time_url": CLOCK_URL} @@ -66,7 +66,7 @@ async def fake_card(base_url, **kw): return {"capabilities": {"extensions": []}} monkeypatch.setattr(mod.protocol_v1, "get_card", fake_card) with pytest.raises(RuntimeError): - await _step(tolerate_missing_sync_time=False)._sync_one(None, "http://gw", "email", CLOCK_URL) + await _step(tolerate_missing_sync_time=False)._sync_one(_cardless(), "http://gw", "email", CLOCK_URL) @pytest.mark.asyncio @@ -144,6 +144,12 @@ async def fetch(client, key, base_url): return json.loads((await gw._serve_env_card()).body) +def _cardless() -> DeployedEnv: + """A record that kept no card, which a sync reads live from the gateway.""" + return DeployedGatewayEnv(env_id="e1", env_version=1, gateway_url="http://gw", mcp_url="http://gw/mcp", db_web_url=None, + sandbox_id="sb-1") + + def _carded(card: dict) -> DeployedEnv: from agent_env.env.gateway.constants import WELL_KNOWN_PATH diff --git a/tst/unit/task_step/test_apply_server_config.py b/tst/unit/task_step/test_apply_server_config.py index 07dec1269..f9badb26c 100644 --- a/tst/unit/task_step/test_apply_server_config.py +++ b/tst/unit/task_step/test_apply_server_config.py @@ -88,14 +88,14 @@ def __init__(self, cards: dict[str, dict]): self.get_calls: list[str] = [] self.post_calls: list[tuple[str, dict]] = [] - async def get(self, url, *, timeout=None, params=None): + async def get(self, url, *, timeout=None, params=None, headers=None): self.get_calls.append(url) for service, card in self.cards.items(): if f"/svc/mcp-{service}/.well-known/agent-env.json" in url: return _resp(200, url, card) return _resp(404, url, {"error": "no card"}) - async def request(self, method, url, *, json=None, timeout=None): + async def request(self, method, url, *, json=None, timeout=None, headers=None): self.post_calls.append((url, json)) return _resp(200, url, {"ok": True, "echo": json}) @@ -269,7 +269,7 @@ async def test_tolerate_unadvertised_still_fails_on_rejected_args(): # fails loud — that's the fail-loud-on-unknown-persona contract. gw = _MockGateway({"slack": _card(_ext(_SET_ACTING_USER_URI, "/agentenv/ext/set_acting_user"))}) - async def rejecting_request(method, url, *, json=None, timeout=None): + async def rejecting_request(method, url, *, json=None, timeout=None, headers=None): gw.post_calls.append((url, json)) return _resp(500, url, {"error": "user_not_found: nobody@x.com"}) @@ -452,7 +452,7 @@ async def test_non_404_card_fetch_still_raises_under_tolerate(): # endpoint that errors (500) is a real fault and must still fail loud. gw = _MockGateway({}) - async def failing_get(url, *, timeout=None, params=None): + async def failing_get(url, *, timeout=None, params=None, headers=None): gw.get_calls.append(url) return _resp(500, url, {"error": "boom"}) @@ -472,7 +472,7 @@ async def failing_get(url, *, timeout=None, params=None): async def test_server_rejects_args_raises_with_status(): gw = _MockGateway({"slack": _card(_ext(_SET_ERRORS_URI, "/agentenv/ext/set_errors"))}) - async def rejecting_request(method, url, *, json=None, timeout=None): + async def rejecting_request(method, url, *, json=None, timeout=None, headers=None): gw.post_calls.append((url, json)) return _resp(400, url, {"error": "error_rate must be between 0.0 and 1.0"}) diff --git a/tst/unit/task_step/test_prompt_agent_triggers.py b/tst/unit/task_step/test_prompt_agent_triggers.py index a34bda4a3..5f26930a6 100644 --- a/tst/unit/task_step/test_prompt_agent_triggers.py +++ b/tst/unit/task_step/test_prompt_agent_triggers.py @@ -47,7 +47,7 @@ def _resp(status: int, body: dict, method: str = "GET") -> httpx.Response: async def test_snapshot_shape_multi_env(): sent = [] - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): sent.append((url, params, timeout)) return _resp(200, {"triggers": [{"id": "v6-rate", "status": "fired"}, {"id": "v6-accept", "status": "armed"}]}) @@ -63,7 +63,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_an_env_whose_card_offers_no_triggers_state_reads_as_empty_with_no_request(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): raise AssertionError("no request for an env whose card offers no triggers state") ctx = _ctx(_env("env-x", "http://gw-x", EXT_CLOCK_URI)) @@ -75,7 +75,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_a_404_from_an_advertised_triggers_state_fails_closed(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): return _resp(404, {}) ctx = _ctx(_env("env-x", "http://gw-x")) @@ -87,7 +87,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_fail_closed_raises_on_persistent_error(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): raise httpx.ConnectError("gateway down") ctx = _ctx(_env("env-x", "http://gw-x")) @@ -101,7 +101,7 @@ async def fake_get(self, url, timeout, params=None): async def test_retries_then_succeeds(): calls = {"n": 0} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): calls["n"] += 1 if calls["n"] == 1: raise httpx.ConnectError("transient") @@ -119,7 +119,7 @@ async def fake_get(self, url, timeout, params=None): async def test_triggers_state_is_read_at_the_stored_cards_address(): sent = [] - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): sent.append(url) return _resp(200, {"triggers": []}) @@ -220,7 +220,7 @@ def _patched_store(store=None, *, key_prefix=""): async def test_capture_uploads_state_and_records_summary(): calls = {"n": 0} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): if url.endswith("/clock/state"): return _resp(200, _CLOCK) calls["n"] += 1 @@ -256,7 +256,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_capture_is_kept_under_the_fixture_prefix(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): return _resp(200, _CLOCK if url.endswith("/clock/state") else _STATE) ctx = TaskStepContext( @@ -274,7 +274,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_capture_records_clock_context_and_per_trigger_detail(): """A recurrence reads "armed" forever, so only fire_count shows it fired.""" - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): return _resp(200, _CLOCK if url.endswith("/clock/state") else _CLOCK_STATE) ctx = TaskStepContext( @@ -313,7 +313,7 @@ async def test_capture_distinguishes_a_failed_coupling_from_one_that_never_fired "failure_count": 2, "last_failure_at": "2026-06-01T12:00:00Z"}, ]} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): return _resp(200, _CLOCK if url.endswith("/clock/state") else failed) ctx = TaskStepContext( @@ -343,7 +343,7 @@ async def test_capture_does_not_wait_on_a_firing_recurring_clock_trigger(): "status": "firing", "fire_count": 3, "next_mark": "2026-06-03T00:00:00Z"}]} calls = {"n": 0} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): if url.endswith("/clock/state"): return _resp(200, _CLOCK) calls["n"] += 1 @@ -366,7 +366,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_capture_tolerates_a_gateway_without_clock_state(): """An unreachable clock state: capture still succeeds, clock recorded as None.""" - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): if url.endswith("/clock/state"): raise httpx.ConnectError("no clock route") return _resp(200, _STATE) @@ -397,7 +397,7 @@ async def fake_get(self, url, timeout, params=None): async def test_capture_records_no_clock_with_no_request_when_the_card_offers_no_clock_state(): sent = [] - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): sent.append(url) return _resp(200, _STATE) @@ -419,7 +419,7 @@ async def fake_get(self, url, timeout, params=None): async def test_capture_reads_triggers_and_clock_at_the_stored_cards_address(): sent = [] - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): sent.append((url, timeout)) return _resp(200, _CLOCK if url.endswith("/clock/state") else _STATE) @@ -488,7 +488,7 @@ def test_settle_poll_does_not_wait_out_a_delay_or_working_hours(): async def test_capture_noop_without_registrations(): calls = {"n": 0} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): calls["n"] += 1 return _resp(200, _STATE) @@ -502,7 +502,7 @@ async def fake_get(self, url, timeout, params=None): @pytest.mark.asyncio async def test_capture_fail_open_records_error(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): raise httpx.ConnectError("gateway gone") ctx = TaskStepContext( @@ -526,7 +526,7 @@ async def test_capture_grace_polls_until_settled(): responses = [firing, _STATE] calls = {"n": 0} - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): if url.endswith("/clock/state"): return _resp(200, _CLOCK) body = responses[min(calls["n"], len(responses) - 1)] @@ -572,7 +572,7 @@ async def hanging_get(self, url, timeout): @pytest.mark.asyncio async def test_capture_put_failure_keeps_inline_summary(): - async def fake_get(self, url, timeout, params=None): + async def fake_get(self, url, timeout, params=None, headers=None): return _resp(200, _STATE) class _FailingStore(_FakeStore): diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index 3d48575a8..de5865b58 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -26,7 +26,7 @@ from agent_env.providers.sandbox_providers.sandbox import Sandbox, VmSandbox from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps import run_code as _run_code_module -from tst.unit.computer_home import CARD, Computer +from tst.unit.computer_home import CARD, TOKEN, Computer from tst.unit.event_loop_probe import on_event_loop from agent_env.task_step.task_steps.run_code import ( RunCodeExecutionError, @@ -1309,13 +1309,15 @@ async def _run_in(computer, body, deployed=None, **step): @pytest.mark.asyncio -async def test_a_script_for_a_computer_runs_inside_it_and_reaches_its_driver_on_loopback(computer): +async def test_a_script_for_a_computer_runs_inside_it_and_reaches_its_driver_on_loopback_with_its_token(computer): body = b"import os\ndef run(inp):\n return {'env': inp['env'], 'cwd': os.getcwd()}\n" ctx = await _run_in(computer, body) result = ctx.metadata["script_results"]["run-code-1"] - assert result["env"] == {"id": "desk", "url": "http://127.0.0.1:18765", "mcp_url": "http://127.0.0.1:18765/mcp"} + assert result["env"] == {"id": "desk", "url": "http://127.0.0.1:18765", "mcp_url": "http://127.0.0.1:18765/mcp", + "headers": {"Authorization": f"Bearer {TOKEN}"}} + assert computer.headers and all(sent == {"Authorization": f"Bearer {TOKEN}"} for sent in computer.headers) assert Path(result["cwd"]).resolve() == computer.home.resolve() run_dir = ".agentenv/run_code/run-code-1" assert [call[1] for call in computer.calls if call[0] == "data/add"] == [ @@ -1349,7 +1351,7 @@ async def test_a_run_the_computer_stopped_at_its_timeout_is_reported_as_one(comp async def test_a_run_whose_call_drops_fails_naming_why_and_the_next_one_runs_afresh(computer, monkeypatch): invoke, dropped = computer.invoke_extension, [] - async def drop_once(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None): + async def drop_once(base_url, card, uri, params=None, timeout=30, verify=True, *, method=None, headers=None): if params["argv"][0] == "python3" and not dropped: dropped.append(params["request_id"]) raise httpx.RemoteProtocolError("the driver went away") diff --git a/tst/unit/task_step/test_snapshot_agent_state.py b/tst/unit/task_step/test_snapshot_agent_state.py index 8682c673e..551f1cd26 100644 --- a/tst/unit/task_step/test_snapshot_agent_state.py +++ b/tst/unit/task_step/test_snapshot_agent_state.py @@ -167,7 +167,7 @@ async def _no_probe(_base): got, legacy = [], [] - async def _get_data(base, timeout=30): + async def _get_data(base, timeout=30, headers=None): got.append(base) return GetDataResponse(parts=[DataPart(data={"rows": 1})]) @@ -200,7 +200,7 @@ async def test_a_service_that_exports_a_file_bundle_is_captured_from_its_export_ monkeypatch.setattr(artifact_mod, "EnvironmentUniverseArtifact", type("U", (), {"get": staticmethod(lambda _id: _StubUniverse())})) monkeypatch.setattr(config_mod, "get_config", lambda: type("Cfg", (), {"get_object_store_at": lambda self, object_url: store})()) - async def _get_data(base, timeout=30): + async def _get_data(base, timeout=30, headers=None): if base.endswith("/mcp-calendar"): return GetDataResponse(parts=[uploaded_file_part("calendar.zip", name="calendar.zip", mime_type="application/zip")]) return GetDataResponse(parts=[DataPart(data={"contacts": 4})]) diff --git a/tst/unit/task_step/test_snapshot_env.py b/tst/unit/task_step/test_snapshot_env.py index b118bc393..877244c37 100644 --- a/tst/unit/task_step/test_snapshot_env.py +++ b/tst/unit/task_step/test_snapshot_env.py @@ -236,6 +236,7 @@ def __init__(self, response=None, error=None): self.response = response self.error = error self.requests = [] + self.headers = [] def __call__(self, **kwargs): return self @@ -246,8 +247,9 @@ async def __aenter__(self): async def __aexit__(self, *exc): return False - def stream(self, method, url, params=None): + def stream(self, method, url, params=None, headers=None): self.requests.append((method, url, params)) + self.headers.append(headers) client = self class _CM: @@ -428,7 +430,7 @@ async def test_trajectory_capture_records_one_sentence_when_the_card_offers_no_t @pytest.mark.asyncio -async def test_trajectory_capture_streams_from_the_stored_cards_address_and_endpoint(monkeypatch): +async def test_trajectory_capture_streams_from_the_stored_cards_address_and_endpoint_with_its_token(monkeypatch): _wire_fake_snapshot(monkeypatch) client = _FakeAsyncClient(response=_FakeTrajectoryResponse([b"{}\n"])) _wire_fake_capture(monkeypatch, client) @@ -436,22 +438,28 @@ async def test_trajectory_capture_streams_from_the_stored_cards_address_and_endp deployed.environment_card_url = f"https://sandbox.example/sb-1{WELL_KNOWN_PATH}" deployed.environment_card = {"capabilities": {"extensions": [{ "uri": "urn:agentenv:trajectory/v1", "params": {"endpoint": "/env/trajectory", "methods": {"get": {"method": "GET"}}}}]}} + deployed.control_token = "t0ken" await _step(include_env_trajectory=True).execute(_ctx(deployed)) assert [(m, u) for m, u, _ in client.requests] == [("GET", "https://sandbox.example/sb-1/env/trajectory")] + assert client.headers == [{"Authorization": "Bearer t0ken"}] @pytest.mark.asyncio @pytest.mark.parametrize("deployments", [[], [_deployed()]], ids=["pure-override", "next-to-a-deployment"]) -async def test_trajectory_capture_with_a_gateway_url_override_reads_the_override(monkeypatch, deployments): +async def test_trajectory_capture_with_a_gateway_url_override_reads_the_override_without_a_token(monkeypatch, + deployments): _wire_fake_snapshot(monkeypatch) client = _FakeAsyncClient(response=_FakeTrajectoryResponse([b"{}\n"])) _wire_fake_capture(monkeypatch, client) + for deployed in deployments: + deployed.control_token = "t0ken" await _step(env_id="env-1", gateway_url="https://override/", include_env_trajectory=True).execute(_ctx(*deployments)) assert [(m, u) for m, u, _ in client.requests] == [("GET", "https://override/trajectory")] + assert "Authorization" not in (client.headers[0] or {}) # ── env changes capture ────────────────────────────────────────────────────── @@ -464,7 +472,7 @@ def __init__(self, rows=None, error=None, page_size=2): self.rows = rows or [] self.page_size = page_size - async def request(self, method, url, params=None): + async def request(self, method, url, params=None, headers=None): self.requests.append((method, url, params)) if self.error: raise self.error @@ -736,7 +744,7 @@ async def test_export_service_to_file_materializes_filepart_zip(monkeypatch, tmp async def fake_supports_v1(base_url): return True - async def fake_get_data(base_url, timeout=None): + async def fake_get_data(base_url, timeout=None, headers=None): return GetDataResponse(parts=[part]) async def fake_get_card(base_url, timeout=10): @@ -782,7 +790,7 @@ async def test_export_service_to_file_resolves_relative_uri_and_streams( async def fake_supports_v1(base_url): return True - async def fake_get_data(base_url, timeout=None): + async def fake_get_data(base_url, timeout=None, headers=None): return GetDataResponse(parts=[part]) async def fake_get_card(base_url, timeout=10): @@ -815,7 +823,7 @@ async def __aenter__(self): async def __aexit__(self, *exc): return False - def stream(self, method, url, timeout=None): + def stream(self, method, url, timeout=None, headers=None): captured["method"] = method captured["url"] = url return _FakeStreamCtx() @@ -840,7 +848,7 @@ def _serve_filepart_uri(monkeypatch, uri): async def fake_supports_v1(base_url): return True - async def fake_get_data(base_url, timeout=None): + async def fake_get_data(base_url, timeout=None, headers=None): return GetDataResponse(parts=[part]) async def fake_get_card(base_url, timeout=10): @@ -975,7 +983,7 @@ async def test_export_service_to_file_filepart_without_bytes_or_uri_raises_value async def fake_supports_v1(base_url): return True - async def fake_get_data(base_url, timeout=None): + async def fake_get_data(base_url, timeout=None, headers=None): return GetDataResponse(parts=[part]) async def fake_get_card(base_url, timeout=10): diff --git a/tst/unit/task_step/verifiers/test_env_outcome_verifier_step.py b/tst/unit/task_step/verifiers/test_env_outcome_verifier_step.py index 8c98dfaee..0ae826cb0 100644 --- a/tst/unit/task_step/verifiers/test_env_outcome_verifier_step.py +++ b/tst/unit/task_step/verifiers/test_env_outcome_verifier_step.py @@ -5,6 +5,7 @@ import pytest from agent_env.artifact import FileArtifact +from agent_env.env.env import DeployedEnv from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.verifiers.env_outcome_verifier import EnvOutcomeVerifierTaskStep, ScoreAggregator from tst.unit.event_loop_probe import on_event_loop @@ -77,7 +78,7 @@ def load(): monkeypatch.setattr(FileArtifact, "get", staticmethod(lambda id, version=None: SimpleNamespace(load=load))) step = EnvOutcomeVerifierTaskStep(id="outcome", version=None, env_id="e", file_artifact_id="script", verifier_id="v") ctx = TaskStepContext() - ctx.deployed_envs.append(SimpleNamespace(env_id="e", mcp_url="http://mcp")) + ctx.deployed_envs.append(DeployedEnv(env_id="e", env_version=1, mcp_url="http://mcp")) asyncio.run(step.execute(ctx))