From 1e338262b37c50cfd249794039145dee9b23fe8e Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 13:34:08 -0800 Subject: [PATCH 1/8] feat(website): record a website env's traffic as PCAP, served through traffic/v1 (FD-4683) This PR's change, rebased as one commit onto #448. Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 4 +- .../examples/website/notes/README.md | 4 + .../src/agentenv_protocol/__init__.py | 14 ++ .../src/agentenv_protocol/traffic.py | 86 ++++++++ .../agentenv-protocol/tests/test_traffic.py | 87 ++++++++ src/agent_env/cli/env/deploy.py | 7 +- src/agent_env/env/envs/_deployment.py | 8 +- src/agent_env/env/envs/multi_env.py | 3 +- src/agent_env/env/envs/website.py | 3 +- src/agent_env/env/gateway/Dockerfile | 2 +- src/agent_env/env/gateway/constants.py | 15 ++ src/agent_env/env/gateway/entrypoint.py | 5 +- src/agent_env/env/gateway/gateway.py | 95 ++++++++- .../env_providers/env_gateway_provider.py | 152 ++++++++++++- src/agent_env/task/episode/rules.py | 3 +- src/agent_env/task/teardown.py | 16 +- src/agent_env/task_step/registry.py | 2 + .../task_step/task_steps/deploy_env.py | 12 +- .../task_step/task_steps/traffic_capture.py | 147 +++++++++++++ tst/unit/env/envs/test_multi_env_name.py | 3 +- tst/unit/env/gateway/traffic_test.py | 200 ++++++++++++++++++ .../env/gateway/website_websocket_test.py | 20 ++ .../env_providers/traffic_capture_test.py | 199 +++++++++++++++++ .../deploy_env_capture_traffic_test.py | 123 +++++++++++ tst/unit/task_step/traffic_capture_test.py | 158 ++++++++++++++ 25 files changed, 1340 insertions(+), 28 deletions(-) create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/traffic.py create mode 100644 packages/agentenv-protocol/tests/test_traffic.py create mode 100644 src/agent_env/task_step/task_steps/traffic_capture.py create mode 100644 tst/unit/env/gateway/traffic_test.py create mode 100644 tst/unit/providers/env_providers/traffic_capture_test.py create mode 100644 tst/unit/task_step/deploy_env_capture_traffic_test.py create mode 100644 tst/unit/task_step/traffic_capture_test.py diff --git a/AGENTS.md b/AGENTS.md index a8b4d6a4d..fb35c3ec8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -117,9 +117,9 @@ and deserialized through a registry. | Package | Contents | |---|---| -| `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`, one site: a server image, `backend_image` in its env.toml, serving the pages and API, ready when its card's `website/v1` `ready` answers a `WebsiteReadiness` that says so, as `agentenv_protocol.website` (`website_readiness`, `env_readiness`) defines a website env's health for the deploy and `DeployedEnv.readiness()` alike, or, for a card without `ready`, when its `health` path answers 200, which is deprecated; the gateway container's `env/gateway/constants.py` mirrors the URI and op, a test keeping them equal; a separate frontend image is deprecated, warned of at put and optional in a bundle, and a stored `backend_docker_image_artifact` reads as the server image), `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`), a run only the forms its sandbox's deploys boot (`booted_forms`: on each link, the first form it boots natively, by the deploy's own rule, `boot.first_native`; a compose of a disk adds its container), and a deploy composes first what it boots that no compose started (`ComputerEnv.compose_unstarted`); 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 (`utils/probing.py`) 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, over its env.toml's `[settings]` and a computer env's `settings` beside it in `apps`, each by top-level field (`params.app_settings`; checked at put, a required one left to the deploy; the deploy record keeps the merged values as `app_settings`, secrets withheld); 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. Its site route, `/website//`, never forwards a site's data plane or card (`/agentenv/…`, `/.well-known/agent-env.json`), which agentenv reaches through `/svc//`; a site with no frontend image is served by its server image there. 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`. | +| `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`, one site: a server image, `backend_image` in its env.toml, serving the pages and API, ready when its card's `website/v1` `ready` answers a `WebsiteReadiness` that says so, as `agentenv_protocol.website` (`website_readiness`, `env_readiness`) defines a website env's health for the deploy and `DeployedEnv.readiness()` alike, or, for a card without `ready`, when its `health` path answers 200, which is deprecated; the gateway container's `env/gateway/constants.py` mirrors the URI and op, a test keeping them equal; a separate frontend image is deprecated, warned of at put and optional in a bundle, and a stored `backend_docker_image_artifact` reads as the server image), `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`), a run only the forms its sandbox's deploys boot (`booted_forms`: on each link, the first form it boots natively, by the deploy's own rule, `boot.first_native`; a compose of a disk adds its container), and a deploy composes first what it boots that no compose started (`ComputerEnv.compose_unstarted`); 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 (`utils/probing.py`) 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, over its env.toml's `[settings]` and a computer env's `settings` beside it in `apps`, each by top-level field (`params.app_settings`; checked at put, a required one left to the deploy; the deploy record keeps the merged values as `app_settings`, secrets withheld); 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. Its site route, `/website//`, never forwards a site's data plane or card (`/agentenv/…`, `/.well-known/agent-env.json`), which agentenv reaches through `/svc//`; a site with no frontend image is served by its server image there. A deploy with `capture_traffic` (`deploy_env`'s, `env deploy --capture-traffic`; a website env, or a multi env with a website, on a VM, refused for any other env and in container mode) records every packet each website's server and frontend containers send and receive: beside each runs a capture service, tcpdump from the gateway's image in that container's network namespace (`network_mode: "service:"`), the only service holding `NET_RAW` and `NET_ADMIN`, rotating its classic pcap files (10 of 20 MB each, the oldest overwritten) onto a volume no website container mounts and the gateway mounts read-only and serves as `urn:agentenv:traffic/v1` (`list`, `GET /traffic/captures`, oldest first; `get`, `GET /traffic/capture?name=`), which `agentenv_protocol.traffic` defines (`TrafficCapture`, `list_traffic_captures`, `fetch_traffic_capture`), the gateway's `constants.py` mirroring its names, a test keeping them equal. Those routes take the per-deploy control token the record keeps (`DeployedEnv.control_token`, the gateway's `AGENTENV_CONTROL_TOKEN`), which the gateway takes off every request before it proxies one, so no site reads it; a capture that isn't running once the gateway is up refuses the deploy. 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/`, among them `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`). | +| `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/`, among them `start_computer_recording` and `end_computer_recording`, `collect_traffic_capture`, which stores each traffic/v1 capture of an env deployed with `capture_traffic` as a pcap `FileArtifact` in one `FileArtifactUniverse` recorded under `context.metadata["traffic_captures"][env_id]`, 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 capturing env still up whose traffic the run didn't collect is collected, 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 namespaces, `@local` and `@hf` (`is_local_id`), each with a documents file of its own: while the CLI or a `build_app` server runs, a local id and what a 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). | diff --git a/packages/agentenv-protocol/examples/website/notes/README.md b/packages/agentenv-protocol/examples/website/notes/README.md index 861d6d643..0d55da1c4 100644 --- a/packages/agentenv-protocol/examples/website/notes/README.md +++ b/packages/agentenv-protocol/examples/website/notes/README.md @@ -99,6 +99,10 @@ agent-env env website put --id notes-site \ Run it from the repo root. `environment_name` comes from the card, which the spec names `notes`. The stored env has a server image and no frontend image. Pass `--skip-validation` to register without deploying. +To record what the site sends and receives, deploy it with `--capture-traffic` (or `"capture_traffic": true` on a +task's `deploy_env` step): tcpdump runs in the server container's network namespace and the gateway serves its pcap +files as `traffic/v1`, which a `collect_traffic_capture` step with the env's `env_id` stores as file artifacts. + `env.toml` describes the same env for a bundle. Its keys are `type`, `spec` and `backend_image`; the env takes its `environment_name`, health path and domains from the spec. A bundle builds `backend_image` with this folder as the build context, but this Dockerfile needs the repo root. Until the protocol package is installed from a diff --git a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py index 61f9cbab9..fc67f4f3a 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/__init__.py @@ -4,6 +4,14 @@ from .client import ExtensionCallError from .errors import Applied, AgentEnvError from .preflight import intake_fit_check +from . import traffic +from .traffic import ( + TRAFFIC_EXTENSION_URI, + TrafficCapture, + TrafficCaptureListing, + fetch_traffic_capture, + list_traffic_captures, +) from .tool_client import ToolResult, ToolSession, ToolSessionError from .agent_env_environment import ( HIDDEN_BY_DEFAULT_META_KEY, @@ -88,6 +96,12 @@ __all__ = [ "client", + "traffic", + "TRAFFIC_EXTENSION_URI", + "TrafficCapture", + "TrafficCaptureListing", + "list_traffic_captures", + "fetch_traffic_capture", "AgentEnvApplication", "AgentEnvEnvironment", "AgentEnvFastMCPApplication", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/traffic.py b/packages/agentenv-protocol/src/agentenv_protocol/traffic.py new file mode 100644 index 000000000..98289a643 --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/traffic.py @@ -0,0 +1,86 @@ +"""Traffic capture, ``urn:agentenv:traffic/v1``: an env that records the packets its containers send and receive serves +them as capture files through its card. ``list`` (``TRAFFIC_LIST_OP``) answers a ``TrafficCaptureListing``, each file a +``TrafficCapture`` named for the container it holds, oldest first; ``get`` (``TRAFFIC_GET_OP``) streams the file a +``list`` entry names, given as its ``name`` param, as ``PCAP_MEDIA_TYPE``. Both are framework-only: they take the env's +control token (``client.control_headers``). ``serves_traffic`` reads whether a card advertises it, +``list_traffic_captures`` and ``fetch_traffic_capture`` call it as a card advertises it.""" +from __future__ import annotations + +from datetime import datetime + +import httpx +from pydantic import BaseModel, ConfigDict + +from .client import ExtensionCallError, find_extension_method, invoke_extension, request_headers + +TRAFFIC_EXTENSION_URI = "urn:agentenv:traffic/v1" +TRAFFIC_LIST_OP = "list" +TRAFFIC_GET_OP = "get" +TRAFFIC_FORMAT = "pcap" +PCAP_MEDIA_TYPE = "application/vnd.tcpdump.pcap" +LIST_TIMEOUT_S = 30.0 +FETCH_TIMEOUT = httpx.Timeout(600.0, connect=10.0) + + +class TrafficCapture(BaseModel): + """One capture file a ``list`` names: its ``name``, which ``get`` takes, the ``container`` whose packets it holds, + its size in ``bytes`` and when it was last written. Fields it doesn't name are ignored, so a reader keeps reading + an env that says more.""" + + model_config = ConfigDict(extra="ignore") + + name: str + container: str + bytes: int + modified_at: datetime + + +class TrafficCaptureListing(BaseModel): + """``list``'s answer: the files' ``format`` and each ``TrafficCapture``, oldest first within a container, so its + files read in the order their packets were captured.""" + + model_config = ConfigDict(extra="ignore") + + format: str = TRAFFIC_FORMAT + captures: list[TrafficCapture] = [] + + +def serves_traffic(card: dict) -> bool: + """Whether ``card`` advertises traffic/v1's ``list``; no network.""" + return isinstance(card, dict) and find_extension_method(card, TRAFFIC_EXTENSION_URI, TRAFFIC_LIST_OP) is not None + + +async def list_traffic_captures(base_url: str, card: dict, *, timeout: float = LIST_TIMEOUT_S, verify: bool = True, + headers: dict[str, str] | None = None) -> TrafficCaptureListing: + """The captures the env whose ``card`` is read from ``base_url`` serves, through its ``list``. Raises + ``ValueError`` when the card doesn't advertise it or the answer isn't a ``TrafficCaptureListing``, httpx's error + when it can't be asked, and ``ExtensionCallError`` for a non-2xx answer.""" + answer = await invoke_extension(base_url, card, TRAFFIC_EXTENSION_URI, timeout=timeout, verify=verify, + method=TRAFFIC_LIST_OP, headers=headers) + return TrafficCaptureListing.model_validate(answer) + + +async def fetch_traffic_capture(base_url: str, card: dict, name: str, path: str, *, + timeout: httpx.Timeout | float = FETCH_TIMEOUT, verify: bool = True, + headers: dict[str, str] | None = None) -> int: + """Stream the capture ``name``, as a ``list`` entry names it, from the env whose ``card`` is read from + ``base_url`` through its ``get`` to the file ``path``, and return its size. Raises ``ValueError`` when the card + doesn't advertise ``get``, httpx's error when it can't be asked, and ``ExtensionCallError`` for a non-2xx answer, + with nothing written.""" + op = find_extension_method(card, TRAFFIC_EXTENSION_URI, TRAFFIC_GET_OP) + if op is None or not op["endpoint"]: + raise ValueError(f"card does not advertise {TRAFFIC_EXTENSION_URI} {TRAFFIC_GET_OP}") + url = f"{base_url}{op['endpoint']}" + params = {"name": name} + request = {"params": params} if op["method"] == "GET" else {"json": params} + size = 0 + async with httpx.AsyncClient(verify=verify, timeout=timeout) as client: + async with client.stream(op["method"], url, headers=request_headers(url, headers), **request) as response: + if response.status_code >= 400: + await response.aread() + raise ExtensionCallError(response) + with open(path, "wb") as out: + async for chunk in response.aiter_bytes(): + out.write(chunk) + size += len(chunk) + return size diff --git a/packages/agentenv-protocol/tests/test_traffic.py b/packages/agentenv-protocol/tests/test_traffic.py new file mode 100644 index 000000000..4c09e3a0c --- /dev/null +++ b/packages/agentenv-protocol/tests/test_traffic.py @@ -0,0 +1,87 @@ +"""traffic/v1's client helpers: a card says whether an env serves it, ``list_traffic_captures`` reads its listing as +typed captures and ``fetch_traffic_capture`` streams one to a file, each at the endpoint and verb the card advertises, +with the headers given or granted.""" +from __future__ import annotations + +from datetime import datetime, timezone + +import httpx +import pytest + +from agentenv_protocol import client as protocol_v1 +from agentenv_protocol.traffic import ( + PCAP_MEDIA_TYPE, + TRAFFIC_EXTENSION_URI, + TrafficCapture, + fetch_traffic_capture, + list_traffic_captures, + serves_traffic, +) + +_RealAsyncClient = httpx.AsyncClient +CARD = {"capabilities": {"extensions": [{"uri": TRAFFIC_EXTENSION_URI, "params": {"methods": { + "list": {"method": "GET", "endpoint": "/svc/x/traffic/captures"}, + "get": {"method": "GET", "endpoint": "/svc/x/traffic/capture"}}}}]}} +LISTING = {"format": "pcap", "extra": 1, "captures": [ + {"name": "site-backend.pcap0", "container": "site-backend", "bytes": 24, "modified_at": "1970-01-01T00:16:40+00:00", + "more": "ignored"}]} + + +@pytest.fixture +def seen(monkeypatch): + requests: list[httpx.Request] = [] + + def env(request: httpx.Request) -> httpx.Response: + requests.append(request) + if request.headers.get("authorization") != "Bearer tok": + return httpx.Response(401, json={"ok": False, "error": {"code": "unauthorized", "message": "no token"}}) + if request.url.path.endswith("/captures"): + return httpx.Response(200, json=LISTING) + return httpx.Response(200, content=b"pcap-bytes", headers={"content-type": PCAP_MEDIA_TYPE}) + + transport = httpx.MockTransport(env) + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + return requests + + +def test_a_card_serves_traffic_only_when_it_advertises_list(): + assert serves_traffic(CARD) + assert not serves_traffic({"capabilities": {"extensions": []}}) + assert not serves_traffic(None) + + +@pytest.mark.asyncio +async def test_list_reads_the_advertised_endpoint_as_typed_captures(seen): + listing = await list_traffic_captures("http://env", CARD, headers=protocol_v1.control_headers("tok")) + assert listing.format == "pcap" + assert listing.captures == [TrafficCapture(name="site-backend.pcap0", container="site-backend", bytes=24, + modified_at=datetime(1970, 1, 1, 0, 16, 40, tzinfo=timezone.utc))] + assert [str(request.url) for request in seen] == ["http://env/svc/x/traffic/captures"] + + +@pytest.mark.asyncio +async def test_fetch_streams_the_named_capture_to_a_file_with_the_granted_headers(seen, tmp_path): + path = tmp_path / "out.pcap" + with protocol_v1.authorized({"http://env": protocol_v1.control_headers("tok")}): + size = await fetch_traffic_capture("http://env", CARD, "site-backend.pcap0", str(path)) + assert (size, path.read_bytes()) == (10, b"pcap-bytes") + assert str(seen[0].url) == "http://env/svc/x/traffic/capture?name=site-backend.pcap0" + + +@pytest.mark.asyncio +async def test_a_refusal_raises_an_extension_call_error_and_writes_nothing(seen, tmp_path): + path = tmp_path / "out.pcap" + with pytest.raises(protocol_v1.ExtensionCallError) as refused: + await fetch_traffic_capture("http://env", CARD, "site-backend.pcap0", str(path), headers={}) + with pytest.raises(protocol_v1.ExtensionCallError): + await list_traffic_captures("http://env", CARD, headers={}) + assert refused.value.code == "unauthorized" and not path.exists() + + +@pytest.mark.asyncio +async def test_a_card_without_traffic_is_refused_before_any_call(seen, tmp_path): + with pytest.raises(ValueError): + await list_traffic_captures("http://env", {}, headers={}) + with pytest.raises(ValueError): + await fetch_traffic_capture("http://env", {}, "x.pcap0", str(tmp_path / "x"), headers={}) + assert seen == [] diff --git a/src/agent_env/cli/env/deploy.py b/src/agent_env/cli/env/deploy.py index f7afb1040..cdd17e6fa 100644 --- a/src/agent_env/cli/env/deploy.py +++ b/src/agent_env/cli/env/deploy.py @@ -51,7 +51,10 @@ @click.option("--keyed-roles", is_flag=True, default=False, help="Draw a role key for the deploy: its gateway then believes a role other than default only from a " "caller that sends the key beside it.") -def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: str, sandbox: str, service_db_env_id: str | None, gateway_env_id: str | None, website_browser_env_id: str | None, disk_size_gb: float, cpu: float | None, memory_mb: int | None, env_state_type: str | None, env_state_instance_id: str | None, keyed_roles: bool): +@click.option("--capture-traffic", is_flag=True, default=False, + help="Record every packet the env's websites' containers send and receive, as rotating pcap files its " + "gateway serves as traffic/v1 (VM deploys of website or multi envs only).") +def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: str, sandbox: str, service_db_env_id: str | None, gateway_env_id: str | None, website_browser_env_id: str | None, disk_size_gb: float, cpu: float | None, memory_mb: int | None, env_state_type: str | None, env_state_instance_id: str | None, keyed_roles: bool, capture_traffic: bool): """Deploy an MCP server environment.""" from agent_env.providers import build_sandbox_provider from agent_env.config import get_config @@ -94,6 +97,8 @@ def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: click.echo(f"Env state instance id: {env_state_instance_id}") if keyed_roles: deploy_kwargs["keyed_roles"] = True + if capture_traffic: + deploy_kwargs["capture_traffic"] = True try: with run_scope(env.id): deployed_env = asyncio.run(env.deploy(**deploy_kwargs)) diff --git a/src/agent_env/env/envs/_deployment.py b/src/agent_env/env/envs/_deployment.py index 7659455b0..6dc06a735 100644 --- a/src/agent_env/env/envs/_deployment.py +++ b/src/agent_env/env/envs/_deployment.py @@ -288,13 +288,19 @@ def provider_or_class(env: Env) -> EnvironmentProvider | type[EnvironmentProvide def deploy_refusal(env: Env, provider: EnvironmentProvider | type[EnvironmentProvider], options: dict) -> str | None: """Why deploying env through provider, or a provider of that class, is refused before anything is built, or None: the provider - can't deploy an env like it (see ``provider_refusal``), or an option it can't take.""" + can't deploy an env like it (see ``provider_refusal``), ``capture_traffic`` for an env that serves no website, or an option it + can't take.""" + from agent_env.env.envs.website import WebsiteEnv + provider_class = provider if isinstance(provider, type) else type(provider) reason = provider_refusal(provider_class, type(env), mcp_names=[child.environment_name for child in getattr(env, "mcp_server_envs", ())], website_names=[child.environment_name for child in getattr(env, "website_envs", ())]) if reason is not None: return f"env '{env.id}' has env_provider_type '{env.env_provider_type}', which {reason}" + if options.get("capture_traffic") and not (isinstance(env, WebsiteEnv) or getattr(env, "website_envs", ())): + return (f"env '{env.id}' serves no website, and capture_traffic records what a website's containers send and " + "receive") return option_refusal(env, provider, options) diff --git a/src/agent_env/env/envs/multi_env.py b/src/agent_env/env/envs/multi_env.py index d1da07009..3789e308b 100644 --- a/src/agent_env/env/envs/multi_env.py +++ b/src/agent_env/env/envs/multi_env.py @@ -157,12 +157,13 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: ctx.refuse(problems) return fields - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False, capture_traffic: bool = False) -> DeployedEnv: deployed_env = await deploy_through_provider( self, environment_name=None, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, artifact_id=artifact_id, artifact_version=artifact_version, **({"keyed_roles": True} if keyed_roles else {}), + **({"capture_traffic": True} if capture_traffic else {}), ) self._gateway_mode = gateway_mode self._mcp_server_name = deployed_env.mcp_server_name diff --git a/src/agent_env/env/envs/website.py b/src/agent_env/env/envs/website.py index c70a5e206..48848e21c 100644 --- a/src/agent_env/env/envs/website.py +++ b/src/agent_env/env/envs/website.py @@ -210,12 +210,13 @@ def accept_toml(cls, data: dict, ctx: AuthoringContext) -> dict: fields["domains"] = spec.site.domains return fields - async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False) -> DeployedEnv: + async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, artifact_id: str | None = None, artifact_version: int | None = None, *, attribution: Optional[Attribution] = None, keyed_roles: bool = False, capture_traffic: bool = False) -> DeployedEnv: return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, artifact_id=artifact_id, artifact_version=artifact_version, **({"keyed_roles": True} if keyed_roles else {}), + **({"capture_traffic": True} if capture_traffic else {}), ) def deploy_refusal(self, *, sandbox_type: str | None = None, **options) -> str | None: diff --git a/src/agent_env/env/gateway/Dockerfile b/src/agent_env/env/gateway/Dockerfile index 73000f01d..4957dd39f 100644 --- a/src/agent_env/env/gateway/Dockerfile +++ b/src/agent_env/env/gateway/Dockerfile @@ -1,6 +1,6 @@ FROM public.ecr.aws/docker/library/python:3.11-slim@sha256:c20888b6acdd1e63e1c433a185bf3ad162c0288fe484616ce062e0d28add2900 -RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/* +RUN apt-get update && apt-get install -y curl tcpdump && rm -rf /var/lib/apt/lists/* WORKDIR /app diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 784eb080f..0af3614a2 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -81,6 +81,11 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int: EXT_NOTIFICATIONS_URI = "urn:agentenv:notifications/v1" EXT_WEBSITE_URI = "urn:agentenv:website/v1" WEBSITE_READY_OP = "ready" +EXT_TRAFFIC_URI = "urn:agentenv:traffic/v1" +TRAFFIC_LIST_OP = "list" +TRAFFIC_GET_OP = "get" +TRAFFIC_CAPTURE_DIR_ENV = "TRAFFIC_CAPTURE_DIR" +PCAP_MEDIA_TYPE = "application/vnd.tcpdump.pcap" WEBSITE_BROWSE_URL_PARAM = "browse_url" # Rows per GET /changes page: the default, and the most a client may ask for. @@ -149,6 +154,7 @@ def sources(names: list[str]) -> str: AGENT_ENV_SITE_ROOT_HEADER = "AgentEnv-Site-Root" GATEWAY_ROLE_KEY_ENV = "AGENTENV_ROLE_KEY" GATEWAY_PROVIDER_HEADERS_ENV = "AGENTENV_PROVIDER_REQUEST_HEADERS" +CONTROL_TOKEN_ENV = "AGENTENV_CONTROL_TOKEN" WEBSITE_ROLE_KEY_ENV = "WEBSITE_ROLE_KEY" # MCP request `_meta` keys the gateway stamps on every child tools/call: the caller's role, and on the MCP path a # key naming the inbound session, so a child can keep per-caller state behind the gateway's one shared session. @@ -272,3 +278,12 @@ def random_mcp_server_name() -> str: "required": False, }, ] +TRAFFIC_EXTENSION = { + "uri": EXT_TRAFFIC_URI, + "description": "Traffic (urn:agentenv:traffic/v1): the packets each website's server and frontend containers sent and received, captured by tcpdump in each container's network namespace as rotating classic pcap files. list names each capture file with the container it holds, its size and when it was last written, oldest first; get streams one as application/vnd.tcpdump.pcap. A container's oldest file is overwritten once its rotation is full. Harness control path: with a control token, each call takes it as its bearer.", + "params": {"endpoint": "/traffic/captures", "format": "pcap", "methods": { + TRAFFIC_LIST_OP: {"method": "GET", "endpoint": "/traffic/captures", "request": {"type": "object", "properties": {}}}, + TRAFFIC_GET_OP: {"method": "GET", "endpoint": "/traffic/capture", "request": {"type": "object", "properties": { + "name": {"type": "string", "minLength": 1}}, "required": ["name"]}}}}, + "required": False, +} diff --git a/src/agent_env/env/gateway/entrypoint.py b/src/agent_env/env/gateway/entrypoint.py index 2912fbdbc..451f718e5 100644 --- a/src/agent_env/env/gateway/entrypoint.py +++ b/src/agent_env/env/gateway/entrypoint.py @@ -7,7 +7,8 @@ import httpx from . import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode, InternalMCPServer -from .constants import DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV, GATEWAY_ROLE_KEY_ENV +from .constants import (CONTROL_TOKEN_ENV, DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV, GATEWAY_ROLE_KEY_ENV, + TRAFFIC_CAPTURE_DIR_ENV) from .gateway import Gateway logger = logging.getLogger(__name__) @@ -154,5 +155,7 @@ def parse_name_url_map(env_str: str) -> dict[str, str]: service_db_url=service_db_url, role_key=os.environ.get(GATEWAY_ROLE_KEY_ENV) or None, provider_headers=os.environ.get(GATEWAY_PROVIDER_HEADERS_ENV, "").split(","), + traffic_dir=os.environ.get(TRAFFIC_CAPTURE_DIR_ENV) or None, + control_token=os.environ.get(CONTROL_TOKEN_ENV) or None, ) gateway.run() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 283fd90a7..0511c9398 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -9,6 +9,7 @@ import logging import os import posixpath +import re import socket import subprocess from collections.abc import Iterable @@ -62,12 +63,14 @@ METHOD_ADD, METHOD_GET, METHOD_RESET, + PCAP_MEDIA_TYPE, PROTOCOL_VERSION, RPC_PATH, STATE_TOOL_NAME_CLASHES_KEY, STEP_SOURCE_POST_TURN_HOOK, TOOL_DISABLE_ACTION, TOOL_ENABLE_ACTION, + TRAFFIC_EXTENSION, WEBSITE_BROWSE_URL_PARAM, WELL_KNOWN_PATH, WILDCARD, @@ -102,6 +105,9 @@ # RFC 9110's safe methods: a request using one asks for no change, so the proxy measures no write window for it. _SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE"}) _STEP_API = {"source": "step_api"} +_CAPTURE_FILE = re.compile(r"(?P[A-Za-z0-9][A-Za-z0-9_.-]*)\.pcap(?P[0-9]*)") +_CONTROL_SCOPE_KEY = "agentenv.control" +_AUTHORIZATION = b"authorization" class ToolNameClash(ValueError): @@ -142,11 +148,19 @@ def __init__( service_db_url: str | None = None, role_key: str | None = None, provider_headers: Iterable[str] = (), + traffic_dir: str | None = None, + control_token: str | None = None, ): """`role_key` is the deploy's secret: with one, a request is believed to act as a role other than `default` only when it carries the key beside the role; without one, as before it, the role header alone is believed. `provider_headers` names the headers the sandbox provider in front of the gateway is sent, its credentials, - which the gateway takes off every request it forwards to a service behind it.""" + which the gateway takes off every request it forwards to a service behind it. + + `traffic_dir` is the folder the deploy's capture sidecars write their websites' pcap files to, mounted here + read-only; with one, the gateway serves traffic/v1 from it. `control_token` is the bearer its framework-only + routes take (traffic/v1's): a request carrying it has it taken off before anything else sees it, so no site or + backing server it is proxied to ever reads it; without one, those routes are open, as an env's are without + ``AGENTENV_CONTROL_TOKEN``.""" self.host = host self.port = port self.server_name = server_name @@ -158,6 +172,8 @@ def __init__( self._service_db_url = service_db_url self._role_key = role_key or None self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) + self._traffic_dir = Path(traffic_dir) if traffic_dir else None + self._control_bearer = f"Bearer {control_token}".encode() if control_token else None self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: @@ -414,6 +430,15 @@ async def notifications_clear(request: Request) -> Response: async def env_trajectory(request: Request) -> Response: return self._serve_trajectory() + if self._traffic_dir is not None: + @self._mcp.custom_route("/traffic/captures", methods=["GET"]) + async def traffic_captures(request: Request) -> Response: + return self._refuse_uncontrolled(request) or await asyncio.to_thread(self._serve_captures) + + @self._mcp.custom_route("/traffic/capture", methods=["GET"]) + async def traffic_capture(request: Request) -> Response: + return self._refuse_uncontrolled(request) or self._serve_capture(request.query_params.get("name", "")) + @self._mcp.custom_route("/changes", methods=["GET"]) async def env_changes(request: Request) -> Response: return await self._serve_changes(request) @@ -457,6 +482,8 @@ def _asgi_app(self): inner.router.lifespan_context = self._process_lifespan(inner.router.lifespan_context) async def app(scope, receive, send): + if scope["type"] in ("http", "websocket") and self._control_bearer is not None: + scope = self._take_control_token(scope) if scope["type"] != "http": await inner(scope, receive, send) return @@ -475,6 +502,22 @@ async def app(scope, receive, send): return app + def _take_control_token(self, scope: dict) -> dict: + """``scope`` without any Authorization header that carries the gateway's control token, and marked when one + did, so a framework-only route can tell and nothing the request is proxied to ever reads the token.""" + headers = scope.get("headers", ()) + kept = [(name, value) for name, value in headers + if not (name == _AUTHORIZATION and hmac.compare_digest(value, self._control_bearer))] + return {**scope, "headers": kept, _CONTROL_SCOPE_KEY: len(kept) != len(headers)} + + def _refuse_uncontrolled(self, request: Request) -> Response | None: + """The refusal of a framework-only route's request that didn't carry the control token, or None: with no token, + every request is let through.""" + if self._control_bearer is None or request.scope.get(_CONTROL_SCOPE_KEY): + return None + return self._json_response({"ok": False, "error": {"code": "unauthorized", "message": + "this route is agent-env's own: it takes the env's control token"}}, 401) + @asynccontextmanager async def _session_lifespan(self, mcp: FastMCP): """Once per inbound MCP session: the session's key joins the caller, and a connect gets to heal a failed @@ -821,7 +864,7 @@ async def _fetch_child_card(self, client: httpx.AsyncClient, key: str, base_url: return None if child is None else self._rewrite_child_card(key, child) def _gateway_extensions(self) -> list[dict]: - return GATEWAY_EXTENSIONS + return [*GATEWAY_EXTENSIONS, TRAFFIC_EXTENSION] if self._traffic_dir is not None else GATEWAY_EXTENSIONS async def _serve_env_card(self) -> Response: """Compose the env card: the gateway's own extensions + each backing server's card nested @@ -1566,6 +1609,54 @@ def _read_changelog(self, after_id: int, count: int) -> list[dict]: row["timestamp"] = row["timestamp"].isoformat() return rows + def _capture_files(self) -> list[tuple[Path, re.Match]]: + """Each capture file in the traffic folder, as tcpdump's rotation names it, with its name's parts; a link is + none, so nothing but what the captures wrote is ever served.""" + try: + entries = list(self._traffic_dir.iterdir()) + except FileNotFoundError: + return [] + return [(path, match) for path in entries + if (match := _CAPTURE_FILE.fullmatch(path.name)) and path.is_file() and not path.is_symlink()] + + def _serve_captures(self) -> Response: + """traffic/v1's list: each capture file's name, the container it holds, its size and when it was last written, + oldest first, so a container's files read in the order their packets were captured.""" + found = [] + for path, match in self._capture_files(): + try: + stat = path.stat() + except FileNotFoundError: + continue + found.append(((match["container"], stat.st_mtime_ns, path.name), { + "name": path.name, "container": match["container"], "bytes": stat.st_size, + "modified_at": datetime.fromtimestamp(stat.st_mtime, timezone.utc).isoformat()})) + captures = [capture for _, capture in sorted(found, key=lambda entry: entry[0])] + return self._json_response({"format": "pcap", "captures": captures}, 200) + + def _serve_capture(self, name: str) -> Response: + """traffic/v1's get: the capture file ``name``, one ``list`` names, streamed as it stands; tcpdump writes each + packet as it captures it, so a file still being written ends on a whole packet, or at worst one torn.""" + if not _CAPTURE_FILE.fullmatch(name): + return self._json_response({"ok": False, "error": {"code": "invalid_params", + "message": f"{name!r} is not a capture's name, as list gives it"}}, 400) + path = self._traffic_dir / name + try: + if not path.is_file() or path.is_symlink(): + raise FileNotFoundError(name) + handle = os.fdopen(os.open(path, os.O_RDONLY | os.O_NOFOLLOW), "rb") + except OSError: + return self._json_response({"ok": False, "error": {"code": "not_found", + "message": f"no capture named {name!r}"}}, 404) + + def chunks(): + with handle: + while chunk := handle.read(65536): + yield chunk + + return StreamingResponse(chunks(), media_type=PCAP_MEDIA_TYPE, + headers={"Content-Disposition": f'attachment; filename="{name}"'}) + def _serve_trajectory(self) -> Response: """Stream the whole env trajectory as JSONL; only a mid-write tail line can arrive torn.""" try: diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 58fee59ab..34f227c38 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -23,10 +23,12 @@ from agent_env.env.env import DeployedGatewayEnv, _mcp_url from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode from agent_env.env.gateway.constants import ( + CONTROL_TOKEN_ENV, DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV, GATEWAY_ROLE_KEY_ENV, STATE_TOOL_NAME_CLASHES_KEY, + TRAFFIC_CAPTURE_DIR_ENV, WEBSITE_ROLE_KEY_ENV, WELL_KNOWN_PATH, random_mcp_server_name, @@ -95,6 +97,13 @@ WEBSITE_READY_TIMEOUT_S = 300.0 WEBSITE_READY_POLL_S = 2.0 +TRAFFIC_CAPTURE_SUFFIX = "traffic" +TRAFFIC_VOLUME = "traffic-captures" +TRAFFIC_GATEWAY_DIR = "/var/lib/agentenv/traffic" +TRAFFIC_SIDECAR_DIR = "/captures" +TRAFFIC_CAPTURE_FILE_MB = 20 +TRAFFIC_CAPTURE_FILES = 10 + @dataclass class MCPServerConfig: @@ -123,6 +132,7 @@ class DeployedGateway: environment_card: dict | None = None environment_card_read_at_utc: str | None = None role_key: str | None = None + control_token: str | None = None @dataclass @@ -298,6 +308,8 @@ def create_docker_compose( extra_hosts: tuple[str, ...] = (), mcp_server_name: str | None = None, role_key: str | None = None, + capture_traffic: bool = False, + control_token: str | None = None, ) -> str: """Generate docker-compose.yml content for a gateway deployment onto a VM/laptop/arbitrary machine. @@ -311,6 +323,11 @@ def create_docker_compose( website_configs: Optional list of website configurations (frontend + backend pairs). role_key: The deploy's role key, given to the gateway as AGENTENV_ROLE_KEY and to each website server as WEBSITE_ROLE_KEY; without one, both believe the AgentEnv-Role header as before it. + capture_traffic: Add a capture service beside each website's server container, and its deprecated + frontend container when it has one (see ``traffic_capture_services``): tcpdump, from the gateway's + image, in that container's network namespace, the only service holding NET_RAW and NET_ADMIN, writing + rotating pcap files to a volume the gateway mounts read-only and serves as traffic/v1. + control_token: The token the gateway's framework-only routes take, given to it as AGENTENV_CONTROL_TOKEN. Returns: docker-compose.yml content as a string. @@ -398,7 +415,8 @@ def create_docker_compose( env-network: driver: bridge """ - problems = _compose_name_problems(mcp_servers, website_configs, sidecars, images=(gateway_image,)) + problems = _compose_name_problems(mcp_servers, website_configs, sidecars, images=(gateway_image,), + capture_traffic=capture_traffic) if mcp_server_name and ("\n" in mcp_server_name or "\r" in mcp_server_name): problems.append(f"MCP server name {mcp_server_name!r} must be one line") if problems: @@ -474,6 +492,9 @@ def create_docker_compose( dep_lines += [f" {name}:", " condition: service_healthy"] for backend_name in website_backend_servers.values(): dep_lines += [f" {backend_name}:", " condition: service_healthy"] + captures = traffic_capture_services(website_configs) if capture_traffic else {} + for capture_name in captures: + dep_lines += [f" {capture_name}:", " condition: service_started"] if dep_lines: lines.append(" depends_on:") lines.extend(dep_lines) @@ -512,6 +533,12 @@ def create_docker_compose( gateway_env.append(f" - {GATEWAY_ROLE_KEY_ENV}={role_key}") if provider_headers := sandbox_request_header_names(): gateway_env.append(f" - {GATEWAY_PROVIDER_HEADERS_ENV}={','.join(provider_headers)}") + if control_token: + gateway_env.append(f" - {CONTROL_TOKEN_ENV}={control_token}") + gateway_volumes = [] + if capture_traffic: + gateway_env.append(f" - {TRAFFIC_CAPTURE_DIR_ENV}={TRAFFIC_GATEWAY_DIR}") + gateway_volumes = [" volumes:", f" - {_traffic_mount(TRAFFIC_GATEWAY_DIR, read_only=True)}"] publish = host_port or (lambda port: port) exposed_gateway_ports = [] if expose_gateway_port: @@ -526,6 +553,7 @@ def create_docker_compose( *gateway_env, *exposed_gateway_ports, *gateway_extra_hosts, + *gateway_volumes, " restart: on-failure", " healthcheck:", f' test: ["CMD-SHELL", "python3 -c \\"import socket; s=socket.create_connection((\'localhost\',{gateway_port}),2); s.close()\\""]', @@ -583,6 +611,9 @@ def create_docker_compose( "", ]) + for capture_name, container in captures.items(): + lines.extend([f" {capture_name}:", *_traffic_capture_service(gateway_image, container), ""]) + # Add sidecar services. # Secrets go through the compose environment block, never a docker-run command line. for sc in sidecars or []: @@ -609,6 +640,8 @@ def create_docker_compose( " env-network:", " driver: bridge", ]) + if capture_traffic: + lines.extend(["volumes:", f" {TRAFFIC_VOLUME}: {{}}"]) return "\n".join(lines) async def deploy( @@ -627,18 +660,24 @@ async def deploy( artifact_id: str | None = None, artifact_version: int | None = None, keyed_roles: bool = False, + capture_traffic: bool = False, ) -> DeployedGatewayEnv: """A gateway in front of the env, with the state store it acquires (local Postgres by default); returns the record. A MultiEnv's gateway fronts its MCP servers and websites, a WebsiteEnv's its website, and any other env's its one MCP server. - With ``keyed_roles``, the deploy draws a role key (see ``create_gateway``). + With ``keyed_roles``, the deploy draws a role key (see ``create_gateway``). With ``capture_traffic``, it captures + its websites' traffic (see ``create_gateway``), and refuses an env that serves no website. """ from agent_env.config import get_config, service_db_env_id from agent_env.env.env import Env from agent_env.providers.env_state import LOCAL_POSTGRES_STATE_TYPE, acquire_state_for_deploy topology = _gateway_topology(env) - if problems := _compose_name_problems(topology.mcp_servers, topology.website_configs): + if capture_traffic and not topology.website_configs: + raise ValueError(f"Can't capture {env.id!r}'s traffic: capture_traffic records what a website's containers " + "send and receive, and the env serves no website") + if problems := _compose_name_problems(topology.mcp_servers, topology.website_configs, + capture_traffic=capture_traffic): raise ValueError(f"Can't deploy {env.id!r} behind a gateway: {'; '.join(problems)}") state_instance = await acquire_state_for_deploy( env_state_type=env_state_type, ttl_seconds=ttl_seconds, name_hint=env.id, env_state_instance_id=env_state_instance_id, @@ -650,7 +689,7 @@ async def deploy( mcp_server_images=topology.mcp_server_images, website_configs=topology.website_configs, website_images=topology.website_images, mcp_server_name=topology.mcp_server_name, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, attribution=attribution, - keyed_roles=keyed_roles, + keyed_roles=keyed_roles, **({"capture_traffic": True} if capture_traffic else {}), ) return DeployedGatewayEnv( env_id=env.id, @@ -668,6 +707,7 @@ async def deploy( gateway_mode=gateway_mode.value, env_state_instance_ids=deployed.env_state_instance_ids, role_key=deployed.role_key, + control_token=deployed.control_token, ) async def create_gateway( @@ -691,6 +731,7 @@ async def create_gateway( *, attribution: Optional[Attribution] = None, keyed_roles: bool = False, + capture_traffic: bool = False, ) -> DeployedGateway: """Create and start a gateway from explicit topology, for a caller that builds its own record rather than calling ``deploy()``. @@ -708,22 +749,29 @@ async def create_gateway( ``keyed_roles`` draws a secret role key for the deploy, given to the gateway and each website server and kept on the record: the gateway then believes a role other than ``default`` only beside it, so every caller acting as one must send it. Off, as before the key, the role header alone is believed. + + ``capture_traffic`` records every packet each website's server and frontend containers send and receive, with + tcpdump in a capture service beside each, as rotating pcap files the gateway serves as traffic/v1, behind a + control token the deploy draws for it and keeps on the record (``control_token``). VM deploys only, since only + they serve websites; off, the stack is as before it. """ return await self._launch( sandbox_provider, env_id, state_instance=state_instance, mcp_servers=mcp_servers, mcp_server_images=mcp_server_images, gateway_port=gateway_port, website_configs=website_configs, website_images=website_images, gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, mcp_server_name=mcp_server_name, - attribution=attribution, keyed_roles=keyed_roles, + attribution=attribution, keyed_roles=keyed_roles, **({"capture_traffic": True} if capture_traffic else {}), ) async def _launch(self, sandbox_provider: SandboxProvider, env_id: str | None, *, mcp_server_name: str | None = None, - keyed_roles: bool = False, **gateway) -> DeployedGateway: - """One gateway, deployed on the sandbox provider or each chained member in turn, then its tools probed.""" + keyed_roles: bool = False, capture_traffic: bool = False, **gateway) -> DeployedGateway: + """One gateway, deployed on the sandbox provider or each chained member in turn, then its tools probed. Its role + key and control token are drawn once, so every chained attempt holds the same.""" attempt = functools.partial( self._deploy_gateway, **gateway, mcp_server_name=mcp_server_name or random_mcp_server_name(), # drawn once, so every chained attempt presents the same name role_key=secrets.token_urlsafe(32) if keyed_roles else None, + **({"capture_traffic": True, "control_token": secrets.token_urlsafe(32)} if capture_traffic else {}), ) earlier_sandbox = self._sandbox result = await self._run(sandbox_provider, env_id, attempt) @@ -776,6 +824,8 @@ async def _deploy_via_vm( sidecars: list[SidecarConfig] | None = None, mcp_server_name: str | None = None, role_key: str | None = None, + capture_traffic: bool = False, + control_token: str | None = None, ) -> DeployedGateway: from agent_env.env.env import Env from agent_env.config import gateway_env_id, get_config, website_browser_env_id @@ -868,6 +918,7 @@ async def _stand_up(_spec): extra_hosts=sandbox.extra_hosts, mcp_server_name=mcp_server_name, role_key=role_key, + **({"capture_traffic": True, "control_token": control_token} if capture_traffic else {}), ) logger.info(f"Generated docker-compose.yml:\n{_redact_compose_secrets(compose_content)}") @@ -916,6 +967,8 @@ async def _stand_up(_spec): if not gateway_ready: down = await self._down_services(sandbox) raise RuntimeError("Gateway did not become ready in time" + (f":\n{down}" if down else "")) + if capture_traffic: + await self._refuse_stopped_captures(sandbox, list(traffic_capture_services(website_configs))) # Check container status only on the sidecars the provider actually rendered sidecar_names = ( @@ -964,6 +1017,7 @@ async def _stand_up(_spec): website_frontend_urls=website_frontend_urls, mcp_server_name=mcp_server_name, env_state_instance_ids=[state_instance.instance_id], environment_card=card, environment_card_read_at_utc=card_read_at_utc, role_key=role_key, + control_token=control_token, ) async def _deploy_via_containers( @@ -1211,6 +1265,32 @@ async def _down_services(self, sandbox: VmSandbox) -> str | None: logger.warning(f"Can't list the stack's services: no answer in {_DOWN_SERVICES_TIMEOUT_S:.0f}s") return None + async def _refuse_stopped_captures(self, sandbox: VmSandbox, capture_services: list[str]) -> None: + """Refuse a deploy whose capture services aren't all running once the gateway is up, as one is whose gateway + image has no tcpdump: it would record nothing. A lookup that fails is a warning, since it says nothing of them.""" + try: + exit_code, stdout, stderr = await sandbox.exec_with_output( + "sudo", "docker", "compose", "-f", DOCKER_COMPOSE_PATH, "ps", "-a", "--format", "json") + except Exception as e: + logger.warning(f"Can't check the traffic captures are running: {e}") + return + if exit_code != 0: + logger.warning(f"Can't check the traffic captures are running: docker compose ps exited {exit_code}: " + f"{stderr.strip()[-500:]}") + return + rows = {row.get("Service"): row for row in _compose_ps_rows(stdout)} + stopped = [] + for name in capture_services: + row = rows.get(name) + if row is None or row.get("State") != "running": + tail = await self._log_tail(sandbox, str((row or {}).get("ID") or "")) + why = _why_down(row) if row is not None else "has no container" + stopped.append(f"{name} {why or row.get('State')}" + (f":\n{tail}" if tail else "")) + if stopped: + raise RuntimeError("The env's traffic captures aren't running, so its traffic wouldn't be recorded (a gateway " + "image built before captures has no tcpdump: put the gateway env again):\n" + + "\n".join(stopped)) + async def _read_down_services(self, sandbox: VmSandbox) -> str | None: try: exit_code, stdout, stderr = await sandbox.exec_with_output( @@ -1415,9 +1495,11 @@ async def _deploy_gateway( mcp_server_name: str | None = None, attribution: Optional[Attribution] = None, role_key: str | None = None, + capture_traffic: bool = False, + control_token: str | None = None, ) -> DeployedGateway: """One gateway deploy on one (unchained) sandbox provider: a VM on one that creates VMs, else a container per - server on its private network.""" + server on its private network, which captures no traffic.""" from agent_env.providers.env_state import LocalPostgresStateProvider, build_state_provider # Every attempt starts from the caller's store: a failed attempt's close() cleared it. @@ -1429,6 +1511,9 @@ async def _deploy_gateway( ) deploy_kwargs = {"attribution": dict(attribution or {})} if not self.in_vm(sandbox_provider): + if capture_traffic: + raise ValueError(f"Can't capture traffic on {type(sandbox_provider).__name__}, which runs the gateway in " + "containers: a capture runs beside its website's container on the gateway's VM") if problem := self.sandbox_problem(sandbox_provider, websites=bool(website_configs)): raise ValueError(f"Can't deploy a gateway on {type(sandbox_provider).__name__}, which {problem}") if existing_sandbox is not None: @@ -1449,6 +1534,7 @@ async def _deploy_gateway( gateway_mode=gateway_mode, ttl_seconds=ttl_seconds, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, existing_sandbox=existing_sandbox, sidecars=sidecars, mcp_server_name=mcp_server_name, role_key=role_key, **deploy_kwargs, + **({"capture_traffic": True, "control_token": control_token} if capture_traffic else {}), ) except BaseException: await self._terminate_new_sandbox(earlier_sandbox=earlier_sandbox, existing_sandbox=existing_sandbox) @@ -1561,6 +1647,47 @@ def compose_env_entry(name: str, value: str) -> str: return '"' + "".join(escaped) + '"' +def traffic_capture_services(website_configs: list[WebsiteConfig] | None) -> dict[str, str]: + """The capture service a deploy capturing its websites' traffic runs for each website container, by name, each + with the compose service of the container whose network namespace it captures in: every website's server, and its + deprecated frontend when it has one.""" + containers = [f"{wc.environment_name}-{suffix}" for wc in website_configs or [] + for suffix in (AGENT_ENV_WEBSITE_BACKEND_SUFFIX, AGENT_ENV_WEBSITE_FRONTEND_SUFFIX) + if suffix == AGENT_ENV_WEBSITE_BACKEND_SUFFIX or wc.frontend_image] + return {f"{container}-{TRAFFIC_CAPTURE_SUFFIX}": container for container in containers} + + +def _traffic_mount(path: str, *, read_only: bool = False) -> str: + """The compose ``volumes:`` entry mounting ``TRAFFIC_VOLUME`` at ``path``.""" + return f"{TRAFFIC_VOLUME}:{path}" + (":ro" if read_only else "") + + +def _traffic_capture_service(gateway_image: str, container: str) -> list[str]: + """The compose lines of the service capturing ``container``'s traffic: tcpdump, from the gateway's image, on the + container's own interface, every packet whole and written as it's captured, rotating through + ``TRAFFIC_CAPTURE_FILES`` files of ``TRAFFIC_CAPTURE_FILE_MB`` million bytes named for the container, its oldest + overwritten once they're full, so a long run holds at most their product per container on the VM's disk. It joins + the container's network namespace, so it sees what the container sends and receives and nothing else, and is the + only service given the capabilities a capture needs. It writes to ``TRAFFIC_VOLUME``, which the gateway mounts + read-only and no website container mounts at all.""" + argv = ["-i", "eth0", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), + "-W", str(TRAFFIC_CAPTURE_FILES), "-w", f"{TRAFFIC_SIDECAR_DIR}/{container}.pcap"] + return [ + f" image: {gateway_image}", + f' network_mode: "service:{container}"', + " depends_on:", + f" - {container}", + " cap_add:", + " - NET_RAW", + " - NET_ADMIN", + ' entrypoint: ["tcpdump"]', + f" command: {json.dumps(argv)}", + " volumes:", + f" - {_traffic_mount(TRAFFIC_SIDECAR_DIR)}", + " restart: unless-stopped", + ] + + def website_config(env: WebsiteEnv) -> WebsiteConfig: """A website env as the gateway runs it: its server image, any deprecated frontend image, and its health path.""" frontend = env.frontend_docker_image_artifact @@ -1608,10 +1735,12 @@ def environment_name_problem(name: str) -> str | None: def _compose_name_problems(mcp_servers: list[MCPServerConfig], website_configs: list[WebsiteConfig] | None, - sidecars: list[SidecarConfig] | None = None, *, images: tuple[str, ...] = ()) -> list[str]: + sidecars: list[SidecarConfig] | None = None, *, images: tuple[str, ...] = (), + capture_traffic: bool = False) -> list[str]: """Why the environment names and images of a gateway deploy can't go into its compose file and shell lines: one can't be a service name, two of its services would share one (a website's are suffixed, so a website may share an - MCP server's name), or one is a service every gateway deploy runs.""" + MCP server's name; so are the capture services of a deploy with ``capture_traffic``), or one is a service every + gateway deploy runs.""" names = [c.environment_name for c in [*mcp_servers, *(website_configs or []), *(sidecars or [])]] refs = [*images, *(s.image for s in mcp_servers), *(s.image for s in sidecars or []), *(ref for w in website_configs or [] for ref in (w.backend_image, w.frontend_image) if ref is not None)] @@ -1619,7 +1748,8 @@ def _compose_name_problems(mcp_servers: list[MCPServerConfig], website_configs: if problem] services = Counter([*(c.environment_name for c in [*mcp_servers, *(sidecars or [])]), *(f"{w.environment_name}-{suffix}" for w in website_configs or [] - for suffix in (AGENT_ENV_WEBSITE_BACKEND_SUFFIX, AGENT_ENV_WEBSITE_FRONTEND_SUFFIX))]) + for suffix in (AGENT_ENV_WEBSITE_BACKEND_SUFFIX, AGENT_ENV_WEBSITE_FRONTEND_SUFFIX)), + *(traffic_capture_services(website_configs) if capture_traffic else ())]) problems += [f"two of its services would both be named {name!r}" for name, count in services.items() if count > 1] problems += [f"environment name {name!r} is one a gateway deploy names its own containers " f"({', '.join(sorted(GATEWAY_SERVICE_NAMES))})" for name in services if name in GATEWAY_SERVICE_NAMES] diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index fc5a77536..875d09946 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -148,7 +148,8 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s UNAFFECTED = frozenset({ # They name no agent. - "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "combine_universe_verdicts", "configure_scheduled_checks", + "aggregate_verifiers", "apply_server_config", "build_mcp_cli", "collect_traffic_capture", "combine_universe_verdicts", + "configure_scheduled_checks", "deploy_env", "deploy_sandbox", "end_computer_recording", "start_computer_recording", "env_outcome_verifier", "modify_env_tool_access", "reset_env", "review", "run_container_unit_tests_verifier", "run_docker_container", "sign_in_browser", "snapshot_env", "snapshot_computer_env", "sync_env_clock", diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index 9f11627c4..a0aa9e709 100644 --- a/src/agent_env/task/teardown.py +++ b/src/agent_env/task/teardown.py @@ -3,7 +3,8 @@ ``teardown_run`` takes down everything a run's context records: its bare sandboxes, its envs and its agents. It removes compute only; the instance, its outputs and the artifacts it collected stay. What has to reach the stores goes first: the changelogs staged on its agents (``before_teardown``), the recordings the run left open, ended with fresh -grants whatever env serves them, and what each env saves as it closes (``Env.close``); a recording still open after +grants whatever env serves them, the traffic captures of each env deployed with ``capture_traffic`` that the run didn't +collect, and what each env saves as it closes (``Env.close``); a recording still open after that is registered from what its namespace holds. Then every sandbox is terminated by id, and a local one's work folder goes once it's down. ``agent-env run`` and ``agent-env eval run`` call it after each run, and the ``teardown_sandboxes`` step takes down part of a run the same way (``tear_down``). ``Task.run`` never does, so a @@ -31,6 +32,7 @@ from agent_env.task.interrupts import tearing_down from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.computer_recording import end_open_recordings, register_open_recordings +from agent_env.task_step.task_steps.traffic_capture import collect_uncollected_captures logger = logging.getLogger(__name__) @@ -116,7 +118,8 @@ async def teardown_run(context: TaskStepContext, *, timeout: float = TERMINATE_T async def tear_down(context: TaskStepContext, sandboxes: Sequence[RecordedSandbox], envs: Sequence[DeployedEnv], *, timeout: float) -> TeardownReport: - """End the recordings the run left open on ``envs`` with fresh grants, close each env still up, register any + """End the recordings the run left open on ``envs`` with fresh grants, collect the traffic captures of each of them + that captures and that the run didn't collect, close each env still up, register any recording still open from its namespace, then terminate each of ``sandboxes`` by id, all at once and each within ``timeout`` seconds, and remove a local one's work folder once it's down. A close that fails is logged: its env's sandboxes are terminated by id regardless, which leaves one the close took down as it is. Cancelling it returns what it reached, with the rest in ``left``.""" @@ -129,8 +132,13 @@ async def tear_down(context: TaskStepContext, sandboxes: Sequence[RecordedSandbo await end_open_recordings(context, env_ids) except TimeoutError: logger.warning("teardown: the run's open recordings weren't all ended within %gs", timeout) - await asyncio.gather(*(_close(env, timeout) for env in envs - if any(sandbox_id in up for sandbox_id, _ in env_sandbox_ids(env)))) + still_up = [env for env in envs if any(sandbox_id in up for sandbox_id, _ in env_sandbox_ids(env))] + try: + async with asyncio.timeout(timeout): + await collect_uncollected_captures(context, still_up) + except TimeoutError: + logger.warning("teardown: the run's uncollected traffic captures weren't all collected within %gs", timeout) + await asyncio.gather(*(_close(env, timeout) for env in still_up)) await asyncio.to_thread(register_open_recordings, context, env_ids) except asyncio.CancelledError: return TeardownReport(left=pending) diff --git a/src/agent_env/task_step/registry.py b/src/agent_env/task_step/registry.py index 0c8ebd7a5..f482a64be 100644 --- a/src/agent_env/task_step/registry.py +++ b/src/agent_env/task_step/registry.py @@ -39,6 +39,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep from agent_env.task_step.task_steps.reset_env import ResetEnvTaskStep from agent_env.task_step.task_steps.teardown_sandboxes import TeardownSandboxesTaskStep + from agent_env.task_step.task_steps.traffic_capture import CollectTrafficCaptureTaskStep from agent_env.task_step.task_steps.verifiers.env_outcome_verifier import EnvOutcomeVerifierTaskStep from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep from agent_env.task_step.task_steps.modify_env_tool_access import ModifyEnvToolAccessStep @@ -91,6 +92,7 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: _get_type(CollectArtifactsTaskStep): CollectArtifactsTaskStep, _get_type(StartComputerRecordingTaskStep): StartComputerRecordingTaskStep, _get_type(EndComputerRecordingTaskStep): EndComputerRecordingTaskStep, + _get_type(CollectTrafficCaptureTaskStep): CollectTrafficCaptureTaskStep, _get_type(DeployAgentTaskStep): DeployAgentTaskStep, _get_type(DeployEnvTaskStep): DeployEnvTaskStep, _get_type(DeployHumanAgentTaskStep): DeployHumanAgentTaskStep, diff --git a/src/agent_env/task_step/task_steps/deploy_env.py b/src/agent_env/task_step/task_steps/deploy_env.py index 22530034f..d12609f75 100644 --- a/src/agent_env/task_step/task_steps/deploy_env.py +++ b/src/agent_env/task_step/task_steps/deploy_env.py @@ -49,6 +49,7 @@ def __init__( artifact_version: Optional[int] = None, env_params: Optional[dict] = None, keyed_roles: bool = False, + capture_traffic: bool = False, ): super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) self.env_id = env_id @@ -68,6 +69,7 @@ def __init__( # The env type's own deploy options, which this step passes to its deploy() without reading them. self.env_params = env_params self.keyed_roles = keyed_roles + self.capture_traffic = capture_traffic def run_state(self, user_overrides: Mapping[str, Any]) -> tuple[str | None, str | None]: """The env state type and attached state instance a run with ``user_overrides`` deploys on: each the run's, @@ -92,6 +94,8 @@ def to_dict(self) -> dict: base["metadata"] = self.metadata base["env_params"] = self.env_params base["keyed_roles"] = self.keyed_roles + if self.capture_traffic: + base["capture_traffic"] = True return base @classmethod @@ -113,6 +117,7 @@ def from_dict(cls, data: dict) -> DeployEnvTaskStep: metadata=dict(data.get("metadata") or {}), env_params=data.get("env_params"), keyed_roles=bool(data.get("keyed_roles", False)), + capture_traffic=bool(data.get("capture_traffic", False)), ) def preflight(self) -> list[str]: @@ -138,7 +143,8 @@ def preflight(self) -> list[str]: def _env_params_option(self) -> dict: return {**({"env_params": self.env_params} if self.env_params else {}), - **({"keyed_roles": True} if self.keyed_roles else {})} + **({"keyed_roles": True} if self.keyed_roles else {}), + **({"capture_traffic": True} if self.capture_traffic else {})} async def execute(self, context: TaskStepContext) -> TaskStepContext: from agent_env.env.env import Env @@ -197,6 +203,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if "keyed_roles" not in params and not accepts_kwargs: raise ValueError(f"deploy_env '{self.id}': {env.type} envs take no keyed_roles") deploy_kwargs["keyed_roles"] = True + if self.capture_traffic: + if "capture_traffic" not in params and not accepts_kwargs: + raise ValueError(f"deploy_env '{self.id}': {env.type} envs take no capture_traffic") + deploy_kwargs["capture_traffic"] = True deployed_env = await env.deploy(**deploy_kwargs) # create_instance runs inside env.deploy(), so an annotation set after it # returns has to be persisted explicitly to reach the instance record. diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py new file mode 100644 index 000000000..aa2e0b65d --- /dev/null +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -0,0 +1,147 @@ +"""Collect a website env's traffic captures. A website or multi env deployed with ``capture_traffic`` records every +packet its websites' containers send and receive as rotating pcap files, which its gateway's card serves as traffic/v1 +behind the deploy's control token. ``collect_traffic_capture`` stores each as a FileArtifact and registers them as one +FileArtifactUniverse, recorded in ``context.metadata[TRAFFIC_KEY][env_id]``. Teardown collects any capturing env a run +didn't, before the env closes (``collect_uncollected_captures``), so a capture isn't lost with its VM.""" +from __future__ import annotations + +import asyncio +import logging +import os +import tempfile +import time +from collections.abc import Sequence +from typing import ClassVar, Optional + +from agentenv_protocol.traffic import ( + PCAP_MEDIA_TYPE, + TRAFFIC_EXTENSION_URI, + TRAFFIC_LIST_OP, + TrafficCapture, + fetch_traffic_capture, + list_traffic_captures, + serves_traffic, +) + +from agent_env.artifact import FileArtifact +from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse +from agent_env.artifact.store import get_artifact_store +from agent_env.config import get_config +from agent_env.entity.refs import EntityRef +from agent_env.env.env import DeployedEnv, EnvCapabilityUnsupported +from agent_env.store.ids import derive_id, is_local_id, validate_local_id +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency + +logger = logging.getLogger(__name__) + +#: ``context.metadata[TRAFFIC_KEY][env_id]``: each env whose captures a run collected, or tried to. +TRAFFIC_KEY = "traffic_captures" +_ARTIFACT_TYPE = "traffic_captures" + + +class CollectTrafficCaptureTaskStep(TaskStep): + """Store the traffic captures of the deployed env ``env_id`` as pcap FileArtifacts. Best effort by default: a + failure is recorded and logged, and fails the task only with ``fail_task_on_error``.""" + + type: ClassVar[str] = "collect_traffic_capture" + entity_refs = (EntityRef.env("env_id"),) + + def __init__(self, id: str, version: Optional[int], env_id: str, + depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = False): + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + if not env_id: + raise ValueError("collect_traffic_capture needs the env_id of an env deployed with capture_traffic") + self.env_id = env_id + + def to_dict(self) -> dict: + return {**super().to_dict(), "env_id": self.env_id} + + @classmethod + def from_dict(cls, data: dict) -> CollectTrafficCaptureTaskStep: + base = cls._base_from_dict(data) + base["fail_task_on_error"] = data.get("fail_task_on_error", False) + return cls(**base, env_id=data["env_id"]) + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) + if deployed is None: + error = RuntimeError(f"env {self.env_id!r} isn't deployed in this run") + context.metadata.setdefault(TRAFFIC_KEY, {})[self.env_id] = _failed(error) + logger.warning("collect_traffic_capture: %s", error) + if self.fail_task_on_error: + raise error + return context + entry = await collect_captures(context, deployed) + if self.fail_task_on_error and entry["state"] != "collected": + raise RuntimeError(f"env {self.env_id!r}'s traffic captures weren't collected: {entry.get('error')}") + return context + + +def captures_traffic(deployed: DeployedEnv) -> bool: + """Whether the stored card of ``deployed`` serves traffic/v1, as a deploy with ``capture_traffic`` does; no + network.""" + return isinstance(deployed, DeployedEnv) and serves_traffic(deployed.environment_card) + + +async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[DeployedEnv]) -> None: + """Collect the captures of each of ``envs`` that serves traffic/v1 and whose captures the run hasn't collected or + tried to, as teardown does before the envs close; ``collect_captures`` records what happened to each.""" + tried = context.metadata.get(TRAFFIC_KEY) or {} + for deployed in envs: + if captures_traffic(deployed) and deployed.env_id not in tried: + await collect_captures(context, deployed) + + +async def collect_captures(context: TaskStepContext, deployed: DeployedEnv) -> dict: + """Store each capture ``deployed`` serves as a pcap FileArtifact, register them as one FileArtifactUniverse, and + record it in ``context.metadata[TRAFFIC_KEY]``. Each file is named for its container and its place among that + container's files, oldest first, so they read in the order their packets were captured. Never raises: the entry, + which it also returns, says what happened.""" + entries = context.metadata.setdefault(TRAFFIC_KEY, {}) + try: + if not captures_traffic(deployed): + raise EnvCapabilityUnsupported(TRAFFIC_EXTENSION_URI, TRAFFIC_LIST_OP, deployed.env_id, deployed.instance_id) + listing = await list_traffic_captures(deployed.environment_url, deployed.environment_card, + headers=deployed.control_headers) + entries[deployed.env_id] = await _store(context, deployed, listing.captures) + except Exception as e: + logger.warning("collect_traffic_capture %s: %s", deployed.env_id, e) + entries[deployed.env_id] = _failed(e) + return entries[deployed.env_id] + + +async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list[TrafficCapture]) -> dict: + universe_id = derive_id(context.instance_id or "run", f"traffic-{deployed.env_id}") + if is_local_id(universe_id): + validate_local_id(universe_id) + get_config().check_local_run_write(universe_id) + store, version = get_artifact_store(), int(time.time()) + files, total, seen = {}, 0, {} + for capture in captures: + container = capture.container + seen[container] = seen.get(container, -1) + 1 + filename = f"{container}-{seen[container]:02d}.pcap" + handle, path = tempfile.mkstemp(suffix=".pcap") + os.close(handle) + try: + total += await fetch_traffic_capture(deployed.environment_url, deployed.environment_card, capture.name, + path, headers=deployed.control_headers) + url = await asyncio.to_thread(store.put_object_file, _ARTIFACT_TYPE, universe_id, version, filename, path, + PCAP_MEDIA_TYPE) + finally: + os.unlink(path) + files[filename] = store.put_document(FileArtifact( + id=derive_id(universe_id, filename), filename=filename, content_type=PCAP_MEDIA_TYPE, object_url=url, + description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}")) + artifact = None + if files: + universe = FileArtifactUniverse.put(id=universe_id, file_artifacts=files) + artifact = {"id": universe.id, "version": universe.version} + logger.info("collect_traffic_capture %s: %d capture files, %d bytes", deployed.env_id, len(files), total) + return {"state": "collected", "artifact": artifact, "files": sorted(files), "bytes": total, + "containers": sorted(seen)} + + +def _failed(error: BaseException) -> dict: + return {"state": "failed", "error": f"{type(error).__name__}: {error}"[:500]} diff --git a/tst/unit/env/envs/test_multi_env_name.py b/tst/unit/env/envs/test_multi_env_name.py index 70a9ea224..7cbc360fb 100644 --- a/tst/unit/env/envs/test_multi_env_name.py +++ b/tst/unit/env/envs/test_multi_env_name.py @@ -67,7 +67,8 @@ async def test_a_deploy_records_the_mcp_url_and_name_its_gateway_card_gives(): card = {"name": "crm", "additionalInterfaces": [{"url": "/mcp", "transport": "mcp"}]} result = SimpleNamespace(gateway_url="https://gw.example", mcp_url="https://gw.example/mcp", db_web_url=None, db_mcp_url=None, website_frontend_urls=None, environment_card=card, environment_card_read_at_utc="t", - env_state_instance_ids=[], mcp_server_name="crm", role_key=None) + env_state_instance_ids=[], mcp_server_name="crm", role_key=None, + control_token=None) async def deploy_gateway(self, sandbox_provider, **kwargs): self._sandbox = SimpleNamespace(sandbox_id="gw", type="modal_vm") diff --git a/tst/unit/env/gateway/traffic_test.py b/tst/unit/env/gateway/traffic_test.py new file mode 100644 index 000000000..a5a410133 --- /dev/null +++ b/tst/unit/env/gateway/traffic_test.py @@ -0,0 +1,200 @@ +"""traffic/v1 on the gateway: a deploy that captures its websites' traffic mounts their pcap files into the gateway, +which lists them oldest first and streams each, only to a caller carrying its control token, which it takes off every +request before anything is proxied, so no site ever reads it; a site's route never reaches the captures. A gateway +without a capture folder serves and advertises none of it.""" + +from __future__ import annotations + +import os + +import httpx +import pytest +from agentenv_protocol import traffic as protocol_traffic + +from agent_env.env.gateway import constants as gw_constants +from agent_env.env.gateway import gateway as gw_mod +from agent_env.env.gateway import rest_proxy as rest_proxy_module +from agent_env.env.gateway.constants import EXT_TRAFFIC_URI, GATEWAY_EXTENSIONS, PCAP_MEDIA_TYPE, WELL_KNOWN_PATH +from agent_env.env.gateway.gateway import Gateway + +TOKEN = "control-secret" +_RealAsyncClient = httpx.AsyncClient +_PCAP_HEADER = bytes.fromhex("d4c3b2a1020004000000000000000000ffff000001000000") + + +@pytest.fixture +def traffic_dir(tmp_path, monkeypatch): + monkeypatch.setattr(gw_mod, "GATEWAY_TRAJECTORY_FILE", str(tmp_path / "trajectory.jsonl")) + folder = tmp_path / "traffic" + folder.mkdir() + return folder + + +def _gateway(traffic_dir=None, control_token=TOKEN, **kw) -> Gateway: + return Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + traffic_dir=str(traffic_dir) if traffic_dir else None, control_token=control_token, **kw) + + +def _write(folder, name: str, body: bytes, mtime: int) -> None: + (folder / name).write_bytes(body) + os.utime(folder / name, (mtime, mtime)) + + +async def _get(gw: Gateway, path: str, headers=None, params=None) -> httpx.Response: + transport = httpx.ASGITransport(app=gw._asgi_app()) + async with _RealAsyncClient(transport=transport, base_url="http://gateway") as client: + return await client.get(path, headers=headers, params=params) + + +def _bearer(token: str = TOKEN) -> dict: + return {"Authorization": f"Bearer {token}"} + + +@pytest.mark.asyncio +async def test_list_names_each_capture_by_its_container_oldest_first_and_get_streams_it(traffic_dir): + _write(traffic_dir, "notes-website-backend.pcap1", _PCAP_HEADER + b"second", 2000) + _write(traffic_dir, "notes-website-backend.pcap0", _PCAP_HEADER + b"first", 1000) + _write(traffic_dir, "notes-website-frontend.pcap0", _PCAP_HEADER, 1500) + _write(traffic_dir, "stray.txt", b"not a capture", 1000) + gw = _gateway(traffic_dir) + + listing = (await _get(gw, "/traffic/captures", headers=_bearer())).json() + got = await _get(gw, "/traffic/capture", headers=_bearer(), params={"name": "notes-website-backend.pcap0"}) + + assert listing["format"] == "pcap" + assert [(c["name"], c["container"], c["bytes"]) for c in listing["captures"]] == [ + ("notes-website-backend.pcap0", "notes-website-backend", len(_PCAP_HEADER) + 5), + ("notes-website-backend.pcap1", "notes-website-backend", len(_PCAP_HEADER) + 6), + ("notes-website-frontend.pcap0", "notes-website-frontend", len(_PCAP_HEADER)), + ] + assert listing["captures"][0]["modified_at"].startswith("1970-01-01T00:16:40") + assert (got.status_code, got.headers["content-type"], got.content) == (200, PCAP_MEDIA_TYPE, _PCAP_HEADER + b"first") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name, status", [ + ("../trajectory.jsonl", 400), ("notes-website-backend.pcap0/../../x", 400), ("stray.txt", 400), + (".pcap0", 400), ("", 400), ("gone-website-backend.pcap0", 404), +]) +async def test_get_serves_only_a_capture_file_of_the_folder(traffic_dir, name, status): + (traffic_dir / "stray.txt").write_bytes(b"x") + response = await _get(_gateway(traffic_dir), "/traffic/capture", headers=_bearer(), params={"name": name}) + assert response.status_code == status, response.text + + +@pytest.mark.asyncio +async def test_an_empty_or_missing_folder_lists_no_captures(traffic_dir, tmp_path): + assert (await _get(_gateway(traffic_dir), "/traffic/captures", headers=_bearer())).json()["captures"] == [] + missing = _gateway(tmp_path / "never-mounted") + assert (await _get(missing, "/traffic/captures", headers=_bearer())).json()["captures"] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("headers", [None, _bearer("guess"), {"Authorization": TOKEN}, {"AgentEnv-Role-Key": TOKEN}], + ids=["none", "wrong", "not-bearer", "role-key-header"]) +async def test_each_traffic_route_refuses_a_caller_without_the_control_token(traffic_dir, headers): + _write(traffic_dir, "notes-website-backend.pcap0", _PCAP_HEADER, 1000) + gw = _gateway(traffic_dir) + listed = await _get(gw, "/traffic/captures", headers=headers) + got = await _get(gw, "/traffic/capture", headers=headers, params={"name": "notes-website-backend.pcap0"}) + assert (listed.status_code, got.status_code) == (401, 401) + assert listed.json()["error"]["code"] == "unauthorized" and _PCAP_HEADER not in got.content + + +@pytest.mark.asyncio +async def test_the_inner_app_alone_refuses_the_traffic_routes_of_a_gateway_with_a_token(traffic_dir): + """Only the served app takes the token off a request and vouches for it, so the routes refuse anything else.""" + transport = httpx.ASGITransport(app=_gateway(traffic_dir)._mcp.streamable_http_app()) + async with _RealAsyncClient(transport=transport, base_url="http://gateway") as client: + assert (await client.get("/traffic/captures", headers=_bearer())).status_code == 401 + + +@pytest.mark.asyncio +async def test_a_gateway_with_no_token_serves_its_captures_to_anyone(traffic_dir): + assert (await _get(_gateway(traffic_dir, control_token=None), "/traffic/captures")).status_code == 200 + + +@pytest.mark.asyncio +async def test_the_card_advertises_traffic_only_with_a_capture_folder(traffic_dir): + capturing = (await _get(_gateway(traffic_dir), "/.well-known/agent-env.json")).json() + plain = _gateway(None, control_token=None) + card = (await _get(plain, "/.well-known/agent-env.json")).json() + + [traffic] = [e for e in capturing["capabilities"]["extensions"] if e["uri"] == EXT_TRAFFIC_URI] + assert {name: (op["method"], op["endpoint"]) for name, op in traffic["params"]["methods"].items()} == { + "list": ("GET", "/traffic/captures"), "get": ("GET", "/traffic/capture")} + assert plain._gateway_extensions() is GATEWAY_EXTENSIONS + assert EXT_TRAFFIC_URI not in [e["uri"] for e in card["capabilities"]["extensions"]] + assert (await _get(plain, "/traffic/captures")).status_code == 404 + + +@pytest.mark.asyncio +async def test_the_control_token_never_reaches_a_site_and_a_sites_route_never_reaches_the_captures(traffic_dir, + monkeypatch): + _write(traffic_dir, "notes-website-backend.pcap0", _PCAP_HEADER, 1000) + seen: list[httpx.Request] = [] + + async def body(): + yield b"the site's own page" + + def site(request: httpx.Request) -> httpx.Response: + if request.url.path == WELL_KNOWN_PATH: + return httpx.Response(404) + seen.append(request) + return httpx.Response(200, content=body()) + + transport = httpx.MockTransport(site) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + gw = _gateway(traffic_dir, website_urls={"notes": "http://notes-website-backend:8000"}, + rest_proxy_urls={"notes": "http://notes-website-backend:8000"}) + + for path in ("/website/notes/traffic/captures", "/website/notes/traffic/capture", "/svc/notes/api/x"): + response = await _get(gw, path, headers={**_bearer(), "Authorization-Extra": "kept"}, + params={"name": "notes-website-backend.pcap0"}) + assert (response.status_code, response.content) == (200, b"the site's own page"), path + assert [request.url.path for request in seen] == ["/traffic/captures", "/traffic/capture", "/api/x"] + assert all("authorization" not in request.headers and request.headers["authorization-extra"] == "kept" + for request in seen) + + +@pytest.mark.asyncio +async def test_a_header_that_isnt_the_token_still_reaches_the_site(traffic_dir, monkeypatch): + seen: list[httpx.Request] = [] + + async def body(): + yield b"{}" + + def site(request: httpx.Request) -> httpx.Response: + if request.url.path == WELL_KNOWN_PATH: + return httpx.Response(404) + seen.append(request) + return httpx.Response(200, content=body()) + + transport = httpx.MockTransport(site) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + gw = _gateway(traffic_dir, website_urls={"notes": "http://notes"}) + for authorization in ("Bearer user-token", f"Bearer {TOKEN}"): + await _get(gw, "/website/notes/api", headers={"Authorization": authorization}) + assert [request.headers.get("authorization") for request in seen] == ["Bearer user-token", None] + + +@pytest.mark.asyncio +async def test_a_link_in_the_folder_is_neither_listed_nor_served(traffic_dir, tmp_path): + secret = tmp_path / "secret" + secret.write_bytes(b"not traffic") + (traffic_dir / "notes-website-backend.pcap0").symlink_to(secret) + (traffic_dir / "notes-website-backend.pcap1").mkdir() + gw = _gateway(traffic_dir) + assert (await _get(gw, "/traffic/captures", headers=_bearer())).json()["captures"] == [] + for name in ("notes-website-backend.pcap0", "notes-website-backend.pcap1"): + response = await _get(gw, "/traffic/capture", headers=_bearer(), params={"name": name}) + assert response.status_code == 404 and b"not traffic" not in response.content + + +def test_the_gateways_traffic_constants_mirror_the_protocols(): + """The gateway's container can't import the protocol, so it keeps its own copies of traffic/v1's names.""" + assert (gw_constants.EXT_TRAFFIC_URI, gw_constants.TRAFFIC_LIST_OP, gw_constants.TRAFFIC_GET_OP, + gw_constants.PCAP_MEDIA_TYPE, gw_constants.TRAFFIC_EXTENSION["params"]["format"]) == ( + protocol_traffic.TRAFFIC_EXTENSION_URI, protocol_traffic.TRAFFIC_LIST_OP, protocol_traffic.TRAFFIC_GET_OP, + protocol_traffic.PCAP_MEDIA_TYPE, protocol_traffic.TRAFFIC_FORMAT) + assert protocol_traffic.serves_traffic({"capabilities": {"extensions": [gw_constants.TRAFFIC_EXTENSION]}}) diff --git a/tst/unit/env/gateway/website_websocket_test.py b/tst/unit/env/gateway/website_websocket_test.py index 76ecac6b4..20c3df49c 100644 --- a/tst/unit/env/gateway/website_websocket_test.py +++ b/tst/unit/env/gateway/website_websocket_test.py @@ -295,3 +295,23 @@ async def test_a_site_that_refuses_the_handshake_has_its_answer_relayed(): await site_server.shutdown() assert refused.value.response.status_code == 403 and unknown.value.response.status_code == 404 + + +@pytest.mark.asyncio +async def test_a_sites_own_authorization_reaches_it_and_the_control_token_never_does(): + site = EchoSite() + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], control_token="control-secret", + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + route = f"ws://127.0.0.1:{gateway_port}/website/notes" + try: + for authorization in ("Bearer user-token", "Bearer control-secret"): + async with connect(f"{route}/live", additional_headers=[("Authorization", authorization)]) as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert [headers.get("authorization") for _, headers in site.handshakes] == ["Bearer user-token", None] diff --git a/tst/unit/providers/env_providers/traffic_capture_test.py b/tst/unit/providers/env_providers/traffic_capture_test.py new file mode 100644 index 000000000..ecbc2f44a --- /dev/null +++ b/tst/unit/providers/env_providers/traffic_capture_test.py @@ -0,0 +1,199 @@ +"""A gateway deploy with ``capture_traffic``: its compose runs tcpdump beside each website container, in that +container's network namespace and holding the only capture capabilities, rotating within a size cap onto a volume the +gateway alone mounts, read-only; the deploy draws the control token the gateway's traffic routes take and keeps it on +the record; and a deploy that can't capture (no website, containers, a capture that isn't running) is refused. Off, +the compose is the one it always was.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +import yaml + +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.env.envs.service_db import ServiceDBConfig +from agent_env.env.envs.website import WebsiteEnv +from agent_env.providers.env_providers import env_gateway_provider +from agent_env.providers.env_providers.constants import GATEWAY_SERVICE_NAME +from agent_env.providers.env_providers.env_gateway_provider import ( + TRAFFIC_CAPTURE_FILE_MB, + TRAFFIC_CAPTURE_FILES, + DeployedGateway, + EnvironmentGatewayProvider, + MCPServerConfig, + WebsiteConfig, + _redact_compose_secrets, + traffic_capture_services, +) +from agent_env.providers.env_state import LocalPostgresStateProvider +from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider +from agent_env.providers.sandbox_providers.sandbox_provider import Accepts +from tst.util.sandbox_providers import VMS_AND_CONTAINERS + +_WEBSITES = [WebsiteConfig(backend_image="notes-be", frontend_image=None, environment_name="notes"), + WebsiteConfig(backend_image="mail-be", frontend_image="mail-fe", environment_name="mail")] + + +def _compose(**kw) -> str: + return EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], gateway_image="agent-gateway", + website_configs=_WEBSITES, state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), **kw) + + +def _vm_provider(**attrs) -> MagicMock: + provider = MagicMock(ON_THIS_MACHINE=False, SUPPORTED_SANDBOX_MODES=VMS_AND_CONTAINERS, + SANDBOX_ACCEPTS=Accepts.LOADABLE, **attrs) + provider.links = (provider,) + return provider + + +def test_each_website_container_gets_a_capture_in_its_own_network_namespace(): + services = yaml.safe_load(_compose(capture_traffic=True, control_token="ctl"))["services"] + + captures = {name: service for name, service in services.items() if name.endswith("-traffic")} + assert traffic_capture_services(_WEBSITES) == { + "notes-website-backend-traffic": "notes-website-backend", "mail-website-backend-traffic": "mail-website-backend", + "mail-website-frontend-traffic": "mail-website-frontend"} + assert set(captures) == set(traffic_capture_services(_WEBSITES)) + for name, container in traffic_capture_services(_WEBSITES).items(): + capture = captures[name] + assert capture["image"] == "agent-gateway" + assert capture["network_mode"] == f"service:{container}" and "networks" not in capture and "ports" not in capture + assert capture["depends_on"] == [container] + assert capture["cap_add"] == ["NET_RAW", "NET_ADMIN"] + assert capture["entrypoint"] == ["tcpdump"] + assert capture["command"] == ["-i", "eth0", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), + "-W", str(TRAFFIC_CAPTURE_FILES), "-w", f"/captures/{container}.pcap"] + assert capture["volumes"] == ["traffic-captures:/captures"] + assert TRAFFIC_CAPTURE_FILE_MB * TRAFFIC_CAPTURE_FILES <= 500 # a long run can't fill the VM's disk + + +def test_only_the_captures_hold_capture_capabilities_and_only_the_gateway_reads_their_volume(): + compose = yaml.safe_load(_compose(capture_traffic=True, control_token="ctl")) + services = compose["services"] + + assert [name for name, s in services.items() if s.get("cap_add")] == sorted( + traffic_capture_services(_WEBSITES), key=list(services).index) + mounting = {name: s["volumes"] for name, s in services.items() + if any(str(v).startswith("traffic-captures:") for v in s.get("volumes") or [])} + assert mounting.pop(GATEWAY_SERVICE_NAME) == ["traffic-captures:/var/lib/agentenv/traffic:ro"] + assert set(mounting) == set(traffic_capture_services(_WEBSITES)) + assert compose["volumes"] == {"traffic-captures": {}} + gateway = services[GATEWAY_SERVICE_NAME] + assert {"AGENTENV_CONTROL_TOKEN=ctl", "TRAFFIC_CAPTURE_DIR=/var/lib/agentenv/traffic"} <= set(gateway["environment"]) + assert {name: gateway["depends_on"][name] for name in traffic_capture_services(_WEBSITES)} == dict.fromkeys( + traffic_capture_services(_WEBSITES), {"condition": "service_started"}) + + +def test_the_logged_compose_masks_the_control_token(): + compose = _compose(capture_traffic=True, control_token="ctl-secret-value") + assert "ctl-secret-value" in compose + assert "ctl-secret-value" not in _redact_compose_secrets(compose) + assert " - AGENTENV_CONTROL_TOKEN=***" in _redact_compose_secrets(compose) + + +def test_without_capture_traffic_the_compose_is_unchanged(): + plain = _compose() + assert plain == _compose(capture_traffic=False, control_token=None) + assert "traffic" not in plain and "tcpdump" not in plain and "cap_add" not in plain and "AGENTENV_CONTROL_TOKEN" not in plain + assert yaml.safe_load(_compose(capture_traffic=True, control_token="t"))["services"].keys() - yaml.safe_load( + plain)["services"].keys() == set(traffic_capture_services(_WEBSITES)) + + +def test_a_name_a_capture_would_take_is_refused(): + with pytest.raises(ValueError, match="'notes-website-backend-traffic'"): + EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp", environment_name="notes-website-backend-traffic")], + website_configs=_WEBSITES, state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), capture_traffic=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("capture_traffic", [True, False], ids=["capturing", "plain"]) +async def test_create_gateway_draws_a_control_token_only_for_a_capturing_deploy(capture_traffic): + gp = EnvironmentGatewayProvider() + with patch.object(gp, "_deploy_via_vm", AsyncMock(return_value=DeployedGateway( + gateway_url="https://gw", mcp_url="https://gw/mcp", db_web_url=None))) as vm, \ + patch.object(env_gateway_provider, "_tool_names", AsyncMock(return_value=set())): + await gp.create_gateway(sandbox_provider=_vm_provider(), mcp_servers=[], mcp_server_images=[], + website_configs=_WEBSITES, env_id="notes", capture_traffic=capture_traffic) + + kwargs = vm.await_args.kwargs + if capture_traffic: + assert kwargs["capture_traffic"] is True and isinstance(kwargs["control_token"], str) + assert len(kwargs["control_token"]) >= 32 + else: + assert "capture_traffic" not in kwargs and "control_token" not in kwargs + + +@pytest.mark.asyncio +async def test_a_capturing_deploy_keeps_its_control_token_on_the_record(): + gp = EnvironmentGatewayProvider() + env = MagicMock(spec=WebsiteEnv, id="notes-env", version=1, environment_name="notes", health_path="/api/health", + settings={}, frontend_docker_image_artifact=None, + server_docker_image_artifact=DockerImageArtifact(id="be", description="d", image_name="notes-be")) + gp._sandbox = MagicMock(sandbox_id="vm-1", type="modal_vm") + launched = DeployedGateway(gateway_url="https://gw", mcp_url="https://gw/mcp", db_web_url=None, + environment_card={"name": "env1"}, control_token="drawn") + with patch("agent_env.providers.env_state.acquire_state_for_deploy", AsyncMock(return_value=MagicMock( + state_type="remote", instance_id="st"))), \ + patch.object(gp, "_launch", AsyncMock(return_value=launched)) as launch: + record = await gp.deploy(env, _vm_provider(), capture_traffic=True) + + assert launch.await_args.kwargs["capture_traffic"] is True + assert record.control_token == "drawn" and record.control_headers == {"Authorization": "Bearer drawn"} + + +@pytest.mark.asyncio +async def test_a_deploy_that_serves_no_website_is_refused_capture_before_anything_starts(): + gp = EnvironmentGatewayProvider() + env = MagicMock(id="slack-env", environment_name="slack", docker_image_artifact=MagicMock(image_name="mcp-slack")) + acquire = AsyncMock() + with patch("agent_env.providers.env_state.acquire_state_for_deploy", acquire), \ + pytest.raises(ValueError, match="serves no website"): + await gp.deploy(env, _vm_provider(), capture_traffic=True) + acquire.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_container_mode_deploy_is_refused_capture(): + gp = EnvironmentGatewayProvider() + with pytest.raises(ValueError, match="Can't capture traffic on ModalSandboxProvider"): + await gp._deploy_gateway(ModalSandboxProvider(), state_instance=None, mcp_servers=[], mcp_server_images=[], + website_configs=_WEBSITES, capture_traffic=True, control_token="t") + + +def _ps(*rows: dict) -> str: + return "\n".join(json.dumps(row) for row in rows) + + +@pytest.mark.asyncio +async def test_a_capture_that_isnt_running_refuses_the_deploy_with_its_output(): + gp, sandbox = EnvironmentGatewayProvider(), MagicMock() + rows = _ps({"Service": "notes-website-backend-traffic", "State": "running", "ID": "a"}, + {"Service": "mail-website-backend-traffic", "State": "restarting", "ExitCode": 127, "ID": "b"}) + + async def exec_with_output(*argv): + if "ps" in argv: + return 0, rows, "" + return 0, 'exec: "tcpdump": executable file not found in $PATH', "" + + sandbox.exec_with_output = exec_with_output + with pytest.raises(RuntimeError) as raised: + await gp._refuse_stopped_captures(sandbox, ["notes-website-backend-traffic", "mail-website-backend-traffic", + "mail-website-frontend-traffic"]) + message = str(raised.value) + assert "mail-website-backend-traffic keeps restarting (last exit 127)" in message and "tcpdump" in message + assert "mail-website-frontend-traffic has no container" in message + assert "notes-website-backend-traffic" not in message + + +@pytest.mark.asyncio +async def test_running_captures_let_the_deploy_through(): + gp, sandbox = EnvironmentGatewayProvider(), MagicMock() + sandbox.exec_with_output = AsyncMock(return_value=(0, _ps( + {"Service": "notes-website-backend-traffic", "State": "running"}), "")) + await gp._refuse_stopped_captures(sandbox, ["notes-website-backend-traffic"]) diff --git a/tst/unit/task_step/deploy_env_capture_traffic_test.py b/tst/unit/task_step/deploy_env_capture_traffic_test.py new file mode 100644 index 000000000..46e40be83 --- /dev/null +++ b/tst/unit/task_step/deploy_env_capture_traffic_test.py @@ -0,0 +1,123 @@ +"""``capture_traffic``, off by default, from where a deploy is asked for to the provider: the deploy_env step and +``env deploy --capture-traffic`` pass it only when set, a website or multi env hands it to its provider, and a deploy +that can't capture is refused before anything is built: an env that serves no website, an env type whose deploy +doesn't take it, or a provider that doesn't.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from click.testing import CliRunner + +from agent_env.cli.env import env as env_cli +from agent_env.env.env import DeployedGatewayEnv +from agent_env.env.envs.multi_env import MultiEnv +from agent_env.env.envs.website import WebsiteEnv +from agent_env.providers.env_providers.env_provider import EnvironmentProvider +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep +from tst.unit.task_step.deploy_env_test import _deployed_env, _env_with_fixed_deploy, _fake_env, _mcp_env + + +def _website(env_provider_type: str = "gateway") -> WebsiteEnv: + return WebsiteEnv(id="notes-env", version=1, server_docker_image_artifact=MagicMock(), environment_name="notes", + env_provider_type=env_provider_type) + + +def _preflight(env, **options) -> list[str]: + with patch("agent_env.env.env.Env.get", return_value=env): + return DeployEnvTaskStep(id="d", version=None, env_id=env.id, **options).preflight() + + +@pytest.mark.parametrize("capture_traffic", [True, False], ids=["capturing", "plain"]) +def test_capture_traffic_roundtrips_and_a_step_without_it_serializes_as_before(capture_traffic): + step = DeployEnvTaskStep(id="t.deploy_env", version=None, env_id="notes-env", capture_traffic=capture_traffic) + assert DeployEnvTaskStep.from_dict(step.to_dict()).capture_traffic is capture_traffic + assert ("capture_traffic" in step.to_dict()) is capture_traffic + assert DeployEnvTaskStep.from_dict({**step.to_dict(), "capture_traffic": None}).capture_traffic is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("capture_traffic", [True, False], ids=["capturing", "plain"]) +async def test_execute_asks_for_a_capture_only_for_a_capturing_step(capture_traffic): + env = _fake_env(_deployed_env(None)) + step = DeployEnvTaskStep(id="t.deploy_env", version=None, env_id="art-1", capture_traffic=capture_traffic) + with patch("agent_env.env.env.Env") as Env: + Env.get.return_value = env + await step.execute(TaskStepContext()) + assert env.deploy.await_args.kwargs.get("capture_traffic") is (True if capture_traffic else None) + + +@pytest.mark.asyncio +async def test_execute_refuses_a_capture_for_an_env_whose_deploy_does_not_take_one(): + env = _env_with_fixed_deploy(_deployed_env(None)) + env.type = "fixed" + step = DeployEnvTaskStep(id="d", version=None, env_id="art-1", capture_traffic=True) + with patch("agent_env.env.env.Env") as Env, pytest.raises(ValueError, match="fixed envs take no capture_traffic"): + Env.get.return_value = env + await step.execute(TaskStepContext()) + + +def test_preflight_passes_a_capture_of_a_website_env_behind_a_gateway(): + assert _preflight(_website(), capture_traffic=True) == [] + + +def test_preflight_refuses_a_capture_of_an_env_that_serves_no_website(): + assert _preflight(_mcp_env("gateway"), capture_traffic=True) == [ + "deploy_env 'd': env 'mcp-email' serves no website, and capture_traffic records what a website's containers " + "send and receive"] + multi = MultiEnv(id="crm", version=1, mcp_server_envs=[_mcp_env("gateway")]) + assert "env 'crm' serves no website" in _preflight(multi, capture_traffic=True)[0] + + +class _NamedOptionsProvider(EnvironmentProvider): + """A plugin's provider whose deploy names its options, and no capture_traffic among them.""" + + type = "named-options" + + async def deploy(self, env, sandbox_provider, *, ttl_seconds: int = 10800, attribution=None): + raise AssertionError("refused before deploying") + + async def close(self) -> None: + pass + + +def test_preflight_refuses_a_capture_a_plugins_provider_does_not_take(): + with patch("agent_env.env.envs.website.provider_or_class", return_value=_NamedOptionsProvider): + [refusal] = _preflight(_website("named-options"), capture_traffic=True) + assert refusal.endswith("which doesn't take capture_traffic; name it in the provider's deploy() or take **options") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("env", [_website(), MultiEnv(id="crm", version=1, mcp_server_envs=[], website_envs=[_website()])], + ids=["website", "multi"]) +@pytest.mark.parametrize("capture_traffic", [True, False], ids=["capturing", "plain"]) +async def test_a_website_or_multi_env_hands_a_capture_to_its_provider_only_when_asked(env, capture_traffic): + module = "website" if isinstance(env, WebsiteEnv) else "multi_env" + record = DeployedGatewayEnv(env_id=env.id, env_version=1, gateway_url="https://gw", sandbox_id="vm") + with patch(f"agent_env.env.envs.{module}.deploy_through_provider", AsyncMock(return_value=record)) as through: + await env.deploy(capture_traffic=capture_traffic) + assert through.await_args.kwargs.get("capture_traffic") is (True if capture_traffic else None) + + +def _cli_deploy_kwargs(args: list[str]) -> dict: + captured: dict = {} + + async def fake_deploy(**kwargs): + captured.update(kwargs) + return DeployedGatewayEnv(env_id="e1", env_version=1, gateway_url="http://gw", mcp_url="http://mcp", + db_web_url=None, sandbox_id="s1") + + fake_env = MagicMock(id="e1", version=1, type="website") + fake_env.deploy = fake_deploy + with patch("agent_env.cli.env.deploy.Env") as Env, patch("agent_env.providers.build_sandbox_provider"): + Env.get.return_value = fake_env + result = CliRunner().invoke(env_cli, ["deploy", "--id", "e1", *args]) + assert result.exit_code == 0, result.output + return captured + + +def test_env_deploy_captures_traffic_only_with_the_flag(): + assert _cli_deploy_kwargs(["--capture-traffic"])["capture_traffic"] is True + assert "capture_traffic" not in _cli_deploy_kwargs([]) diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py new file mode 100644 index 000000000..576010234 --- /dev/null +++ b/tst/unit/task_step/traffic_capture_test.py @@ -0,0 +1,158 @@ +"""collect_traffic_capture against a real gateway serving traffic/v1 from a capture folder: each capture file stored +as a pcap FileArtifact, named for its container and its place among that container's files, in one universe the run +records; what's recorded when the env doesn't capture, isn't deployed or refuses the token; and teardown collecting a +capturing env the run didn't, before the env closes.""" + +from __future__ import annotations + +import os +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from agent_env.artifact import Artifact +from agent_env.artifact.store import get_artifact_store +from agent_env.env.env import DeployedGatewayEnv +from agent_env.env.gateway import gateway as gw_mod +from agent_env.env.gateway.constants import PCAP_MEDIA_TYPE, TRAFFIC_EXTENSION +from agent_env.env.gateway.gateway import Gateway +from agent_env.task.teardown import recorded_sandboxes, tear_down +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps.traffic_capture import ( + TRAFFIC_KEY, + CollectTrafficCaptureTaskStep, + collect_uncollected_captures, +) + +TOKEN = "control-secret" +_RealAsyncClient = httpx.AsyncClient +_PCAP = bytes.fromhex("d4c3b2a1020004000000000000000000ffff000001000000") + + +@pytest.fixture +def gateway(tmp_path, monkeypatch, local_stores): + """A capturing gateway every httpx client in the test reaches, its folder holding two rotated files of the + server's capture, the newer named first, and one of its frontend's.""" + monkeypatch.setattr(gw_mod, "GATEWAY_TRAJECTORY_FILE", str(tmp_path / "trajectory.jsonl")) + folder = tmp_path / "traffic" + folder.mkdir() + for name, body, mtime in (("notes-website-backend.pcap0", b"newer", 2000), ("notes-website-backend.pcap1", b"older", 1000), + ("notes-website-frontend.pcap0", b"front", 1500)): + (folder / name).write_bytes(_PCAP + body) + os.utime(folder / name, (mtime, mtime)) + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], traffic_dir=str(folder), + control_token=TOKEN) + app = gw._asgi_app() + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: _RealAsyncClient( + transport=httpx.ASGITransport(app=app), **{k: v for k, v in kw.items() if k != "verify"})) + yield gw + gw._close_trajectory_file() + + +def _deployed(env_id: str = "notes-env", *, token: str | None = TOKEN, capturing: bool = True) -> DeployedGatewayEnv: + card = {"name": "env1", "capabilities": {"extensions": [TRAFFIC_EXTENSION] if capturing else []}} + return DeployedGatewayEnv(env_id=env_id, env_version=1, gateway_url="http://gateway", sandbox_id="vm-1", + environment_card_url="http://gateway/.well-known/agent-env.json", environment_card=card, + control_token=token) + + +def _context(*envs) -> TaskStepContext: + return TaskStepContext(instance_id="run-1", deployed_envs=list(envs), metadata={}) + + +def _step(**kw) -> CollectTrafficCaptureTaskStep: + return CollectTrafficCaptureTaskStep(id="collect", version=None, env_id="notes-env", **kw) + + +def test_the_step_is_registered_and_roundtrips(): + step = _step(fail_task_on_error=True) + assert get_task_step_registry()["collect_traffic_capture"] is CollectTrafficCaptureTaskStep + again = CollectTrafficCaptureTaskStep.from_dict(step.to_dict()) + assert (again.env_id, again.fail_task_on_error) == ("notes-env", True) + assert CollectTrafficCaptureTaskStep.from_dict({"id": "c", "env_id": "e"}).fail_task_on_error is False + with pytest.raises(ValueError, match="needs the env_id"): + CollectTrafficCaptureTaskStep(id="c", version=None, env_id="") + + +@pytest.mark.asyncio +async def test_each_capture_is_stored_as_a_pcap_in_one_universe_oldest_first(gateway): + context = await _step().execute(_context(_deployed())) + + entry = context.metadata[TRAFFIC_KEY]["notes-env"] + assert entry["state"] == "collected" + assert entry["files"] == ["notes-website-backend-00.pcap", "notes-website-backend-01.pcap", + "notes-website-frontend-00.pcap"] + assert entry["containers"] == ["notes-website-backend", "notes-website-frontend"] + assert entry["bytes"] == 3 * len(_PCAP) + 15 + universe = Artifact.get(entry["artifact"]["id"], version=entry["artifact"]["version"]) + files = universe.get_file_artifacts() + assert sorted(files) == entry["files"] + first = files["notes-website-backend-00.pcap"] + assert first.content_type == PCAP_MEDIA_TYPE + assert get_artifact_store().get_object(first.object_url) == _PCAP + b"older" + assert get_artifact_store().get_object(files["notes-website-backend-01.pcap"].object_url) == _PCAP + b"newer" + + +@pytest.mark.asyncio +async def test_an_env_that_doesnt_capture_is_recorded_as_failed_and_fails_the_task_only_when_asked(gateway): + context = await _step().execute(_context(_deployed(capturing=False))) + assert context.metadata[TRAFFIC_KEY]["notes-env"]["state"] == "failed" + assert "traffic/v1" in context.metadata[TRAFFIC_KEY]["notes-env"]["error"] + with pytest.raises(RuntimeError, match="weren't collected"): + await _step(fail_task_on_error=True).execute(_context(_deployed(capturing=False))) + + +@pytest.mark.asyncio +async def test_a_record_without_the_gateways_token_is_refused(gateway): + context = await _step().execute(_context(_deployed(token="stale"))) + entry = context.metadata[TRAFFIC_KEY]["notes-env"] + assert entry["state"] == "failed" and "401" in entry["error"] + + +@pytest.mark.asyncio +async def test_an_env_the_run_didnt_deploy_is_recorded(gateway): + context = await _step().execute(_context()) + assert context.metadata[TRAFFIC_KEY]["notes-env"] == { + "state": "failed", "error": "RuntimeError: env 'notes-env' isn't deployed in this run"} + with pytest.raises(RuntimeError, match="isn't deployed"): + await _step(fail_task_on_error=True).execute(_context()) + + +@pytest.mark.asyncio +async def test_teardown_collects_only_capturing_envs_the_run_didnt_collect(gateway): + capturing, plain, done = _deployed(), _deployed("plain-env", capturing=False), _deployed("done-env") + context = _context(capturing, plain, done) + context.metadata[TRAFFIC_KEY] = {"done-env": {"state": "collected", "artifact": None}} + + await collect_uncollected_captures(context, context.deployed_envs) + + assert context.metadata[TRAFFIC_KEY]["notes-env"]["state"] == "collected" + assert "plain-env" not in context.metadata[TRAFFIC_KEY] + assert context.metadata[TRAFFIC_KEY]["done-env"] == {"state": "collected", "artifact": None} + + +@pytest.mark.asyncio +async def test_tear_down_collects_a_capture_before_its_env_closes(gateway): + context = _context(_deployed()) + collected_at_close: list[bool] = [] + + async def close(env, timeout): + collected_at_close.append(context.metadata[TRAFFIC_KEY]["notes-env"]["state"] == "collected") + + with patch("agent_env.task.teardown._close", close), \ + patch("agent_env.task.teardown.terminate", AsyncMock(return_value=False)): + report = await tear_down(context, recorded_sandboxes(context), context.deployed_envs, timeout=30) + + assert collected_at_close == [True] and [s.sandbox_id for s in report.terminated] == ["vm-1"] + + +@pytest.mark.asyncio +async def test_tear_down_leaves_an_env_already_torn_down_alone(gateway): + context = _context(_deployed()) + context.metadata["torn_down_sandbox_ids"] = ["vm-1"] + with patch("agent_env.task.teardown._close", AsyncMock()), \ + patch("agent_env.task.teardown.terminate", AsyncMock(return_value=False)): + await tear_down(context, recorded_sandboxes(context), context.deployed_envs, timeout=30) + assert TRAFFIC_KEY not in context.metadata From e197792dc6f949923cf7ff97351e1afab29ae756 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 15:02:38 -0800 Subject: [PATCH 2/8] fix(website): serve each capture as a whole copy, upload it under a prefix of its own, and capture loopback too (FD-4683) - The gateway's traffic/v1 get copies a capture before sending any of it, copying again when tcpdump truncates the file to reuse it, and answers 503 when every copy was rewritten. - collect_traffic_capture writes each collection under its own attempt prefix, and the upload owns its temp file (finish_on_thread), so a cancelled collection never removes a file an upload still reads. - tcpdump captures every interface of the container, its loopback among them. - The read size and error limits are named. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/env/gateway/gateway.py | 46 ++++++++++++-- .../env_providers/env_gateway_provider.py | 13 ++-- .../task_step/task_steps/traffic_capture.py | 37 ++++++++--- tst/unit/env/gateway/traffic_test.py | 62 +++++++++++++++++++ .../env_providers/traffic_capture_test.py | 2 +- tst/unit/task_step/traffic_capture_test.py | 43 +++++++++++++ 6 files changed, 182 insertions(+), 21 deletions(-) diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 0511c9398..427b6320e 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -12,6 +12,7 @@ import re import socket import subprocess +import tempfile from collections.abc import Iterable from contextlib import asynccontextmanager from datetime import datetime, timezone @@ -135,6 +136,8 @@ class Gateway: REST_PROXY_TIMEOUT_S = float(DATA_PLANE_LOAD_TIMEOUT_MAX_S) # data-plane loads REST_PROXY_CONNECT_TIMEOUT_S = 10.0 # a service that never answers the TCP connect fails in seconds, not an hour REST_PROXY_TIMEOUT = httpx.Timeout(REST_PROXY_TIMEOUT_S, connect=REST_PROXY_CONNECT_TIMEOUT_S) + CAPTURE_READ_BYTES = 64 * 1024 + CAPTURE_SNAPSHOT_ATTEMPTS = 3 def __init__( self, @@ -437,7 +440,8 @@ async def traffic_captures(request: Request) -> Response: @self._mcp.custom_route("/traffic/capture", methods=["GET"]) async def traffic_capture(request: Request) -> Response: - return self._refuse_uncontrolled(request) or self._serve_capture(request.query_params.get("name", "")) + return self._refuse_uncontrolled(request) or await asyncio.to_thread( + self._serve_capture, request.query_params.get("name", "")) @self._mcp.custom_route("/changes", methods=["GET"]) async def env_changes(request: Request) -> Response: @@ -1635,8 +1639,10 @@ def _serve_captures(self) -> Response: return self._json_response({"format": "pcap", "captures": captures}, 200) def _serve_capture(self, name: str) -> Response: - """traffic/v1's get: the capture file ``name``, one ``list`` names, streamed as it stands; tcpdump writes each - packet as it captures it, so a file still being written ends on a whole packet, or at worst one torn.""" + """traffic/v1's get: the capture file ``name``, one ``list`` names, as a copy taken whole before any of it is + sent (``_snapshot_capture``), so a file tcpdump's rotation rewrites while it downloads never arrives mixed or + cut short; 503 when every copy was rewritten under it. A file still being written ends on a whole packet, or at + worst one torn.""" if not _CAPTURE_FILE.fullmatch(name): return self._json_response({"ok": False, "error": {"code": "invalid_params", "message": f"{name!r} is not a capture's name, as list gives it"}}, 400) @@ -1648,14 +1654,42 @@ def _serve_capture(self, name: str) -> Response: except OSError: return self._json_response({"ok": False, "error": {"code": "not_found", "message": f"no capture named {name!r}"}}, 404) + with handle: + snapshot = self._snapshot_capture(handle) + if snapshot is None: + return self._json_response({"ok": False, "error": {"code": "unavailable", + "message": f"capture {name!r} kept being rewritten while it was copied; " + "ask again"}}, 503) + size = os.fstat(snapshot.fileno()).st_size def chunks(): - with handle: - while chunk := handle.read(65536): + with snapshot: + while chunk := snapshot.read(self.CAPTURE_READ_BYTES): yield chunk return StreamingResponse(chunks(), media_type=PCAP_MEDIA_TYPE, - headers={"Content-Disposition": f'attachment; filename="{name}"'}) + headers={"Content-Disposition": f'attachment; filename="{name}"', + "Content-Length": str(size)}) + + def _snapshot_capture(self, handle: Any) -> Any | None: + """A copy of the open capture ``handle`` in an unlinked temp file, read back from its start, or None when each + of ``CAPTURE_SNAPSHOT_ATTEMPTS`` copies found the file shorter than what it had copied: tcpdump truncated it to + reuse it, so the copy would mix two rounds of packets. The copy is gone once it's closed.""" + for _ in range(self.CAPTURE_SNAPSHOT_ATTEMPTS): + handle.seek(0) + copy = tempfile.TemporaryFile() + copied, rewritten = 0, False + while chunk := handle.read(self.CAPTURE_READ_BYTES): + copy.write(chunk) + copied += len(chunk) + if os.fstat(handle.fileno()).st_size < copied: + rewritten = True + break + if not rewritten and os.fstat(handle.fileno()).st_size >= copied: + copy.seek(0) + return copy + copy.close() + return None def _serve_trajectory(self) -> Response: """Stream the whole env trajectory as JSONL; only a mid-write tail line can arrive torn.""" diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 34f227c38..7a877706c 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -85,6 +85,7 @@ # How much of a service's output a failed deploy's error carries: enough for the end of a traceback. _LOG_TAIL_LINES = 30 _LOG_LINE_CHARS = 300 +_CAPTURE_CHECK_ERROR_CHARS = 500 # Gathering it may stall with the connection; past this, the deploy raises the error it had. _DOWN_SERVICES_TIMEOUT_S = 60.0 @@ -1276,7 +1277,7 @@ async def _refuse_stopped_captures(self, sandbox: VmSandbox, capture_services: l return if exit_code != 0: logger.warning(f"Can't check the traffic captures are running: docker compose ps exited {exit_code}: " - f"{stderr.strip()[-500:]}") + f"{stderr.strip()[-_CAPTURE_CHECK_ERROR_CHARS:]}") return rows = {row.get("Service"): row for row in _compose_ps_rows(stdout)} stopped = [] @@ -1663,14 +1664,14 @@ def _traffic_mount(path: str, *, read_only: bool = False) -> str: def _traffic_capture_service(gateway_image: str, container: str) -> list[str]: - """The compose lines of the service capturing ``container``'s traffic: tcpdump, from the gateway's image, on the - container's own interface, every packet whole and written as it's captured, rotating through - ``TRAFFIC_CAPTURE_FILES`` files of ``TRAFFIC_CAPTURE_FILE_MB`` million bytes named for the container, its oldest - overwritten once they're full, so a long run holds at most their product per container on the VM's disk. It joins + """The compose lines of the service capturing ``container``'s traffic: tcpdump, from the gateway's image, on every + interface of the container's, its loopback among them, every packet whole and written as it's captured, rotating + through ``TRAFFIC_CAPTURE_FILES`` files of ``TRAFFIC_CAPTURE_FILE_MB`` million bytes named for the container, its + oldest overwritten once they're full, so a long run holds at most their product per container on the VM's disk. It joins the container's network namespace, so it sees what the container sends and receives and nothing else, and is the only service given the capabilities a capture needs. It writes to ``TRAFFIC_VOLUME``, which the gateway mounts read-only and no website container mounts at all.""" - argv = ["-i", "eth0", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), + argv = ["-i", "any", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), "-W", str(TRAFFIC_CAPTURE_FILES), "-w", f"{TRAFFIC_SIDECAR_DIR}/{container}.pcap"] return [ f" image: {gateway_image}", diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index aa2e0b65d..661c6932b 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -5,11 +5,10 @@ didn't, before the env closes (``collect_uncollected_captures``), so a capture isn't lost with its VM.""" from __future__ import annotations -import asyncio +import functools import logging import os import tempfile -import time from collections.abc import Sequence from typing import ClassVar, Optional @@ -32,12 +31,14 @@ from agent_env.store.ids import derive_id, is_local_id, validate_local_id from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency +from agent_env.task_step.thread_work import finish_on_thread logger = logging.getLogger(__name__) #: ``context.metadata[TRAFFIC_KEY][env_id]``: each env whose captures a run collected, or tried to. TRAFFIC_KEY = "traffic_captures" _ARTIFACT_TYPE = "traffic_captures" +_ERROR_CHARS = 500 class CollectTrafficCaptureTaskStep(TaskStep): @@ -116,7 +117,8 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list if is_local_id(universe_id): validate_local_id(universe_id) get_config().check_local_run_write(universe_id) - store, version = get_artifact_store(), int(time.time()) + store = get_artifact_store() + prefix = store.attempt_prefix(_ARTIFACT_TYPE, universe_id) files, total, seen = {}, 0, {} for capture in captures: container = capture.container @@ -127,10 +129,12 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list try: total += await fetch_traffic_capture(deployed.environment_url, deployed.environment_card, capture.name, path, headers=deployed.control_headers) - url = await asyncio.to_thread(store.put_object_file, _ARTIFACT_TYPE, universe_id, version, filename, path, - PCAP_MEDIA_TYPE) - finally: - os.unlink(path) + except BaseException: + _remove(path) + raise + url = await finish_on_thread(functools.partial(_upload, universe_id, prefix + filename, path), + f"Uploading traffic capture {filename} of env {deployed.env_id!r}", + if_never_run=functools.partial(_remove, path)) files[filename] = store.put_document(FileArtifact( id=derive_id(universe_id, filename), filename=filename, content_type=PCAP_MEDIA_TYPE, object_url=url, description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}")) @@ -143,5 +147,22 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list "containers": sorted(seen)} +def _upload(universe_id: str, object_url: str, path: str) -> str: + """Upload the capture at ``path`` to ``object_url``, write-once, and remove it: the upload owns the file, so a + cancelled collection never removes it from under the upload.""" + try: + return get_config().get_object_store_to_write(object_url, universe_id).put_file_at(object_url, path, + PCAP_MEDIA_TYPE) + finally: + _remove(path) + + +def _remove(path: str) -> None: + try: + os.unlink(path) + except FileNotFoundError: + pass + + def _failed(error: BaseException) -> dict: - return {"state": "failed", "error": f"{type(error).__name__}: {error}"[:500]} + return {"state": "failed", "error": f"{type(error).__name__}: {error}"[:_ERROR_CHARS]} diff --git a/tst/unit/env/gateway/traffic_test.py b/tst/unit/env/gateway/traffic_test.py index a5a410133..17f11abd5 100644 --- a/tst/unit/env/gateway/traffic_test.py +++ b/tst/unit/env/gateway/traffic_test.py @@ -71,6 +71,68 @@ async def test_list_names_each_capture_by_its_container_oldest_first_and_get_str assert (got.status_code, got.headers["content-type"], got.content) == (200, PCAP_MEDIA_TYPE, _PCAP_HEADER + b"first") +class _Rewritten: + """An open capture file tcpdump reuses ``times`` times, each time just after a read of it: it truncates the file + and writes its next round's first packets, shorter than what was read.""" + + def __init__(self, path, times: int): + self.path, self.times, self.handle = path, times, open(path, "rb") + self.round = path.read_bytes() + + def read(self, size: int) -> bytes: + chunk = self.handle.read(size) + if chunk and self.times: + self.times -= 1 + self.path.write_bytes(_PCAP_HEADER + b"next") + return chunk + + def seek(self, offset: int) -> None: + if self.times: + self.path.write_bytes(self.round) + self.handle.seek(offset) + + def fileno(self) -> int: + return self.handle.fileno() + + +def test_a_capture_rewritten_while_it_is_copied_is_copied_again_whole(traffic_dir): + path = traffic_dir / "notes-website-backend.pcap0" + path.write_bytes(_PCAP_HEADER + b"the round tcpdump overwrites") + gw = _gateway(traffic_dir) + + with gw._snapshot_capture(_Rewritten(path, times=1)) as copy: + assert copy.read() == _PCAP_HEADER + b"next" + path.write_bytes(_PCAP_HEADER + b"the round tcpdump overwrites") + assert gw._snapshot_capture(_Rewritten(path, times=gw.CAPTURE_SNAPSHOT_ATTEMPTS)) is None + + +@pytest.mark.asyncio +async def test_get_answers_503_when_every_copy_of_a_capture_was_rewritten(traffic_dir, monkeypatch): + _write(traffic_dir, "notes-website-backend.pcap0", _PCAP_HEADER + b"x", 1000) + monkeypatch.setattr(Gateway, "_snapshot_capture", lambda self, handle: None) + response = await _get(_gateway(traffic_dir), "/traffic/capture", headers=_bearer(), + params={"name": "notes-website-backend.pcap0"}) + assert response.status_code == 503 and response.json()["error"]["code"] == "unavailable" + + +@pytest.mark.asyncio +async def test_get_sends_the_copy_it_took_whatever_the_file_becomes_after(traffic_dir, monkeypatch): + path = traffic_dir / "notes-website-backend.pcap0" + path.write_bytes(_PCAP_HEADER + b"whole") + real = Gateway._snapshot_capture + + def snapshot_then_rewrite(self, handle): + copy = real(self, handle) + path.write_bytes(_PCAP_HEADER + b"the next round, longer") + return copy + + monkeypatch.setattr(Gateway, "_snapshot_capture", snapshot_then_rewrite) + response = await _get(_gateway(traffic_dir), "/traffic/capture", headers=_bearer(), + params={"name": "notes-website-backend.pcap0"}) + assert (response.content, response.headers["content-length"]) == (_PCAP_HEADER + b"whole", + str(len(_PCAP_HEADER) + 5)) + + @pytest.mark.asyncio @pytest.mark.parametrize("name, status", [ ("../trajectory.jsonl", 400), ("notes-website-backend.pcap0/../../x", 400), ("stray.txt", 400), diff --git a/tst/unit/providers/env_providers/traffic_capture_test.py b/tst/unit/providers/env_providers/traffic_capture_test.py index ecbc2f44a..6d75cf617 100644 --- a/tst/unit/providers/env_providers/traffic_capture_test.py +++ b/tst/unit/providers/env_providers/traffic_capture_test.py @@ -65,7 +65,7 @@ def test_each_website_container_gets_a_capture_in_its_own_network_namespace(): assert capture["depends_on"] == [container] assert capture["cap_add"] == ["NET_RAW", "NET_ADMIN"] assert capture["entrypoint"] == ["tcpdump"] - assert capture["command"] == ["-i", "eth0", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), + assert capture["command"] == ["-i", "any", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), "-W", str(TRAFFIC_CAPTURE_FILES), "-w", f"/captures/{container}.pcap"] assert capture["volumes"] == ["traffic-captures:/captures"] assert TRAFFIC_CAPTURE_FILE_MB * TRAFFIC_CAPTURE_FILES <= 500 # a long run can't fill the VM's disk diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index 576010234..3fa74601c 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -5,7 +5,9 @@ from __future__ import annotations +import asyncio import os +import threading from unittest.mock import AsyncMock, patch import httpx @@ -20,6 +22,7 @@ from agent_env.task.teardown import recorded_sandboxes, tear_down from agent_env.task_step.context import TaskStepContext from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps import traffic_capture as traffic_capture_module from agent_env.task_step.task_steps.traffic_capture import ( TRAFFIC_KEY, CollectTrafficCaptureTaskStep, @@ -95,6 +98,46 @@ async def test_each_capture_is_stored_as_a_pcap_in_one_universe_oldest_first(gat assert get_artifact_store().get_object(files["notes-website-backend-01.pcap"].object_url) == _PCAP + b"newer" +@pytest.mark.asyncio +async def test_two_collections_of_one_env_in_the_same_second_are_each_stored_under_a_prefix_of_their_own(gateway): + context = _context(_deployed()) + first = dict(await traffic_capture_module.collect_captures(context, _deployed())) + second = await traffic_capture_module.collect_captures(context, _deployed()) + + assert (first["state"], second["state"]) == ("collected", "collected") + urls = [{name: artifact.object_url for name, artifact in + Artifact.get(entry["artifact"]["id"], version=entry["artifact"]["version"]).get_file_artifacts().items()} + for entry in (first, second)] + assert sorted(urls[0]) == sorted(urls[1]) and not set(urls[0].values()) & set(urls[1].values()) + + +@pytest.mark.asyncio +async def test_a_collection_cancelled_mid_upload_leaves_the_upload_its_file_to_finish(gateway, monkeypatch): + started, release, uploaded = threading.Event(), threading.Event(), [] + real = traffic_capture_module._upload + + def held(universe_id, object_url, path): + started.set() + release.wait(10) + uploaded.append((path, os.path.exists(path))) + return real(universe_id, object_url, path) + + monkeypatch.setattr(traffic_capture_module, "_upload", held) + collecting = asyncio.ensure_future(traffic_capture_module.collect_captures(_context(_deployed()), _deployed())) + assert await asyncio.to_thread(started.wait, 10) + collecting.cancel() + with pytest.raises(asyncio.CancelledError): + await collecting + release.set() + for _ in range(200): + if uploaded and not os.path.exists(uploaded[0][0]): + break + await asyncio.sleep(0.01) + + [(path, existed)] = uploaded + assert existed and not os.path.exists(path) + + @pytest.mark.asyncio async def test_an_env_that_doesnt_capture_is_recorded_as_failed_and_fails_the_task_only_when_asked(gateway): context = await _step().execute(_context(_deployed(capturing=False))) From 3e2e3cef4f576f898b6fc9e9eaffaf07c165f503 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 16:06:07 -0800 Subject: [PATCH 3/8] fix(website): teardown tries a failed traffic collection again before the env closes (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/task_step/task_steps/traffic_capture.py | 9 +++++---- tst/unit/task_step/traffic_capture_test.py | 5 +++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index 661c6932b..a429d3a73 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -86,11 +86,12 @@ def captures_traffic(deployed: DeployedEnv) -> bool: async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[DeployedEnv]) -> None: - """Collect the captures of each of ``envs`` that serves traffic/v1 and whose captures the run hasn't collected or - tried to, as teardown does before the envs close; ``collect_captures`` records what happened to each.""" - tried = context.metadata.get(TRAFFIC_KEY) or {} + """Collect the captures of each of ``envs`` that serves traffic/v1 and whose captures the run hasn't collected, a + collection that failed tried again, as teardown does before the envs close; ``collect_captures`` records what + happened to each.""" + entries = context.metadata.get(TRAFFIC_KEY) or {} for deployed in envs: - if captures_traffic(deployed) and deployed.env_id not in tried: + if captures_traffic(deployed) and (entries.get(deployed.env_id) or {}).get("state") != "collected": await collect_captures(context, deployed) diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index 3fa74601c..729e91342 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -164,10 +164,11 @@ async def test_an_env_the_run_didnt_deploy_is_recorded(gateway): @pytest.mark.asyncio -async def test_teardown_collects_only_capturing_envs_the_run_didnt_collect(gateway): +async def test_teardown_collects_each_capturing_env_the_run_didnt_collect_a_failed_collection_again(gateway): capturing, plain, done = _deployed(), _deployed("plain-env", capturing=False), _deployed("done-env") context = _context(capturing, plain, done) - context.metadata[TRAFFIC_KEY] = {"done-env": {"state": "collected", "artifact": None}} + context.metadata[TRAFFIC_KEY] = {"done-env": {"state": "collected", "artifact": None}, + "notes-env": {"state": "failed", "error": "ConnectError: refused"}} await collect_uncollected_captures(context, context.deployed_envs) From 9953df5d353e26ccfdb41d912e031ee2cf706109 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 16:43:47 -0800 Subject: [PATCH 4/8] fix(website): teardown saves the traffic it collects to the run, and collects a retry's rolled-back deployments too (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- src/agent_env/task/store.py | 15 +++++ src/agent_env/task/task.py | 6 +- src/agent_env/task/teardown.py | 9 ++- .../task_step/task_steps/traffic_capture.py | 61 +++++++++++++------ tst/unit/task_step/retry_test.py | 2 + tst/unit/task_step/traffic_capture_test.py | 41 ++++++++++++- 7 files changed, 110 insertions(+), 26 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fb35c3ec8..d1971eb4d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -119,7 +119,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`, one site: a server image, `backend_image` in its env.toml, serving the pages and API, ready when its card's `website/v1` `ready` answers a `WebsiteReadiness` that says so, as `agentenv_protocol.website` (`website_readiness`, `env_readiness`) defines a website env's health for the deploy and `DeployedEnv.readiness()` alike, or, for a card without `ready`, when its `health` path answers 200, which is deprecated; the gateway container's `env/gateway/constants.py` mirrors the URI and op, a test keeping them equal; a separate frontend image is deprecated, warned of at put and optional in a bundle, and a stored `backend_docker_image_artifact` reads as the server image), `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`), a run only the forms its sandbox's deploys boot (`booted_forms`: on each link, the first form it boots natively, by the deploy's own rule, `boot.first_native`; a compose of a disk adds its container), and a deploy composes first what it boots that no compose started (`ComputerEnv.compose_unstarted`); 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 (`utils/probing.py`) 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, over its env.toml's `[settings]` and a computer env's `settings` beside it in `apps`, each by top-level field (`params.app_settings`; checked at put, a required one left to the deploy; the deploy record keeps the merged values as `app_settings`, secrets withheld); 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. Its site route, `/website//`, never forwards a site's data plane or card (`/agentenv/…`, `/.well-known/agent-env.json`), which agentenv reaches through `/svc//`; a site with no frontend image is served by its server image there. A deploy with `capture_traffic` (`deploy_env`'s, `env deploy --capture-traffic`; a website env, or a multi env with a website, on a VM, refused for any other env and in container mode) records every packet each website's server and frontend containers send and receive: beside each runs a capture service, tcpdump from the gateway's image in that container's network namespace (`network_mode: "service:"`), the only service holding `NET_RAW` and `NET_ADMIN`, rotating its classic pcap files (10 of 20 MB each, the oldest overwritten) onto a volume no website container mounts and the gateway mounts read-only and serves as `urn:agentenv:traffic/v1` (`list`, `GET /traffic/captures`, oldest first; `get`, `GET /traffic/capture?name=`), which `agentenv_protocol.traffic` defines (`TrafficCapture`, `list_traffic_captures`, `fetch_traffic_capture`), the gateway's `constants.py` mirroring its names, a test keeping them equal. Those routes take the per-deploy control token the record keeps (`DeployedEnv.control_token`, the gateway's `AGENTENV_CONTROL_TOKEN`), which the gateway takes off every request before it proxies one, so no site reads it; a capture that isn't running once the gateway is up refuses the deploy. 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/`, among them `start_computer_recording` and `end_computer_recording`, `collect_traffic_capture`, which stores each traffic/v1 capture of an env deployed with `capture_traffic` as a pcap `FileArtifact` in one `FileArtifactUniverse` recorded under `context.metadata["traffic_captures"][env_id]`, 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 capturing env still up whose traffic the run didn't collect is collected, 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/`, among them `start_computer_recording` and `end_computer_recording`, `collect_traffic_capture`, which stores each traffic/v1 capture of an env deployed with `capture_traffic` as a pcap `FileArtifact` in one `FileArtifactUniverse` recorded under `context.metadata["traffic_captures"][env_id]`, 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 capturing env still up whose traffic the run didn't collect is collected, a retry's rolled-back ones too (kept as `capturing_envs` in its `retry_resets` marker, each under `-rolled-back-`), and what it records saved to the stored instance's `context.metadata`, 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 namespaces, `@local` and `@hf` (`is_local_id`), each with a documents file of its own: while the CLI or a `build_app` server runs, a local id and what a 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). | diff --git a/src/agent_env/task/store.py b/src/agent_env/task/store.py index 869db45d0..1a720a82a 100644 --- a/src/agent_env/task/store.py +++ b/src/agent_env/task/store.py @@ -789,6 +789,12 @@ def record_task_cancelled_sync( ), ) + def set_context_metadata_sync(self, instance_id: str, key: str, value: object) -> None: + """Set ``context.metadata[key]`` of a stored instance outside any step, as teardown records what it did after + the run's last step saved; not journaled, since nothing after teardown undoes it.""" + self._doc_store.update(TASK_INSTANCES_COLLECTION, Filter.of(instance_id=instance_id), + UpdateSpec(set={f"context.metadata.{key}": value})) + def append_step_attempt_failure_sync(self, instance_id: str, failure: "StepAttemptFailure") -> None: """Best-effort append of one attempt's failure, designed not to conflict with the ``record_task_failure`` write. ``add_to_set`` creates the array on the first append and @@ -898,6 +904,15 @@ def seed_task_instance_context(instance_id: str, context: "TaskStepContext") -> logger.warning("Failed to seed task instance context", exc_info=True) +def record_context_metadata(instance_id: str, key: str, value: object) -> None: + """Save ``context.metadata[key]`` of the stored instance ``instance_id``, as teardown does once the run's steps + have saved theirs; a failure is logged, never raised.""" + try: + get_task_instance_store().set_context_metadata_sync(instance_id, key, value) + except Exception: + logger.warning("Failed to save context.metadata[%r] of task instance %s", key, instance_id, exc_info=True) + + async def record_step_complete( instance_id: str, step_id: str, diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 6840a58d8..c6ce57022 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Callable, ClassVar, Optional from agentenv_protocol import client as protocol_v1 +from agentenv_protocol.traffic import serves_traffic from agent_env.config import run_env_ids from agent_env.plugins import _registration @@ -44,7 +45,9 @@ def _rolled_back_sandboxes(before: dict, after: dict) -> dict: context ``before`` an undo and absent ``after`` it, shaped like the context's own slices so the executor's cleanup can scan the marker exactly as it scans a context. The undo is what removes these from the live and persisted context, so without this record an - orphaned sandbox from the failed attempt is discoverable nowhere. + orphaned sandbox from the failed attempt is discoverable nowhere. ``capturing_envs`` keeps + the whole record of each env among them that serves traffic/v1, so teardown can still + collect its captures before its sandboxes go. """ def gone(key: str) -> list[dict]: kept = {_identity(e) for e in after.get(key) or []} @@ -64,6 +67,7 @@ def gone(key: str) -> list[dict]: {"sandbox_id": x.get("sandbox_id"), "sandbox_type": x.get("sandbox_type")} for x in gone("deployed_sandboxes") ], + "capturing_envs": [e for e in gone("deployed_envs") if serves_traffic(e.get("environment_card") or {})], } diff --git a/src/agent_env/task/teardown.py b/src/agent_env/task/teardown.py index a0aa9e709..2468b70c1 100644 --- a/src/agent_env/task/teardown.py +++ b/src/agent_env/task/teardown.py @@ -32,7 +32,7 @@ from agent_env.task.interrupts import tearing_down from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.computer_recording import end_open_recordings, register_open_recordings -from agent_env.task_step.task_steps.traffic_capture import collect_uncollected_captures +from agent_env.task_step.task_steps.traffic_capture import collect_uncollected_captures, rolled_back_deployments logger = logging.getLogger(__name__) @@ -119,7 +119,8 @@ async def teardown_run(context: TaskStepContext, *, timeout: float = TERMINATE_T async def tear_down(context: TaskStepContext, sandboxes: Sequence[RecordedSandbox], envs: Sequence[DeployedEnv], *, timeout: float) -> TeardownReport: """End the recordings the run left open on ``envs`` with fresh grants, collect the traffic captures of each of them - that captures and that the run didn't collect, close each env still up, register any + that captures and that the run didn't collect, and of each capturing env a retry rolled back whose sandboxes are + among ``sandboxes``, close each env still up, register any recording still open from its namespace, then terminate each of ``sandboxes`` by id, all at once and each within ``timeout`` seconds, and remove a local one's work folder once it's down. A close that fails is logged: its env's sandboxes are terminated by id regardless, which leaves one the close took down as it is. Cancelling it returns what it reached, with the rest in ``left``.""" @@ -133,9 +134,11 @@ async def tear_down(context: TaskStepContext, sandboxes: Sequence[RecordedSandbo except TimeoutError: logger.warning("teardown: the run's open recordings weren't all ended within %gs", timeout) still_up = [env for env in envs if any(sandbox_id in up for sandbox_id, _ in env_sandbox_ids(env))] + rolled_back = [(key, env) for key, env in rolled_back_deployments(context) + if any(sandbox_id in up for sandbox_id, _ in env_sandbox_ids(env))] try: async with asyncio.timeout(timeout): - await collect_uncollected_captures(context, still_up) + await collect_uncollected_captures(context, still_up, rolled_back=rolled_back) except TimeoutError: logger.warning("teardown: the run's uncollected traffic captures weren't all collected within %gs", timeout) await asyncio.gather(*(_close(env, timeout) for env in still_up)) diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index a429d3a73..4a93e3919 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -2,9 +2,11 @@ packet its websites' containers send and receive as rotating pcap files, which its gateway's card serves as traffic/v1 behind the deploy's control token. ``collect_traffic_capture`` stores each as a FileArtifact and registers them as one FileArtifactUniverse, recorded in ``context.metadata[TRAFFIC_KEY][env_id]``. Teardown collects any capturing env a run -didn't, before the env closes (``collect_uncollected_captures``), so a capture isn't lost with its VM.""" +didn't, a retry's rolled-back ones included, before the env closes (``collect_uncollected_captures``), so a capture +isn't lost with its VM, and saves what it recorded to the run's stored instance.""" from __future__ import annotations +import asyncio import functools import logging import os @@ -85,36 +87,57 @@ def captures_traffic(deployed: DeployedEnv) -> bool: return isinstance(deployed, DeployedEnv) and serves_traffic(deployed.environment_card) -async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[DeployedEnv]) -> None: +async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[DeployedEnv], *, + rolled_back: Sequence[tuple[str, DeployedEnv]] = ()) -> None: """Collect the captures of each of ``envs`` that serves traffic/v1 and whose captures the run hasn't collected, a - collection that failed tried again, as teardown does before the envs close; ``collect_captures`` records what - happened to each.""" + collection that failed tried again, as teardown does before the envs close, and of each ``rolled_back`` deployment, + a retry's, under its own key (``rolled_back_deployments``); ``collect_captures`` records what happened to each, and + the entries are saved to the run's stored instance, since teardown runs after its steps saved theirs.""" entries = context.metadata.get(TRAFFIC_KEY) or {} - for deployed in envs: - if captures_traffic(deployed) and (entries.get(deployed.env_id) or {}).get("state") != "collected": - await collect_captures(context, deployed) - - -async def collect_captures(context: TaskStepContext, deployed: DeployedEnv) -> dict: + tried = False + for key, deployed in [*((deployed.env_id, deployed) for deployed in envs), *rolled_back]: + if captures_traffic(deployed) and (entries.get(key) or {}).get("state") != "collected": + await collect_captures(context, deployed, key=key) + tried = True + if tried and context.instance_id: + from agent_env.task.store import record_context_metadata + + await asyncio.to_thread(record_context_metadata, context.instance_id, TRAFFIC_KEY, + context.metadata[TRAFFIC_KEY]) + + +def rolled_back_deployments(context: TaskStepContext) -> list[tuple[str, DeployedEnv]]: + """Each capturing env a retry rolled back, with the key its captures are recorded under: + ``-rolled-back-`` for the ``n``-th of the run's ``retry_resets``.""" + found = [] + for n, reset in enumerate(context.metadata.get("retry_resets") or [], start=1): + for record in (reset.get("rolled_back_sandboxes") or {}).get("capturing_envs") or []: + deployed = DeployedEnv.from_dict(record) + found.append((f"{deployed.env_id}-rolled-back-{n}", deployed)) + return found + + +async def collect_captures(context: TaskStepContext, deployed: DeployedEnv, *, key: Optional[str] = None) -> dict: """Store each capture ``deployed`` serves as a pcap FileArtifact, register them as one FileArtifactUniverse, and - record it in ``context.metadata[TRAFFIC_KEY]``. Each file is named for its container and its place among that - container's files, oldest first, so they read in the order their packets were captured. Never raises: the entry, - which it also returns, says what happened.""" + record it in ``context.metadata[TRAFFIC_KEY][key]``, by default its env_id. Each file is named for its container + and its place among that container's files, oldest first, so they read in the order their packets were captured. + Never raises: the entry, which it also returns, says what happened.""" + key = key or deployed.env_id entries = context.metadata.setdefault(TRAFFIC_KEY, {}) try: if not captures_traffic(deployed): raise EnvCapabilityUnsupported(TRAFFIC_EXTENSION_URI, TRAFFIC_LIST_OP, deployed.env_id, deployed.instance_id) listing = await list_traffic_captures(deployed.environment_url, deployed.environment_card, headers=deployed.control_headers) - entries[deployed.env_id] = await _store(context, deployed, listing.captures) + entries[key] = await _store(context, deployed, listing.captures, key) except Exception as e: - logger.warning("collect_traffic_capture %s: %s", deployed.env_id, e) - entries[deployed.env_id] = _failed(e) - return entries[deployed.env_id] + logger.warning("collect_traffic_capture %s: %s", key, e) + entries[key] = _failed(e) + return entries[key] -async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list[TrafficCapture]) -> dict: - universe_id = derive_id(context.instance_id or "run", f"traffic-{deployed.env_id}") +async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list[TrafficCapture], key: str) -> dict: + universe_id = derive_id(context.instance_id or "run", f"traffic-{key}") if is_local_id(universe_id): validate_local_id(universe_id) get_config().check_local_run_write(universe_id) diff --git a/tst/unit/task_step/retry_test.py b/tst/unit/task_step/retry_test.py index 46c67094e..66dd3ea8d 100644 --- a/tst/unit/task_step/retry_test.py +++ b/tst/unit/task_step/retry_test.py @@ -173,6 +173,7 @@ async def test_retry_rolls_back_span_and_redispatches_on_fresh_world(store): "deployed_envs": [{"sandbox_id": "sb-s0-1", "sandbox_type": "fake", "sandbox_ids": {}}], "deployed_agents": [], "deployed_sandboxes": [], + "capturing_envs": [], } assert [f.get("retried") for f in result.metadata["failed_steps"]] == [True] @@ -773,6 +774,7 @@ def test_rolled_back_sandboxes_is_the_difference_between_stored_and_replayed(): "deployed_envs": [{"sandbox_id": "sb-e", "sandbox_type": "beta", "sandbox_ids": {"vm": "sb-vm"}}], "deployed_agents": [{"sandbox_id": "sb-a", "sandbox_type": "local"}], "deployed_sandboxes": [{"sandbox_id": "sb-x", "sandbox_type": None}], + "capturing_envs": [], } diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index 729e91342..f3a5a8c0b 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -19,6 +19,8 @@ from agent_env.env.gateway import gateway as gw_mod from agent_env.env.gateway.constants import PCAP_MEDIA_TYPE, TRAFFIC_EXTENSION from agent_env.env.gateway.gateway import Gateway +from agent_env.task.store import find_task_instance, get_task_instance_store +from agent_env.task.task import _rolled_back_sandboxes from agent_env.task.teardown import recorded_sandboxes, tear_down from agent_env.task_step.context import TaskStepContext from agent_env.task_step.registry import get_task_step_registry @@ -54,9 +56,10 @@ def gateway(tmp_path, monkeypatch, local_stores): gw._close_trajectory_file() -def _deployed(env_id: str = "notes-env", *, token: str | None = TOKEN, capturing: bool = True) -> DeployedGatewayEnv: +def _deployed(env_id: str = "notes-env", *, token: str | None = TOKEN, capturing: bool = True, + sandbox_id: str = "vm-1") -> DeployedGatewayEnv: card = {"name": "env1", "capabilities": {"extensions": [TRAFFIC_EXTENSION] if capturing else []}} - return DeployedGatewayEnv(env_id=env_id, env_version=1, gateway_url="http://gateway", sandbox_id="vm-1", + return DeployedGatewayEnv(env_id=env_id, env_version=1, gateway_url="http://gateway", sandbox_id=sandbox_id, environment_card_url="http://gateway/.well-known/agent-env.json", environment_card=card, control_token=token) @@ -200,3 +203,37 @@ async def test_tear_down_leaves_an_env_already_torn_down_alone(gateway): patch("agent_env.task.teardown.terminate", AsyncMock(return_value=False)): await tear_down(context, recorded_sandboxes(context), context.deployed_envs, timeout=30) assert TRAFFIC_KEY not in context.metadata + + +@pytest.mark.asyncio +async def test_what_teardown_collects_is_saved_to_the_runs_stored_instance(gateway): + instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) + context = _context(_deployed()) + context.instance_id = instance.instance_id + + await collect_uncollected_captures(context, context.deployed_envs) + + stored = find_task_instance(instance.instance_id).context["metadata"][TRAFFIC_KEY] + assert stored == context.metadata[TRAFFIC_KEY] and stored["notes-env"]["state"] == "collected" + + +@pytest.mark.asyncio +async def test_tear_down_collects_a_retrys_rolled_back_deployment_under_its_own_key_before_its_vm_goes(gateway): + rolled_back, live = _deployed(sandbox_id="vm-1"), _deployed(sandbox_id="vm-2") + gone = _rolled_back_sandboxes(TaskStepContext(deployed_envs=[rolled_back]).to_dict(), TaskStepContext().to_dict()) + context = _context(live) + context.metadata["retry_resets"] = [{"attempt": 2, "rolled_back_sandboxes": gone}] + collected_before_terminate: list[set] = [] + + async def terminate(sandbox): + collected_before_terminate.append({key for key, entry in context.metadata[TRAFFIC_KEY].items() + if entry["state"] == "collected"}) + return False + + with patch("agent_env.task.teardown._close", AsyncMock()), patch("agent_env.task.teardown.terminate", terminate): + report = await tear_down(context, recorded_sandboxes(context), context.deployed_envs, timeout=30) + + assert sorted(s.sandbox_id for s in report.terminated) == ["vm-1", "vm-2"] + assert collected_before_terminate == [{"notes-env", "notes-env-rolled-back-1"}] * 2 + assert context.metadata[TRAFFIC_KEY]["notes-env-rolled-back-1"]["artifact"] != \ + context.metadata[TRAFFIC_KEY]["notes-env"]["artifact"] From 29e406ab1ddee5ea97108516287d57dc2f0117fd Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 17:13:00 -0800 Subject: [PATCH 5/8] fix(website): teardown saves the captures it collected even when its time runs out partway (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_step/task_steps/traffic_capture.py | 23 ++++++++++--------- tst/unit/task_step/traffic_capture_test.py | 21 +++++++++++++++++ 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index 4a93e3919..7d19485f2 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -6,7 +6,6 @@ isn't lost with its VM, and saves what it recorded to the run's stored instance.""" from __future__ import annotations -import asyncio import functools import logging import os @@ -92,18 +91,20 @@ async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[ """Collect the captures of each of ``envs`` that serves traffic/v1 and whose captures the run hasn't collected, a collection that failed tried again, as teardown does before the envs close, and of each ``rolled_back`` deployment, a retry's, under its own key (``rolled_back_deployments``); ``collect_captures`` records what happened to each, and - the entries are saved to the run's stored instance, since teardown runs after its steps saved theirs.""" + the entries are saved to the run's stored instance, since teardown runs after its steps saved theirs, even when + teardown's time runs out partway, so what was collected by then stays linked to the run.""" entries = context.metadata.get(TRAFFIC_KEY) or {} tried = False - for key, deployed in [*((deployed.env_id, deployed) for deployed in envs), *rolled_back]: - if captures_traffic(deployed) and (entries.get(key) or {}).get("state") != "collected": - await collect_captures(context, deployed, key=key) - tried = True - if tried and context.instance_id: - from agent_env.task.store import record_context_metadata - - await asyncio.to_thread(record_context_metadata, context.instance_id, TRAFFIC_KEY, - context.metadata[TRAFFIC_KEY]) + try: + for key, deployed in [*((deployed.env_id, deployed) for deployed in envs), *rolled_back]: + if captures_traffic(deployed) and (entries.get(key) or {}).get("state") != "collected": + tried = True + await collect_captures(context, deployed, key=key) + finally: + if tried and context.instance_id and TRAFFIC_KEY in context.metadata: + from agent_env.task.store import record_context_metadata + + record_context_metadata(context.instance_id, TRAFFIC_KEY, context.metadata[TRAFFIC_KEY]) def rolled_back_deployments(context: TaskStepContext) -> list[tuple[str, DeployedEnv]]: diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index f3a5a8c0b..758830166 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -217,6 +217,27 @@ async def test_what_teardown_collects_is_saved_to_the_runs_stored_instance(gatew assert stored == context.metadata[TRAFFIC_KEY] and stored["notes-env"]["state"] == "collected" +@pytest.mark.asyncio +async def test_what_teardown_collected_before_its_time_ran_out_is_still_saved_to_the_run(gateway, monkeypatch): + instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) + context = _context(_deployed(), _deployed("slow-env")) + context.instance_id = instance.instance_id + real = traffic_capture_module.collect_captures + + async def collect(context, deployed, *, key=None): + if deployed.env_id == "slow-env": + await asyncio.sleep(30) + return await real(context, deployed, key=key) + + monkeypatch.setattr(traffic_capture_module, "collect_captures", collect) + with pytest.raises(TimeoutError): + async with asyncio.timeout(1): + await collect_uncollected_captures(context, context.deployed_envs) + + stored = find_task_instance(instance.instance_id).context["metadata"][TRAFFIC_KEY] + assert stored["notes-env"]["state"] == "collected" and "slow-env" not in stored + + @pytest.mark.asyncio async def test_tear_down_collects_a_retrys_rolled_back_deployment_under_its_own_key_before_its_vm_goes(gateway): rolled_back, live = _deployed(sandbox_id="vm-1"), _deployed(sandbox_id="vm-2") From 14ff94f8a853dde49b50b8dedba47b32e78f8e33 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 17:40:04 -0800 Subject: [PATCH 6/8] fix(website): a traffic collection stopped partway keeps the files it stored linked to the run as partial (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_step/task_steps/traffic_capture.py | 56 ++++++++++++------- tst/unit/task_step/traffic_capture_test.py | 29 ++++++++++ 2 files changed, 64 insertions(+), 21 deletions(-) diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index 7d19485f2..ae9a2fff1 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -122,7 +122,9 @@ async def collect_captures(context: TaskStepContext, deployed: DeployedEnv, *, k """Store each capture ``deployed`` serves as a pcap FileArtifact, register them as one FileArtifactUniverse, and record it in ``context.metadata[TRAFFIC_KEY][key]``, by default its env_id. Each file is named for its container and its place among that container's files, oldest first, so they read in the order their packets were captured. - Never raises: the entry, which it also returns, says what happened.""" + Never raises: the entry, which it also returns, says what happened. A collection that stops partway, failed or + cancelled, still registers the files it stored by then and records them as ``partial``, which teardown collects + again.""" key = key or deployed.env_id entries = context.metadata.setdefault(TRAFFIC_KEY, {}) try: @@ -133,7 +135,8 @@ async def collect_captures(context: TaskStepContext, deployed: DeployedEnv, *, k entries[key] = await _store(context, deployed, listing.captures, key) except Exception as e: logger.warning("collect_traffic_capture %s: %s", key, e) - entries[key] = _failed(e) + partial = entries.get(key) or {} + entries[key] = {**partial, **_failed(e), "state": "partial"} if partial.get("state") == "partial" else _failed(e) return entries[key] @@ -145,29 +148,40 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list store = get_artifact_store() prefix = store.attempt_prefix(_ARTIFACT_TYPE, universe_id) files, total, seen = {}, 0, {} - for capture in captures: - container = capture.container - seen[container] = seen.get(container, -1) + 1 - filename = f"{container}-{seen[container]:02d}.pcap" - handle, path = tempfile.mkstemp(suffix=".pcap") - os.close(handle) - try: - total += await fetch_traffic_capture(deployed.environment_url, deployed.environment_card, capture.name, - path, headers=deployed.control_headers) - except BaseException: - _remove(path) - raise - url = await finish_on_thread(functools.partial(_upload, universe_id, prefix + filename, path), - f"Uploading traffic capture {filename} of env {deployed.env_id!r}", - if_never_run=functools.partial(_remove, path)) - files[filename] = store.put_document(FileArtifact( - id=derive_id(universe_id, filename), filename=filename, content_type=PCAP_MEDIA_TYPE, object_url=url, - description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}")) + try: + for capture in captures: + container = capture.container + seen[container] = seen.get(container, -1) + 1 + filename = f"{container}-{seen[container]:02d}.pcap" + handle, path = tempfile.mkstemp(suffix=".pcap") + os.close(handle) + try: + total += await fetch_traffic_capture(deployed.environment_url, deployed.environment_card, capture.name, + path, headers=deployed.control_headers) + except BaseException: + _remove(path) + raise + url = await finish_on_thread(functools.partial(_upload, universe_id, prefix + filename, path), + f"Uploading traffic capture {filename} of env {deployed.env_id!r}", + if_never_run=functools.partial(_remove, path)) + files[filename] = store.put_document(FileArtifact( + id=derive_id(universe_id, filename), filename=filename, content_type=PCAP_MEDIA_TYPE, object_url=url, + description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}")) + except BaseException as e: + if files: + context.metadata[TRAFFIC_KEY][key] = {**_entry(universe_id, files, total, seen), "state": "partial", + "error": _failed(e)["error"]} + raise + logger.info("collect_traffic_capture %s: %d capture files, %d bytes", deployed.env_id, len(files), total) + return _entry(universe_id, files, total, seen) + + +def _entry(universe_id: str, files: dict, total: int, seen: dict) -> dict: + """The ``collected`` entry of ``files``, registered as one FileArtifactUniverse when there are any.""" artifact = None if files: universe = FileArtifactUniverse.put(id=universe_id, file_artifacts=files) artifact = {"id": universe.id, "version": universe.version} - logger.info("collect_traffic_capture %s: %d capture files, %d bytes", deployed.env_id, len(files), total) return {"state": "collected", "artifact": artifact, "files": sorted(files), "bytes": total, "containers": sorted(seen)} diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index 758830166..c5aad7190 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -238,6 +238,35 @@ async def collect(context, deployed, *, key=None): assert stored["notes-env"]["state"] == "collected" and "slow-env" not in stored +@pytest.mark.asyncio +async def test_a_collection_stopped_partway_keeps_the_files_it_stored_linked_as_partial(gateway, monkeypatch): + instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) + context = _context(_deployed()) + context.instance_id = instance.instance_id + real, fetched = traffic_capture_module.fetch_traffic_capture, [] + + async def fetch(*args, **kwargs): + fetched.append(args[2]) + if len(fetched) == 2: + await asyncio.sleep(30) + return await real(*args, **kwargs) + + monkeypatch.setattr(traffic_capture_module, "fetch_traffic_capture", fetch) + with pytest.raises(TimeoutError): + async with asyncio.timeout(1): + await collect_uncollected_captures(context, context.deployed_envs) + + entry = context.metadata[TRAFFIC_KEY]["notes-env"] + assert entry["state"] == "partial" and entry["files"] == ["notes-website-backend-00.pcap"] + universe = Artifact.get(entry["artifact"]["id"], version=entry["artifact"]["version"]) + assert sorted(universe.get_file_artifacts()) == ["notes-website-backend-00.pcap"] + assert find_task_instance(instance.instance_id).context["metadata"][TRAFFIC_KEY]["notes-env"] == entry + + monkeypatch.setattr(traffic_capture_module, "fetch_traffic_capture", real) + await collect_uncollected_captures(context, context.deployed_envs) + assert context.metadata[TRAFFIC_KEY]["notes-env"]["state"] == "collected" + + @pytest.mark.asyncio async def test_tear_down_collects_a_retrys_rolled_back_deployment_under_its_own_key_before_its_vm_goes(gateway): rolled_back, live = _deployed(sandbox_id="vm-1"), _deployed(sandbox_id="vm-2") From 46aa361311878b725c03f397277dc7c3eb7bfed8 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 20:03:33 -0800 Subject: [PATCH 7/8] fix(website): traffic collection makes every store write on a worker thread, so teardown's time limit holds (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_step/task_steps/traffic_capture.py | 70 +++++++++++++------ tst/unit/task_step/traffic_capture_test.py | 66 ++++++++++++++++- 2 files changed, 112 insertions(+), 24 deletions(-) diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index ae9a2fff1..072192990 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -6,6 +6,8 @@ isn't lost with its VM, and saves what it recorded to the run's stored instance.""" from __future__ import annotations +import asyncio +import copy import functools import logging import os @@ -39,6 +41,7 @@ #: ``context.metadata[TRAFFIC_KEY][env_id]``: each env whose captures a run collected, or tried to. TRAFFIC_KEY = "traffic_captures" _ARTIFACT_TYPE = "traffic_captures" +_SAVING: set[asyncio.Future] = set() _ERROR_CHARS = 500 @@ -100,11 +103,31 @@ async def collect_uncollected_captures(context: TaskStepContext, envs: Sequence[ if captures_traffic(deployed) and (entries.get(key) or {}).get("state") != "collected": tried = True await collect_captures(context, deployed, key=key) - finally: - if tried and context.instance_id and TRAFFIC_KEY in context.metadata: - from agent_env.task.store import record_context_metadata + except BaseException: + if tried: + _save_in_background(context) + raise + if tried: + await _save(context) + + +async def _save(context: TaskStepContext) -> None: + """Save the run's ``traffic_captures`` entries to its stored instance on a worker thread, which finishes the + write even when teardown's time runs out while it waits.""" + if context.instance_id and TRAFFIC_KEY in context.metadata: + from agent_env.task.store import record_context_metadata - record_context_metadata(context.instance_id, TRAFFIC_KEY, context.metadata[TRAFFIC_KEY]) + await finish_on_thread(functools.partial(record_context_metadata, context.instance_id, TRAFFIC_KEY, + copy.deepcopy(context.metadata[TRAFFIC_KEY])), + f"Saving the traffic captures of task instance {context.instance_id}") + + +def _save_in_background(context: TaskStepContext) -> None: + """Start ``_save`` without waiting for it, as a collection teardown's time ran out on does, so the entries + collected by then are saved while teardown goes on to terminate the sandboxes.""" + saving = asyncio.ensure_future(_save(context)) + _SAVING.add(saving) + saving.add_done_callback(_SAVING.discard) def rolled_back_deployments(context: TaskStepContext) -> list[tuple[str, DeployedEnv]]: @@ -122,9 +145,9 @@ async def collect_captures(context: TaskStepContext, deployed: DeployedEnv, *, k """Store each capture ``deployed`` serves as a pcap FileArtifact, register them as one FileArtifactUniverse, and record it in ``context.metadata[TRAFFIC_KEY][key]``, by default its env_id. Each file is named for its container and its place among that container's files, oldest first, so they read in the order their packets were captured. - Never raises: the entry, which it also returns, says what happened. A collection that stops partway, failed or - cancelled, still registers the files it stored by then and records them as ``partial``, which teardown collects - again.""" + Never raises: the entry, which it also returns, says what happened. Each file stored is registered in the universe + as it lands, every store write on a worker thread, so a collection that stops partway, failed or cancelled, keeps + the files it stored by then recorded as ``partial``, which teardown collects again.""" key = key or deployed.env_id entries = context.metadata.setdefault(TRAFFIC_KEY, {}) try: @@ -147,7 +170,7 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list get_config().check_local_run_write(universe_id) store = get_artifact_store() prefix = store.attempt_prefix(_ARTIFACT_TYPE, universe_id) - files, total, seen = {}, 0, {} + files, total, seen, artifact = {}, 0, {}, None try: for capture in captures: container = capture.container @@ -164,26 +187,31 @@ async def _store(context: TaskStepContext, deployed: DeployedEnv, captures: list url = await finish_on_thread(functools.partial(_upload, universe_id, prefix + filename, path), f"Uploading traffic capture {filename} of env {deployed.env_id!r}", if_never_run=functools.partial(_remove, path)) - files[filename] = store.put_document(FileArtifact( + document = FileArtifact( id=derive_id(universe_id, filename), filename=filename, content_type=PCAP_MEDIA_TYPE, object_url=url, - description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}")) + description=f"Traffic capture of {container} in env {deployed.env_id!r}, task instance {context.instance_id}") + files[filename] = await finish_on_thread(functools.partial(store.put_document, document), + f"Storing traffic capture {filename} of env {deployed.env_id!r}") + artifact = await finish_on_thread(functools.partial(_register, universe_id, dict(files)), + f"Registering the traffic captures of env {deployed.env_id!r}") + context.metadata[TRAFFIC_KEY][key] = _entry("partial", artifact, files, total, seen) except BaseException as e: - if files: - context.metadata[TRAFFIC_KEY][key] = {**_entry(universe_id, files, total, seen), "state": "partial", + if artifact is not None: + context.metadata[TRAFFIC_KEY][key] = {**_entry("partial", artifact, files, total, seen), "error": _failed(e)["error"]} raise logger.info("collect_traffic_capture %s: %d capture files, %d bytes", deployed.env_id, len(files), total) - return _entry(universe_id, files, total, seen) + return _entry("collected", artifact, files, total, seen) + + +def _register(universe_id: str, files: dict) -> dict: + """Register ``files`` as the next version of the FileArtifactUniverse ``universe_id``.""" + universe = FileArtifactUniverse.put(id=universe_id, file_artifacts=files) + return {"id": universe.id, "version": universe.version} -def _entry(universe_id: str, files: dict, total: int, seen: dict) -> dict: - """The ``collected`` entry of ``files``, registered as one FileArtifactUniverse when there are any.""" - artifact = None - if files: - universe = FileArtifactUniverse.put(id=universe_id, file_artifacts=files) - artifact = {"id": universe.id, "version": universe.version} - return {"state": "collected", "artifact": artifact, "files": sorted(files), "bytes": total, - "containers": sorted(seen)} +def _entry(state: str, artifact: Optional[dict], files: dict, total: int, seen: dict) -> dict: + return {"state": state, "artifact": artifact, "files": sorted(files), "bytes": total, "containers": sorted(seen)} def _upload(universe_id: str, object_url: str, path: str) -> str: diff --git a/tst/unit/task_step/traffic_capture_test.py b/tst/unit/task_step/traffic_capture_test.py index c5aad7190..025437f60 100644 --- a/tst/unit/task_step/traffic_capture_test.py +++ b/tst/unit/task_step/traffic_capture_test.py @@ -8,12 +8,13 @@ import asyncio import os import threading +import time from unittest.mock import AsyncMock, patch import httpx import pytest -from agent_env.artifact import Artifact +from agent_env.artifact import Artifact, FileArtifact from agent_env.artifact.store import get_artifact_store from agent_env.env.env import DeployedGatewayEnv from agent_env.env.gateway import gateway as gw_mod @@ -217,6 +218,16 @@ async def test_what_teardown_collects_is_saved_to_the_runs_stored_instance(gatew assert stored == context.metadata[TRAFFIC_KEY] and stored["notes-env"]["state"] == "collected" +async def _saved(instance_id: str, timeout_s: float = 10.0) -> dict: + """The run's saved ``traffic_captures`` entries, once a save teardown left running has landed.""" + deadline = asyncio.get_running_loop().time() + timeout_s + while True: + stored = find_task_instance(instance_id).context.get("metadata", {}).get(TRAFFIC_KEY) + if stored is not None or asyncio.get_running_loop().time() > deadline: + return stored + await asyncio.sleep(0.05) + + @pytest.mark.asyncio async def test_what_teardown_collected_before_its_time_ran_out_is_still_saved_to_the_run(gateway, monkeypatch): instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) @@ -234,7 +245,7 @@ async def collect(context, deployed, *, key=None): async with asyncio.timeout(1): await collect_uncollected_captures(context, context.deployed_envs) - stored = find_task_instance(instance.instance_id).context["metadata"][TRAFFIC_KEY] + stored = await _saved(instance.instance_id) assert stored["notes-env"]["state"] == "collected" and "slow-env" not in stored @@ -260,13 +271,62 @@ async def fetch(*args, **kwargs): assert entry["state"] == "partial" and entry["files"] == ["notes-website-backend-00.pcap"] universe = Artifact.get(entry["artifact"]["id"], version=entry["artifact"]["version"]) assert sorted(universe.get_file_artifacts()) == ["notes-website-backend-00.pcap"] - assert find_task_instance(instance.instance_id).context["metadata"][TRAFFIC_KEY]["notes-env"] == entry + assert (await _saved(instance.instance_id))["notes-env"] == entry monkeypatch.setattr(traffic_capture_module, "fetch_traffic_capture", real) await collect_uncollected_captures(context, context.deployed_envs) assert context.metadata[TRAFFIC_KEY]["notes-env"]["state"] == "collected" +@pytest.mark.asyncio +async def test_a_slow_document_store_doesnt_hold_teardowns_time_and_what_was_stored_still_lands(gateway, monkeypatch): + instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) + context = _context(_deployed()) + context.instance_id = instance.instance_id + store = get_artifact_store() + real, stored = store.put_document, [] + + def slow_put(document): + if isinstance(document, FileArtifact): + if stored: + time.sleep(2) + stored.append(document.id) + return real(document) + + monkeypatch.setattr(store, "put_document", slow_put) + started = time.monotonic() + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.5): + await collect_uncollected_captures(context, context.deployed_envs) + + assert time.monotonic() - started < 1.5 + entry = (await _saved(instance.instance_id))["notes-env"] + assert entry["state"] == "partial" and entry["files"] == ["notes-website-backend-00.pcap"] + + +@pytest.mark.asyncio +async def test_a_slow_run_save_doesnt_hold_teardowns_time_and_still_lands(gateway, monkeypatch): + from agent_env.task import store as task_store + + instance = get_task_instance_store().create_instance("task-1", 1, total_steps=1) + context = _context(_deployed()) + context.instance_id = instance.instance_id + real = task_store.record_context_metadata + + def slow_record(*args): + time.sleep(2) + real(*args) + + monkeypatch.setattr(task_store, "record_context_metadata", slow_record) + started = time.monotonic() + with pytest.raises(TimeoutError): + async with asyncio.timeout(1): + await collect_uncollected_captures(context, context.deployed_envs) + + assert time.monotonic() - started < 2 + assert (await _saved(instance.instance_id))["notes-env"]["state"] == "collected" + + @pytest.mark.asyncio async def test_tear_down_collects_a_retrys_rolled_back_deployment_under_its_own_key_before_its_vm_goes(gateway): rolled_back, live = _deployed(sandbox_id="vm-1"), _deployed(sandbox_id="vm-2") From ffa4b7b1282f3d48098cb2c2e41b83b1d4b92eeb Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 3 Nov 2026 21:22:59 -0800 Subject: [PATCH 8/8] docs(website): traffic capture records what a website's containers send from just after each starts, not before (FD-4683) Co-Authored-By: Claude Opus 5.5 (1M context) --- AGENTS.md | 2 +- src/agent_env/cli/env/deploy.py | 4 ++-- src/agent_env/env/gateway/constants.py | 2 +- .../providers/env_providers/env_gateway_provider.py | 9 +++++---- src/agent_env/task_step/task_steps/traffic_capture.py | 2 +- 5 files changed, 10 insertions(+), 9 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index d1971eb4d..5f7c4e8e8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -117,7 +117,7 @@ and deserialized through a registry. | Package | Contents | |---|---| -| `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`, one site: a server image, `backend_image` in its env.toml, serving the pages and API, ready when its card's `website/v1` `ready` answers a `WebsiteReadiness` that says so, as `agentenv_protocol.website` (`website_readiness`, `env_readiness`) defines a website env's health for the deploy and `DeployedEnv.readiness()` alike, or, for a card without `ready`, when its `health` path answers 200, which is deprecated; the gateway container's `env/gateway/constants.py` mirrors the URI and op, a test keeping them equal; a separate frontend image is deprecated, warned of at put and optional in a bundle, and a stored `backend_docker_image_artifact` reads as the server image), `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`), a run only the forms its sandbox's deploys boot (`booted_forms`: on each link, the first form it boots natively, by the deploy's own rule, `boot.first_native`; a compose of a disk adds its container), and a deploy composes first what it boots that no compose started (`ComputerEnv.compose_unstarted`); 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 (`utils/probing.py`) 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, over its env.toml's `[settings]` and a computer env's `settings` beside it in `apps`, each by top-level field (`params.app_settings`; checked at put, a required one left to the deploy; the deploy record keeps the merged values as `app_settings`, secrets withheld); 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. Its site route, `/website//`, never forwards a site's data plane or card (`/agentenv/…`, `/.well-known/agent-env.json`), which agentenv reaches through `/svc//`; a site with no frontend image is served by its server image there. A deploy with `capture_traffic` (`deploy_env`'s, `env deploy --capture-traffic`; a website env, or a multi env with a website, on a VM, refused for any other env and in container mode) records every packet each website's server and frontend containers send and receive: beside each runs a capture service, tcpdump from the gateway's image in that container's network namespace (`network_mode: "service:"`), the only service holding `NET_RAW` and `NET_ADMIN`, rotating its classic pcap files (10 of 20 MB each, the oldest overwritten) onto a volume no website container mounts and the gateway mounts read-only and serves as `urn:agentenv:traffic/v1` (`list`, `GET /traffic/captures`, oldest first; `get`, `GET /traffic/capture?name=`), which `agentenv_protocol.traffic` defines (`TrafficCapture`, `list_traffic_captures`, `fetch_traffic_capture`), the gateway's `constants.py` mirroring its names, a test keeping them equal. Those routes take the per-deploy control token the record keeps (`DeployedEnv.control_token`, the gateway's `AGENTENV_CONTROL_TOKEN`), which the gateway takes off every request before it proxies one, so no site reads it; a capture that isn't running once the gateway is up refuses the deploy. 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`. | +| `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`, one site: a server image, `backend_image` in its env.toml, serving the pages and API, ready when its card's `website/v1` `ready` answers a `WebsiteReadiness` that says so, as `agentenv_protocol.website` (`website_readiness`, `env_readiness`) defines a website env's health for the deploy and `DeployedEnv.readiness()` alike, or, for a card without `ready`, when its `health` path answers 200, which is deprecated; the gateway container's `env/gateway/constants.py` mirrors the URI and op, a test keeping them equal; a separate frontend image is deprecated, warned of at put and optional in a bundle, and a stored `backend_docker_image_artifact` reads as the server image), `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`), a run only the forms its sandbox's deploys boot (`booted_forms`: on each link, the first form it boots natively, by the deploy's own rule, `boot.first_native`; a compose of a disk adds its container), and a deploy composes first what it boots that no compose started (`ComputerEnv.compose_unstarted`); 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 (`utils/probing.py`) 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, over its env.toml's `[settings]` and a computer env's `settings` beside it in `apps`, each by top-level field (`params.app_settings`; checked at put, a required one left to the deploy; the deploy record keeps the merged values as `app_settings`, secrets withheld); 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. Its site route, `/website//`, never forwards a site's data plane or card (`/agentenv/…`, `/.well-known/agent-env.json`), which agentenv reaches through `/svc//`; a site with no frontend image is served by its server image there. A deploy with `capture_traffic` (`deploy_env`'s, `env deploy --capture-traffic`; a website env, or a multi env with a website, on a VM, refused for any other env and in container mode) records every packet each website's server and frontend containers send and receive from just after each starts (what a server sends while it starts up can come before its capture attaches): beside each runs a capture service, tcpdump from the gateway's image in that container's network namespace (`network_mode: "service:"`), the only service holding `NET_RAW` and `NET_ADMIN`, rotating its classic pcap files (10 of 20 MB each, the oldest overwritten) onto a volume no website container mounts and the gateway mounts read-only and serves as `urn:agentenv:traffic/v1` (`list`, `GET /traffic/captures`, oldest first; `get`, `GET /traffic/capture?name=`), which `agentenv_protocol.traffic` defines (`TrafficCapture`, `list_traffic_captures`, `fetch_traffic_capture`), the gateway's `constants.py` mirroring its names, a test keeping them equal. Those routes take the per-deploy control token the record keeps (`DeployedEnv.control_token`, the gateway's `AGENTENV_CONTROL_TOKEN`), which the gateway takes off every request before it proxies one, so no site reads it; a capture that isn't running once the gateway is up refuses the deploy. 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/`, among them `start_computer_recording` and `end_computer_recording`, `collect_traffic_capture`, which stores each traffic/v1 capture of an env deployed with `capture_traffic` as a pcap `FileArtifact` in one `FileArtifactUniverse` recorded under `context.metadata["traffic_captures"][env_id]`, 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 capturing env still up whose traffic the run didn't collect is collected, a retry's rolled-back ones too (kept as `capturing_envs` in its `retry_resets` marker, each under `-rolled-back-`), and what it records saved to the stored instance's `context.metadata`, 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 namespaces, `@local` and `@hf` (`is_local_id`), each with a documents file of its own: while the CLI or a `build_app` server runs, a local id and what a 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/`. | diff --git a/src/agent_env/cli/env/deploy.py b/src/agent_env/cli/env/deploy.py index cdd17e6fa..c68103da7 100644 --- a/src/agent_env/cli/env/deploy.py +++ b/src/agent_env/cli/env/deploy.py @@ -52,8 +52,8 @@ help="Draw a role key for the deploy: its gateway then believes a role other than default only from a " "caller that sends the key beside it.") @click.option("--capture-traffic", is_flag=True, default=False, - help="Record every packet the env's websites' containers send and receive, as rotating pcap files its " - "gateway serves as traffic/v1 (VM deploys of website or multi envs only).") + help="Record every packet the env's websites' containers send and receive from just after each starts, as " + "rotating pcap files its gateway serves as traffic/v1 (VM deploys of website or multi envs only).") def deploy(env_id: str, env_version: int | None, ttl_seconds: int, gateway_mode: str, sandbox: str, service_db_env_id: str | None, gateway_env_id: str | None, website_browser_env_id: str | None, disk_size_gb: float, cpu: float | None, memory_mb: int | None, env_state_type: str | None, env_state_instance_id: str | None, keyed_roles: bool, capture_traffic: bool): """Deploy an MCP server environment.""" from agent_env.providers import build_sandbox_provider diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index 0af3614a2..03e910a3c 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/src/agent_env/env/gateway/constants.py @@ -280,7 +280,7 @@ def random_mcp_server_name() -> str: ] TRAFFIC_EXTENSION = { "uri": EXT_TRAFFIC_URI, - "description": "Traffic (urn:agentenv:traffic/v1): the packets each website's server and frontend containers sent and received, captured by tcpdump in each container's network namespace as rotating classic pcap files. list names each capture file with the container it holds, its size and when it was last written, oldest first; get streams one as application/vnd.tcpdump.pcap. A container's oldest file is overwritten once its rotation is full. Harness control path: with a control token, each call takes it as its bearer.", + "description": "Traffic (urn:agentenv:traffic/v1): the packets each website's server and frontend containers sent and received, captured by tcpdump in each container's network namespace, from just after the container starts, as rotating classic pcap files. list names each capture file with the container it holds, its size and when it was last written, oldest first; get streams one as application/vnd.tcpdump.pcap. A container's oldest file is overwritten once its rotation is full. Harness control path: with a control token, each call takes it as its bearer.", "params": {"endpoint": "/traffic/captures", "format": "pcap", "methods": { TRAFFIC_LIST_OP: {"method": "GET", "endpoint": "/traffic/captures", "request": {"type": "object", "properties": {}}}, TRAFFIC_GET_OP: {"method": "GET", "endpoint": "/traffic/capture", "request": {"type": "object", "properties": { diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index 7a877706c..ee7aa6f28 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -751,8 +751,8 @@ async def create_gateway( the record: the gateway then believes a role other than ``default`` only beside it, so every caller acting as one must send it. Off, as before the key, the role header alone is believed. - ``capture_traffic`` records every packet each website's server and frontend containers send and receive, with - tcpdump in a capture service beside each, as rotating pcap files the gateway serves as traffic/v1, behind a + ``capture_traffic`` records every packet each website's server and frontend containers send and receive once + its capture attaches, just after the container starts, with tcpdump in a capture service beside each, as rotating pcap files the gateway serves as traffic/v1, behind a control token the deploy draws for it and keeps on the record (``control_token``). VM deploys only, since only they serve websites; off, the stack is as before it. """ @@ -1668,8 +1668,9 @@ def _traffic_capture_service(gateway_image: str, container: str) -> list[str]: interface of the container's, its loopback among them, every packet whole and written as it's captured, rotating through ``TRAFFIC_CAPTURE_FILES`` files of ``TRAFFIC_CAPTURE_FILE_MB`` million bytes named for the container, its oldest overwritten once they're full, so a long run holds at most their product per container on the VM's disk. It joins - the container's network namespace, so it sees what the container sends and receives and nothing else, and is the - only service given the capabilities a capture needs. It writes to ``TRAFFIC_VOLUME``, which the gateway mounts + the container's network namespace, so it sees what the container sends and receives and nothing else, from when it + attaches, which is just after the container starts: what a server sends while it starts up can come before it, and + isn't recorded. It is the only service given the capabilities a capture needs. It writes to ``TRAFFIC_VOLUME``, which the gateway mounts read-only and no website container mounts at all.""" argv = ["-i", "any", "-n", "-s", "0", "-U", "-Z", "root", "-C", str(TRAFFIC_CAPTURE_FILE_MB), "-W", str(TRAFFIC_CAPTURE_FILES), "-w", f"{TRAFFIC_SIDECAR_DIR}/{container}.pcap"] diff --git a/src/agent_env/task_step/task_steps/traffic_capture.py b/src/agent_env/task_step/task_steps/traffic_capture.py index 072192990..f6dd48cd3 100644 --- a/src/agent_env/task_step/task_steps/traffic_capture.py +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -1,5 +1,5 @@ """Collect a website env's traffic captures. A website or multi env deployed with ``capture_traffic`` records every -packet its websites' containers send and receive as rotating pcap files, which its gateway's card serves as traffic/v1 +packet its websites' containers send and receive, from just after each starts, as rotating pcap files, which its gateway's card serves as traffic/v1 behind the deploy's control token. ``collect_traffic_capture`` stores each as a FileArtifact and registers them as one FileArtifactUniverse, recorded in ``context.metadata[TRAFFIC_KEY][env_id]``. Teardown collects any capturing env a run didn't, a retry's rolled-back ones included, before the env closes (``collect_uncollected_captures``), so a capture