diff --git a/AGENTS.md b/AGENTS.md index a8b4d6a4..5f7c4e8e 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 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`, 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, 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/packages/agentenv-protocol/examples/website/notes/README.md b/packages/agentenv-protocol/examples/website/notes/README.md index 861d6d64..0d55da1c 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 61f9cbab..fc67f4f3 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 00000000..98289a64 --- /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 00000000..4c09e3a0 --- /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 f7afb104..c68103da 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 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 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 7659455b..6dc06a73 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 d1da0700..3789e308 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 c70a5e20..48848e21 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 73000f01..4957dd39 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 784eb080..03e910a3 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, 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": { + "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 2912fbdb..451f718e 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 283fd90a..427b6320 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -9,8 +9,10 @@ import logging import os import posixpath +import re import socket import subprocess +import tempfile from collections.abc import Iterable from contextlib import asynccontextmanager from datetime import datetime, timezone @@ -62,12 +64,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 +106,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): @@ -129,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, @@ -142,11 +151,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 +175,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 +433,16 @@ 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 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: return await self._serve_changes(request) @@ -457,6 +486,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 +506,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 +868,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 +1613,84 @@ 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, 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) + 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) + 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 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}"', + "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.""" 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 58fee59a..ee7aa6f2 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, @@ -83,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 @@ -95,6 +98,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 +133,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 +309,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 +324,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 +416,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 +493,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 +534,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 +554,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 +612,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 +641,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 +661,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 +690,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 +708,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 +732,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 +750,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 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. """ 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 +825,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 +919,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 +968,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 +1018,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 +1266,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()[-_CAPTURE_CHECK_ERROR_CHARS:]}") + 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 +1496,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 +1512,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 +1535,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 +1648,48 @@ 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 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, 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"] + 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 +1737,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 +1750,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 fc5a7753..875d0994 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/store.py b/src/agent_env/task/store.py index 869db45d..1a720a82 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 6840a58d..c6ce5702 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 9f11627c..2468b70c 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, rolled_back_deployments logger = logging.getLogger(__name__) @@ -116,7 +118,9 @@ 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, 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``.""" @@ -129,8 +133,15 @@ 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))] + 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, 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)) 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 0c8ebd7a..f482a64b 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 22530034..d12609f7 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 00000000..f6dd48cd --- /dev/null +++ b/src/agent_env/task_step/task_steps/traffic_capture.py @@ -0,0 +1,235 @@ +"""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, 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 +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 +import tempfile +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 +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" +_SAVING: set[asyncio.Future] = set() +_ERROR_CHARS = 500 + + +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], *, + 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, 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, 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 + 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) + 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 + + 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]]: + """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][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. 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: + 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[key] = await _store(context, deployed, listing.captures, key) + except Exception as e: + logger.warning("collect_traffic_capture %s: %s", key, e) + partial = entries.get(key) or {} + entries[key] = {**partial, **_failed(e), "state": "partial"} if partial.get("state") == "partial" else _failed(e) + return entries[key] + + +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) + store = get_artifact_store() + prefix = store.attempt_prefix(_ARTIFACT_TYPE, universe_id) + files, total, seen, artifact = {}, 0, {}, None + 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)) + 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}") + 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 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("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(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: + """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}"[:_ERROR_CHARS]} diff --git a/tst/unit/env/envs/test_multi_env_name.py b/tst/unit/env/envs/test_multi_env_name.py index 70a9ea22..7cbc360f 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 00000000..17f11abd --- /dev/null +++ b/tst/unit/env/gateway/traffic_test.py @@ -0,0 +1,262 @@ +"""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") + + +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), + (".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 76ecac6b..20c3df49 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 00000000..6d75cf61 --- /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", "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 + + +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 00000000..46e40be8 --- /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/retry_test.py b/tst/unit/task_step/retry_test.py index 46c67094..66dd3ea8 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 new file mode 100644 index 00000000..025437f6 --- /dev/null +++ b/tst/unit/task_step/traffic_capture_test.py @@ -0,0 +1,349 @@ +"""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 asyncio +import os +import threading +import time +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +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 +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 +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, + 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, + 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=sandbox_id, + 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_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))) + 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_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}, + "notes-env": {"state": "failed", "error": "ConnectError: refused"}} + + 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 + + +@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" + + +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) + 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 = await _saved(instance.instance_id) + 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 (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") + 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"]