diff --git a/.env.example b/.env.example index 53d6564..e113b00 100644 --- a/.env.example +++ b/.env.example @@ -39,6 +39,11 @@ SANDBOX_REDIS_URL=redis://127.0.0.1:6379/0 SANDBOX_LOCAL_ROOT=./data/sandboxes SANDBOX_SHARED_ROOT= SANDBOX_IDLE_TTL_SECONDS=1800 +SANDBOX_SUSPENDED_RETENTION_SECONDS=2592000 +SANDBOX_SUSPEND_SNAPSHOT=auto +# A workspace directory this worker no longer owns (resumed on another worker, +# or released while this worker was away) is deleted after this long; 0 keeps it. +SANDBOX_ORPHAN_DORMANT_DIR_TTL_SECONDS=86400 # Concurrent commands per sandbox; only executions with a distinct exec_scope # run in parallel, and an unscoped command stays exclusive. SANDBOX_MAX_PARALLEL_EXECS_PER_SANDBOX=16 diff --git a/.github/ci/compose.apparmor.yaml b/.github/ci/compose.apparmor.yaml new file mode 100644 index 0000000..964702a --- /dev/null +++ b/.github/ci/compose.apparmor.yaml @@ -0,0 +1,13 @@ +# CI only. Layered over compose.yaml on the GitHub-hosted runners, whose Docker +# enforces its `docker-default` AppArmor profile. That profile denies the +# `mount` Bubblewrap performs to make `/` a slave mount inside its new mount +# namespace, so every isolation level fails the startup probe with +# `bwrap: Failed to make / slave: Permission denied` and the service exits. +# The Docker Desktop and OrbStack hosts the quick start is written against have +# no AppArmor, which is why compose.yaml itself does not carry this: the +# Compose file stays the narrow default, and a host that enforces AppArmor opts +# in deliberately (docs/ISOLATION.md, "Hosts that enforce AppArmor"). +services: + agent-sandbox: + security_opt: + - apparmor=unconfined diff --git a/.github/ci/compose.fleet.apparmor.yaml b/.github/ci/compose.fleet.apparmor.yaml new file mode 100644 index 0000000..16b6743 --- /dev/null +++ b/.github/ci/compose.fleet.apparmor.yaml @@ -0,0 +1,12 @@ +# CI only. See compose.apparmor.yaml: the same Docker AppArmor profile, applied +# to each replica of the fleet file. +services: + worker-a: + security_opt: + - apparmor=unconfined + worker-b: + security_opt: + - apparmor=unconfined + worker-c: + security_opt: + - apparmor=unconfined diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b5467fb..8efabe4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -70,6 +70,17 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 + - name: Allow unprivileged user namespaces + # ubuntu-latest is Ubuntu 24.04, which ships + # kernel.apparmor_restrict_unprivileged_userns=1: a process that creates + # a user namespace loses its capabilities inside it, so Bubblewrap cannot + # write its uid map ("setting up uid map: Permission denied") and no + # isolation level can be negotiated. That is the host this job runs on, + # not something the code under test can work around; the service reports + # it as a preflight warning. Switching the restriction off for the + # duration of the job is what the sandbox's own documentation tells an + # operator to do on such a host. + run: sudo sysctl -w kernel.apparmor_restrict_unprivileged_userns=0 - name: Build the base image run: docker build -f Dockerfile.base -t agent-sandbox-base:latest . - name: Build the test image @@ -114,13 +125,24 @@ jobs: run: python -m pip install uv - name: Install project run: uv sync --frozen + - name: Allow unprivileged user namespaces + # ubuntu-latest is Ubuntu 24.04, which ships + # kernel.apparmor_restrict_unprivileged_userns=1: a process that creates + # a user namespace loses its capabilities inside it, so Bubblewrap cannot + # write its uid map ("setting up uid map: Permission denied") and no + # isolation level can be negotiated. That is the host this job runs on, + # not something the code under test can work around; the service reports + # it as a preflight warning. Switching the restriction off for the + # duration of the job is what the sandbox's own documentation tells an + # operator to do on such a host. + run: sudo sysctl -w kernel.apparmor_restrict_unprivileged_userns=0 - name: Build the base image run: docker build -f Dockerfile.base -t agent-sandbox-base:latest . - name: Start replicas over a shared database and registry # The third replica runs with a different SANDBOX_PROFILE_HASH, behind the # fleet file's `other-profile` profile, so the checks that depend on a # fleet that is not uniform have one to look at. - run: docker compose -f compose.fleet.yaml --profile other-profile up -d --wait --build + run: docker compose -f compose.fleet.yaml -f .github/ci/compose.fleet.apparmor.yaml --profile other-profile up -d --wait --build - name: Verify what a fleet promises # The single-replica deployment job cannot see any of this: a request # forwarded to the replica that owns the sandbox, a template built on one @@ -132,10 +154,10 @@ jobs: run: uv run python scripts/verify-fleet.py --strict --short-graces - name: Fleet logs on failure if: failure() - run: docker compose -f compose.fleet.yaml --profile other-profile logs --tail 200 + run: docker compose -f compose.fleet.yaml -f .github/ci/compose.fleet.apparmor.yaml --profile other-profile logs --tail 200 - name: Stop if: always() - run: docker compose -f compose.fleet.yaml --profile other-profile down -v + run: docker compose -f compose.fleet.yaml -f .github/ci/compose.fleet.apparmor.yaml --profile other-profile down -v package: name: Release artifacts @@ -192,6 +214,14 @@ jobs: deployment: name: Deployment verification runs-on: ubuntu-latest + env: + # The README's compose.yaml plus the one AppArmor opt-in this runner's + # Docker needs; see the file for why it is not in compose.yaml itself. + COMPOSE_FILE: compose.yaml:.github/ci/compose.apparmor.yaml + # Job level because Compose interpolates the whole file whenever it runs, + # including `config --images`, `logs` and `down`, and a missing token is an + # error there rather than an empty string. + SANDBOX_INTERNAL_TOKEN: ci-verification-token steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -201,14 +231,23 @@ jobs: run: python -m pip install uv - name: Install project run: uv sync --frozen + - name: Allow unprivileged user namespaces + # ubuntu-latest is Ubuntu 24.04, which ships + # kernel.apparmor_restrict_unprivileged_userns=1: a process that creates + # a user namespace loses its capabilities inside it, so Bubblewrap cannot + # write its uid map ("setting up uid map: Permission denied") and no + # isolation level can be negotiated. That is the host this job runs on, + # not something the code under test can work around; the service reports + # it as a preflight warning. Switching the restriction off for the + # duration of the job is what the sandbox's own documentation tells an + # operator to do on such a host. + run: sudo sysctl -w kernel.apparmor_restrict_unprivileged_userns=0 - name: Build the base image # The two commands the README gives a first-time reader, in order. The # service image is FROM agent-sandbox-base, so `compose up` alone fails # on a machine that has never built it. run: docker build -f Dockerfile.base -t agent-sandbox-base:latest . - name: Start the deployment the README describes - env: - SANDBOX_INTERNAL_TOKEN: ci-verification-token run: docker compose up -d --wait --build - name: Scan the image for internal markers # What the build context put in the image, rather than the whole @@ -216,8 +255,6 @@ jobs: # produce tens of thousands of findings on their own. The script targets # /app and refuses to run when /app holds no application code, so a # "clean" answer cannot mean "looked at nothing". - env: - SANDBOX_INTERNAL_TOKEN: ci-verification-token run: ./scripts/scan-image.sh - name: Prove the image scan fails on an image that carries one # A gate that has only ever been seen pass may be wired to nothing. This @@ -252,8 +289,6 @@ jobs: # --strict so a skipped check fails: a deployment that verified half of # what it claims is not verified. This is also the only thing that runs # `examples/`, so an example cannot rot into a script that misleads. - env: - SANDBOX_INTERNAL_TOKEN: ci-verification-token run: uv run python scripts/verify-deployment.py --strict - name: Deployment logs on failure if: failure() diff --git a/.github/workflows/polyglot.yml b/.github/workflows/polyglot.yml index 33186d2..1f63075 100644 --- a/.github/workflows/polyglot.yml +++ b/.github/workflows/polyglot.yml @@ -43,6 +43,17 @@ jobs: run: python -m pip install uv - name: Install project run: uv sync --frozen + - name: Allow unprivileged user namespaces + # ubuntu-latest is Ubuntu 24.04, which ships + # kernel.apparmor_restrict_unprivileged_userns=1: a process that creates + # a user namespace loses its capabilities inside it, so Bubblewrap cannot + # write its uid map ("setting up uid map: Permission denied") and no + # isolation level can be negotiated. That is the host this job runs on, + # not something the code under test can work around; the service reports + # it as a preflight warning. Switching the restriction off for the + # duration of the job is what the sandbox's own documentation tells an + # operator to do on such a host. + run: sudo sysctl -w kernel.apparmor_restrict_unprivileged_userns=0 - name: Free the disk the toolchain images need # Three images of roughly two gigabytes each, plus the build cache that # produced them, against a runner whose root filesystem also carries an diff --git a/README.md b/README.md index aa2998b..b547e90 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,8 @@ execution-backend SPI for isolate, VM, or remote backends. See reusable environments cross workers on, and where a caller keeps checkpoint archives. Snapshot and restore are the caller's to orchestrate. - Generation fencing so stale workers cannot mutate reassigned workspaces. +- Suspend and idempotent resume: a waiting agent gives its slot back and keeps + its workspace, with a retention period and disk reclamation for dormant ones. - REST and stateless Streamable HTTP MCP APIs. - Reusable environment templates, built once and mounted read-only everywhere. - A dependency-free fleet console at `/admin`: no build step, no CDN. @@ -481,6 +483,30 @@ File API calls lock only the target path, so reads and writes to different files stay concurrent with running commands. Writes go through a temporary file and an atomic rename, so a reader never observes a partial file. +## Suspend and resume + +An agent waiting on something slow does not need to hold a slot. Suspend +releases the sandbox's capacity slot and keeps its workspace; resume takes a +slot again and continues with the same files, idempotently. + +```bash +# Give the slot back while waiting. Refused with 409 while anything runs. +POST /api/v1/sandboxes/my-agent/suspend {"generation":3} + +# Continue later. Idempotent; use the generation it returns from then on. +POST /api/v1/sandboxes/my-agent/resume +``` + +Suspend and exec admission are decided on the same route row, so a running +command refuses the suspend (`409 SANDBOX_SUSPEND_BUSY`) and a suspended sandbox +refuses commands (`409 SANDBOX_SUSPENDED`); other sandboxes and scopes are never +touched. Resume reuses the directory on the original worker when it can, and +otherwise moves the sandbox to another worker from a shared workspace or an +object-store snapshot. Suspended sandboxes are released after +`SANDBOX_SUSPENDED_RETENTION_SECONDS` (30 days by default). No process is +frozen and no memory is restored. See [Suspend and resume](docs/SUSPEND_RESUME.md) +for the state machine, snapshots, and disk reclamation. + ## Environment templates Every sandbox starts with an empty `/envs`, so an application that needs NumPy, diff --git a/README.zh-CN.md b/README.zh-CN.md index ccdd0c6..afb8688 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -43,6 +43,7 @@ Sandloom 是面向 AI Agent 的高密度执行服务。一个可信的 Worker - 支持内存或 Redis Worker 注册与发现。 - 支持兼容 S3 的对象存储，使用标准 AWS 凭据链，用于跨 Worker 分发可复用环境以及保存调用方的检查点归档；快照和恢复流程由调用方编排。 - 通过 generation 隔离令牌，阻止旧 Worker 修改已经重新分配的工作空间。 +- 挂起与幂等恢复：等待中的 Agent 归还名额并保留工作空间，休眠工作空间有保留期限和磁盘回收。 - 提供 REST 和无状态 Streamable HTTP MCP API。 - 环境模板构建一次，即可在多个沙箱中以只读方式挂载。 - `/admin` 提供无依赖的集群管理页面，无需前端构建，也不依赖 CDN。 @@ -308,6 +309,20 @@ POST /api/v1/sandboxes/my-agent/files/move {"generation":1,"source":"/workspac 文件 API 只锁定目标路径，因此不同文件的读写可以与正在执行的命令并行。写入通过临时文件和原子重命名完成，读取者不会看到写了一半的文件。 +## 挂起与恢复 + +等待耗时任务的 Agent 无需一直占用名额。挂起会释放沙箱的容量名额并保留工作空间；恢复会重新申请名额，并以幂等方式接着使用原来的文件继续工作。 + +```bash +# 等待期间归还名额。有命令在运行时返回 409。 +POST /api/v1/sandboxes/my-agent/suspend {"generation":3} + +# 稍后继续。幂等；之后使用返回的 generation。 +POST /api/v1/sandboxes/my-agent/resume +``` + +挂起与 exec 准入在同一行路由记录上原子判定：有命令在运行时挂起会被拒绝（`409 SANDBOX_SUSPEND_BUSY`），已挂起的沙箱会拒绝命令（`409 SANDBOX_SUSPENDED`），不会影响其他沙箱或其他 scope。恢复时优先在原 Worker 上复用目录；否则通过共享工作空间或对象存储快照迁移到其他 Worker。挂起的沙箱在 `SANDBOX_SUSPENDED_RETENTION_SECONDS`（默认 30 天）后被释放。不冻结进程，也不恢复内存。状态机、快照和磁盘回收见[挂起与恢复](docs/SUSPEND_RESUME.md)。 + ## 环境模板 每个沙箱最初的 `/envs` 都为空，依赖 NumPy、`node_modules` 或额外 Rust 工具链的应用需要重复安装。模板将这些目录构建一次后归档，再以只读方式挂载到需要它的沙箱中。 diff --git a/docs/ADAPTERS.md b/docs/ADAPTERS.md index 40d0dcd..09dcb9e 100644 --- a/docs/ADAPTERS.md +++ b/docs/ADAPTERS.md @@ -23,8 +23,29 @@ chain. Object storage holds archives rather than filesystem calls: environment templates, and checkpoint archives a caller keeps. `BlobStore.upload_checkpoint(sandbox_id, archive)` names one at `checkpoints//` and returns the URI that `get`/`download_to`/`delete` address it by, which is the whole of what this -project does about checkpoints — it does not snapshot or restore a workspace for -you. See [Architecture](ARCHITECTURE.md). +project does about caller checkpoints — it does not snapshot or restore a running +workspace for you. Suspend is the one exception: it archives a dormant workspace +under `checkpoints//dormant/` so a resume can land on another worker +([Suspend and resume](SUSPEND_RESUME.md)). A metadata store opts in to suspend by +implementing `DormantRouteStore`, and an execution backend by implementing +`DormantLifecycle`; without them the API answers 501 +`SANDBOX_SUSPEND_UNSUPPORTED`. See [Architecture](ARCHITECTURE.md). + +### S3-compatible stores + +The adapter speaks plain S3 with path-style addressing and signature v4, so any +compatible store works. Two have been run against it end to end: LocalStack and +RustFS (`rustfs/rustfs`). RustFS needs its access-key and secret-key +variables set, a data directory writable by UID 10001, and its API port +(9000) as `BLOBSTORE_ENDPOINT`; the bucket does not need to exist, because the +workers create it at startup. Against RustFS, the following behaved exactly as +against AWS S3 and needed no change: put/get/head/delete, deleting a missing key +(204), a multipart upload of a 20 MB object (ETag `-`), listing with +continuation across pages and past 1,000 keys, presigned GET and PUT, +`If-None-Match: *` (a second write answers 412), a mismatched client region, +default and `when_required` boto3 checksum modes, and a wrong secret (403). +Virtual-host addressing was not exercised, because the endpoint was an IP +address. ## Plugin discovery diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index de93e7b..d4243f1 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -32,8 +32,9 @@ flowchart LR - The execution backend creates short-lived isolated processes over a stable workspace. - The workspace is local or shared POSIX storage. The object store holds - environment templates and any checkpoint archives a caller puts there; the - service itself neither snapshots nor restores a workspace. + environment templates, any checkpoint archives a caller puts there, and the + snapshot of a suspended sandbox, which is the only workspace the service + itself archives and restores ([Suspend and resume](SUSPEND_RESUME.md)). ## Sandbox lifecycle @@ -43,7 +44,10 @@ flowchart LR 3. Every `exec` obtains a workspace lock, creates a fresh Bubblewrap process, and persists an idempotency record. 4. The route moves through `ASSIGNED`, `READY`, `RUNNING`, `RELEASING`, and - `RELEASED` states using compare-and-set updates. + `RELEASED` states using compare-and-set updates. An idle `READY` route can + also be parked as `SUSPENDING`/`SUSPENDED`, which releases its capacity slot + and keeps its workspace until a resume; see + [Suspend and resume](SUSPEND_RESUME.md). 5. Reassignment increments `generation`. Stale generations are rejected at both the control API and worker API. @@ -67,6 +71,15 @@ sandbox is pinned rather than evicted: its execution record is marked `WORKER_LOST` with it. The generation advances, so the replacement worker — and any client still holding the old generation — cannot mistake the reclaimed sandbox for the one that was running. +- **Suspended past retention.** A route suspended for longer than + `SANDBOX_SUSPENDED_RETENTION_SECONDS` (30 days by default) is released + with reason `SUSPEND_EXPIRED`, deleting its directory and snapshot. Under + disk pressure the same cycle first drops local copies of suspended sandboxes + that have a snapshot. See [Suspend and resume](SUSPEND_RESUME.md). +- **Directories a worker no longer owns.** A local workspace directory whose + route is released or names another worker (a resume landed elsewhere) is + deleted by its worker after `SANDBOX_ORPHAN_DORMANT_DIR_TTL_SECONDS` + (24 hours by default), timed by a marker on disk so a restart keeps the clock. A client that was waiting on the execution sees its connection to that worker drop. That is not recoverable at the point of failure: the guarantee is that @@ -181,7 +194,8 @@ a rollback has to stay possible. - A shared POSIX workspace has one logical writer. An advisory filesystem lock protects the failover overlap window. - Local storage failover preserves the logical sandbox ID but not physical - files. Callers needing recovery must restore from a checkpoint. + files. Callers needing recovery must restore from a checkpoint; a suspended + sandbox with a snapshot is the exception, and resumes on another worker. ## Non-goals diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index b00b4ba..56bc442 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -82,6 +82,9 @@ profile identity. | `SANDBOX_UID_START` | `20000` | First uid handed to a sandbox. | | `SANDBOX_UID_END` | `59999` | Last one. Must be at least `SANDBOX_UID_START`, which must be at least 1. | | `SANDBOX_IDLE_TTL_SECONDS` | `1800` | A sandbox with no activity for this long is reclaimed. | +| `SANDBOX_SUSPENDED_RETENTION_SECONDS` | `2592000` (30 days) | A suspended sandbox is kept this long, then released with reason `SUSPEND_EXPIRED` and its directory and snapshot deleted. `0` keeps it until a client releases it. An object-store lifecycle rule on `checkpoints/`, if you add one as a backstop, must expire later than this. Suspended sandboxes are not subject to the idle TTL. See [Suspend and resume](SUSPEND_RESUME.md). | +| `SANDBOX_SUSPEND_SNAPSHOT` | `auto` | Whether a suspend archives `/workspace`, `/home` and `/envs` to the object store so a resume can land on another worker: `auto` (local storage with a store configured), `always` (refuse to suspend without a store), or `never`. | +| `SANDBOX_ORPHAN_DORMANT_DIR_TTL_SECONDS` | `86400` (24 hours) | A local workspace directory whose route is released or names another worker (the sandbox resumed elsewhere from its snapshot, or was reassigned while this worker was away) is deleted by this worker's maintenance cycle once it has stayed orphaned this long. The clock is a marker on disk, so a restart does not reset it. A directory whose route names this worker, or that has no route, is never deleted; shared workspaces are skipped. `0` keeps such directories. | | `SANDBOX_WORKER_CAPACITY` | `32` | Sandboxes this worker admits. Placement prefers the worker with the fewest running sessions relative to its capacity, from the last heartbeat. | | `SANDBOX_HEARTBEAT_INTERVAL_SECONDS` | `10.0` | How often a worker re-registers itself. | | `SANDBOX_HEARTBEAT_TTL_SECONDS` | `30` | A registration older than this is stale. It must exceed twice the interval, or a single missed heartbeat would look like a dead worker. | diff --git a/docs/INTEGRATION_TESTING.md b/docs/INTEGRATION_TESTING.md index 406c070..3a9e62b 100644 --- a/docs/INTEGRATION_TESTING.md +++ b/docs/INTEGRATION_TESTING.md @@ -77,6 +77,12 @@ caught only by executing the interpreter. Databases use `tmpfs` for their data directories: the suite drops and recreates the schema per test, so persistence would only slow it down. +Any S3-compatible store can stand in for LocalStack through the +`SANDBOX_TEST_MINIO_ENDPOINT`, `SANDBOX_TEST_MINIO_ACCESS_KEY` and +`SANDBOX_TEST_MINIO_SECRET_KEY` variables below; the suite also passes against RustFS +(`rustfs/rustfs`, started with its access-key and secret-key variables, data +directory writable by UID 10001). + LocalStack is pinned to a community release. MinIO's Docker Hub images now require authentication to pull, and LocalStack's `stable` and `latest` tags resolve to the Pro image, which exits at startup without a license token. diff --git a/docs/ISOLATION.md b/docs/ISOLATION.md index 386202e..b74d315 100644 --- a/docs/ISOLATION.md +++ b/docs/ISOLATION.md @@ -205,6 +205,31 @@ Docker's default set, but `capsh --drop=cap_setfcap` and a tightened `CapabilityBoundingSet=` remove it. A non-root caller is unaffected — prefer one where there is the choice. +### Hosts that enforce AppArmor + +Docker on a host that enforces AppArmor (Ubuntu, including GitHub's hosted +runners) confines every container with its `docker-default` profile, which +denies the `mount` Bubblewrap performs to make `/` a slave mount inside its new +mount namespace. The startup probe then fails at every level with +`bwrap: Failed to make / slave: Permission denied` and the service exits. +Docker Desktop and OrbStack have no AppArmor, so the quick start does not hit it +and the shipped Compose files do not carry the opt-in. On an AppArmor host add +it to the trusted manager container, next to the `seccomp` setting: + +```yaml +# compose.override.yaml +services: + agent-sandbox: + security_opt: + - apparmor=unconfined +``` + +This widens the trusted manager's confinement in the same way as +`seccomp=unconfined`; a platform that can load a narrower reviewed profile that +permits those mounts should prefer it. The repository's own CI does exactly the +override above (`.github/ci/compose.apparmor.yaml`) together with the sysctl +in the previous section. + ### Docker and compatible runtimes Docker's default seccomp profile blocks the nested namespace syscalls used by diff --git a/docs/OPERATIONS.md b/docs/OPERATIONS.md index 40df92c..1630938 100644 --- a/docs/OPERATIONS.md +++ b/docs/OPERATIONS.md @@ -51,7 +51,12 @@ worker; it does not promise every in-flight command completes. Sandbox routes are sticky. Local-only workspace data is not automatically recovered onto another worker. Shared RWX POSIX storage must preserve UID -ownership; S3 checkpoints require caller-driven restore orchestration. +ownership; S3 checkpoints require caller-driven restore orchestration. The +exception is a suspended sandbox whose suspend uploaded a snapshot: resuming it +restores the workspace on another worker ([Suspend and resume](SUSPEND_RESUME.md)). Suspended +sandboxes and their snapshots are kept for `SANDBOX_SUSPENDED_RETENTION_SECONDS` +(30 days by default); an object-store lifecycle rule on `checkpoints/`, if used +as a backstop, must expire later than that. Client requests must carry the current generation so stale owners are refused. ## Backups and recovery diff --git a/docs/SUSPEND_RESUME.md b/docs/SUSPEND_RESUME.md new file mode 100644 index 0000000..26f6868 --- /dev/null +++ b/docs/SUSPEND_RESUME.md @@ -0,0 +1,249 @@ +# Suspend and resume + +An agent that starts a long task and then waits (for CI, a human review, or an +upstream job) holds a sandbox slot the whole time. Suspend gives that slot +back and keeps the workspace. Resume takes a slot again and continues with the +same files. Nothing is frozen: suspend is refused while a command runs, and a +resumed sandbox starts with no processes. It is a lifecycle layer over the +workspace, not a process checkpoint. + +```bash +# Give the slot back while waiting. Refused with 409 while anything runs. +POST /api/v1/sandboxes/my-agent/suspend {"generation":3} + +# Continue later. Idempotent; use the generation it returns from then on. +POST /api/v1/sandboxes/my-agent/resume +``` + +MCP exposes the same two operations as `sandbox_suspend(sandbox_id, generation)` +and `sandbox_resume(sandbox_id)`. Resolving a suspended sandbox through +`/api/v1/sandboxes/resolve` (or `sandbox_resolve`) also wakes it, so a client +that knows nothing about suspend keeps working. + +## States + +```text +READY ──suspend──▶ SUSPENDING ──worker released the slot──▶ SUSPENDED + ▲ │ refused (busy, snapshot failed) │ + └────────────────────┘ │ + ▲ │ + └── READY ◀── ASSIGNED ◀──────────────resume─────────────────┘ + │ worker refused or unreachable: back to SUSPENDED +``` + +Every arrow is a compare-and-set on `(sandbox_id, generation, status)` in the +metadata store. There is no new table or column: `SUSPENDING` and `SUSPENDED` +are values of the existing route status. + +| Call | In state | Result | +| --- | --- | --- | +| suspend | `READY` | `SUSPENDED`, slot released | +| suspend | `SUSPENDED` | 200 with `"suspended": false`; nothing changes | +| suspend | `RUNNING` | 409 `SANDBOX_SUSPEND_BUSY`; nothing changes | +| suspend | `ASSIGNED` (never created) | 409 `SANDBOX_NOT_READY` | +| suspend | stale generation | 409 `STALE_SANDBOX_GENERATION` | +| resume | `SUSPENDED` | `READY`, `"resumed": true` | +| resume | `ASSIGNED`/`READY`/`RUNNING` | 200 with `"resumed": false`; nothing changes | +| resume | `SUSPENDING` | 409 `SANDBOX_SUSPEND_IN_PROGRESS`; retry | +| resume | `RELEASED` (e.g. retention expired) | 409 `STALE_SANDBOX_ROUTE`; `resolve` starts a new, empty sandbox | +| exec, file call, connect | `SUSPENDING`/`SUSPENDED` | 409 `SANDBOX_SUSPENDED` | +| release | `SUSPENDED` | released, directory and snapshot deleted | + +## Concurrency with exec + +Suspend affects only the sandbox it names, and it never interrupts work in +that sandbox. + +- **Admission and suspend decide on the same row.** Exec admission locks the + route row and moves it to `RUNNING`; suspend claims it only from `READY` with + no active execution. The admission write re-checks the status, so on SQLite + (which has no `FOR UPDATE`) a suspend that committed first still wins. One of + the two wins; the other gets a 409 that says why. +- **The worker re-checks before releasing the slot.** It takes the sandbox's + exclusive lifecycle lock without waiting. Any command or file call still + inside the sandbox (in any `exec_scope`) makes the suspend fail with + `SANDBOX_SUSPEND_BUSY`, and the route goes back to `READY`. One agent thread + suspending therefore can never kill another thread's command. +- **No torn state.** The sandbox leaves the worker's live set while that lock is + held. A file call that obtained its handle earlier is refused with + `SANDBOX_SUSPENDED` when it takes its own lock, so nothing writes into a + dormant workspace or into a snapshot being taken. + +## Where a resume lands + +1. **The original worker, if it is live, on the current profile, and has a free + slot.** The directory is reused in place. The generation is unchanged, so a + client keeps using the one it had, unless the worker restarted (new epoch), + in which case the generation advances like any reassignment. +2. **Any worker, for a shared workspace** (`SANDBOX_SHARED_ROOT`). The + directory is already reachable; the generation advances. +3. **Another worker, from the snapshot**, for a local workspace whose suspend + uploaded one. The tree is rebuilt in a staging directory, handed to the + sandbox UID without following links, and renamed into place; the generation + advances, which fences the old owner. +4. Otherwise the call fails and the sandbox stays `SUSPENDED`: + `503 NO_SANDBOX_WORKER_AVAILABLE` when the original worker is live but full, + `503 SANDBOX_DORMANT_WORKSPACE_UNAVAILABLE` when the only copy is on a worker + that is not running. Both are retryable; nothing is deleted. + +A resume that fails after it claimed the route (worker refused, restore failed, +network error) puts the route back to `SUSPENDED`, so the next attempt starts +over. Attached templates come back on the same worker; a sandbox that moved has +to attach them again. + +## Snapshots + +`SANDBOX_SUSPEND_SNAPSHOT` decides whether a suspend also archives the +workspace to the object store: + +- `auto` (default): for local storage when an object store is configured. +- `always`: refuse to suspend (501 `SANDBOX_SNAPSHOT_UNAVAILABLE`) without one. +- `never`: the dormant workspace lives only on its worker. + +A snapshot covers `/workspace`, `/home` and `/envs`; `/cache` and the logs are +recreated empty. It is written under `checkpoints//dormant/` with +the same bounded, descriptor-relative archiver templates use, so the template +snapshot limits (`SANDBOX_TEMPLATE_MAX_EXTRACT_BYTES`, +`SANDBOX_TEMPLATE_MAX_ARCHIVE_BYTES`, `SANDBOX_TEMPLATE_SNAPSHOT_MAX_ENTRIES`, +`SANDBOX_TEMPLATE_SNAPSHOT_MAX_DEPTH` and +`SANDBOX_TEMPLATE_SNAPSHOT_TIMEOUT_SECONDS`) apply. A +manifest written last makes a snapshot visible only once complete, and parts +are verified against its digests before they are extracted. A workspace the +archiver refuses (a socket or FIFO, a tree past the limits) fails the suspend +with 409 `SANDBOX_SNAPSHOT_FAILED` and the sandbox stays `READY`. A shared +workspace is never snapshotted. The snapshot is deleted when the sandbox +resumes or is released. + +Sandloom deletes snapshots itself, so no bucket lifecycle rule is required. If +you add one on `checkpoints/` as a backstop against a +deletion that failed, make it expire objects **later** than +`SANDBOX_SUSPENDED_RETENTION_SECONDS` (30 days by default), with margin for the +maintenance interval and release retries; for example 35 days. A rule that +expires earlier deletes the snapshot of a sandbox that is still suspended, and +a resume that needs it (on another worker, or after disk eviction) then fails +with `SANDBOX_WORKSPACE_LOST`. Raise the rule whenever you raise the retention. + +## Retention and disk + +- **Retention.** A suspended sandbox is kept for + `SANDBOX_SUSPENDED_RETENTION_SECONDS` (default 2592000, 30 days) from the + moment it was suspended. The maintenance cycle then releases it with reason + `SUSPEND_EXPIRED`, which deletes its directory on the worker and its + snapshot. `0` keeps suspended sandboxes until a client releases them. + Suspended sandboxes are not subject to `SANDBOX_IDLE_TTL_SECONDS`, and + nothing a suspended sandbox refuses resets the clock. + + An agent that never comes back therefore does not hold anything forever: + the expiry is an ordinary release, the same path an idle timeout or a + client release takes. The route is claimed (`RELEASING`), the owning worker + deletes the workspace directory, the route is marked `RELEASED` with an + audit entry (`release_reason: SUSPEND_EXPIRED`, + `released_by: system:orphan-reaper`, visible in + `GET /api/v1/sandboxes/{id}/audit`), and the snapshot is deleted. It is + counted in the reaper's `released_total` like any other reclamation. If the + worker cannot be reached, the route stays `RELEASING` and the reaper's + normal retry (`RELEASE_RETRY`) finishes it, snapshot included; a worker + that never comes back leaves its copy to the orphaned-directory cleanup + below. After expiry, `resume` answers 409 `STALE_SANDBOX_ROUTE`, and + `resolve` starts a new, empty sandbox under the same id. +- **Disk pressure.** When the disk is over `SANDBOX_DISK_HIGH_WATERMARK_PERCENT` + or under `SANDBOX_MIN_FREE_BYTES`, the maintenance cycle deletes the local + copies of dormant sandboxes that have a snapshot, oldest first, until the + disk recovers. They stay `SUSPENDED` and resume from the snapshot. A dormant + sandbox without a snapshot is never evicted for disk; its directory is the + only copy. The set of dormant sandboxes a worker may evict is kept in its + memory; after a restart, the first cycle that finds the disk over its + watermark asks the control plane which suspended routes still name this + worker and have a snapshot, and takes those back before it evicts. +- **Copies left on the original worker.** A sandbox that resumed on another + worker from its snapshot leaves its old directory on the worker it left (as + does a sandbox released while its worker was down). Each worker's + maintenance cycle lists its local directories and checks their routes: one + whose route is `RELEASED` or names another worker is orphaned. It is + deleted once it has stayed orphaned for + `SANDBOX_ORPHAN_DORMANT_DIR_TTL_SECONDS` (default 86400, 24 hours; `0` + keeps it). The clock is a marker file under the template cache root + (`.orphan-workspaces/`), written the first time the directory is + seen orphaned, so a worker restart does not reset it; a directory that + becomes owned again loses its marker. A directory is never deleted while + its route names this worker, whether the sandbox is active, suspended, or + being created, even after a restart has emptied the worker's memory; nor is + one with no route, or a shared workspace. The route decides, not the + worker's memory: a suspended sandbox released while its worker was paused + or unreachable leaves a stale dormant entry there, and its directory is + still reclaimed. The deletion re-checks under the + sandbox's lifecycle lock, so a sandbox that came back to this worker in the + meantime keeps its directory. +- **Stalled suspends.** A route left in `SUSPENDING` (for example, the worker + was unreachable) for longer than `SANDBOX_ORPHAN_RELEASE_GRACE_SECONDS` is + finished by the next maintenance cycle. + +`/healthz` and the admin overview report the retention period under +`reclamation.suspended_retention_seconds`, and the orphaned-directory period +under `reclamation.orphan_dormant_dir_ttl_seconds`. + +## Verified combinations + +Each row was brought up as a real Docker deployment (the Sandloom image, the +named middleware, one or two workers; heartbeat 2 s, retention 90 s, orphan +directory TTL 20 s) and driven through the same scenario set over HTTP: +create and write, suspend (slot released, repeat is a no-op), `409 +SANDBOX_SUSPENDED` on exec, read, write and connect, resume (reuse, repeat, +`resolve` auto-wake), suspend against a long exec (`409 SANDBOX_SUSPEND_BUSY`, +other threads and sandboxes unaffected, plus eight exec/suspend races with +exactly one winner each), capacity (`503` while full, retry succeeds once a slot +frees), and retention expiry through the normal release (`SUSPEND_EXPIRED`, +directory and snapshot gone, `409 STALE_SANDBOX_ROUTE` afterwards). Rows with +two workers add the cross-worker case their storage allows. + +| Metadata | Registry | Object store | Storage | Workers | Also verified | +| --- | --- | --- | --- | --- | --- | +| SQLite | memory | none | local | 1 | `snapshot=false` | +| SQLite | memory | S3 (LocalStack) | local | 1 | snapshot written on suspend, deleted on resume and on expiry | +| PostgreSQL 16 | Redis 7 | S3 (LocalStack) | local | 2 | owner stopped: restore on the other worker, snapshot deleted, stale directory reclaimed after the orphan TTL while active and suspended neighbours are untouched; worker unreachable at expiry: `RELEASING`, then `RELEASE_RETRY` deletes the snapshot | +| MySQL 8.4 | Redis 7 | S3 (LocalStack) | local | 2 | same as the PostgreSQL row | +| PostgreSQL 16 | Redis 7 | none | local | 2 | owner stopped: `503 SANDBOX_DORMANT_WORKSPACE_UNAVAILABLE`, still `SUSPENDED`; owner back: resume reuses the directory | +| PostgreSQL 16 | Redis 7 | S3 (RustFS 1.0) | local | 2 | same as the LocalStack rows, plus `scripts/verify-deployment.py` (templates built on one worker and mounted on another through the same store) | +| PostgreSQL 16 | Redis 7 | S3 (RustFS 1.0) | local | 3 (one on a 128 MB disk, one more sharing it with `SANDBOX_SUSPEND_SNAPSHOT=never`) | disk pressure (below) | +| MySQL 8.4 | Redis 7 | none | shared | 2 | `snapshot=false`; owner stopped: the other worker resumes the same directory; the old worker never deletes it | + +The metadata contract (`DormantRouteStore`) also runs against real MySQL and +PostgreSQL servers, and the snapshot archive against a real S3 API, in +`tests/integration/test_dormant_adapters.py` (`./scripts/integration-test.sh +middleware`). A metadata store without `DormantRouteStore`, or an execution +backend without `DormantLifecycle`, was run as a plugin: suspend answers `501 +SANDBOX_SUSPEND_UNSUPPORTED`, the sandbox stays `READY` and keeps working. + +Disk pressure was run with the worker's workspace directory on a 128 MB ext4 +volume and `SANDBOX_DISK_HIGH_WATERMARK_PERCENT=60`. Over the watermark, only the +oldest snapshotted dormant copy was removed, and eviction stopped as soon as the +disk was back under it. With the disk pushed further over, the remaining +snapshotted copies went in order of suspension while an active sandbox and a +dormant sandbox without a snapshot were left alone, even though the disk stayed +over the watermark. Every evicted sandbox stayed `SUSPENDED`, kept its snapshot +until it resumed, and resumed with `workspace_source=restored` and an identical +file digest; one that was not evicted resumed with `reused`. A worker restart kept +the evictions on disk, and a worker that had restarted evicted again once the +disk was over the watermark. + +Not exercised against a live deployment: S3 implementations other than +LocalStack and RustFS, and object stores other than S3 (none ship with the +project). + +## Limitations + +- No process freezing and no memory restore. Background commands must have + finished, and their results stay readable through the execution record. +- Disk eviction is tracked in the worker's memory. A worker that restarts + forgets which dormant copies it evicted (their resume still restores from the + snapshot because the directory is missing) and which it may still evict (it + relearns those from the control plane, but only once the disk is over its + watermark). +- With local storage and no snapshot, a worker that never comes back takes the + dormant workspace with it, exactly as it would take a running one. Retention + eventually releases the route. +- Restored files keep their modes, except that setuid/setgid bits and + group/other write bits are dropped by the extractor. +- A resume racing another resume returns as soon as one of them has claimed the + route; the loser may see `ASSIGNED` for the moment before the winner's worker + marks it `READY`. diff --git a/pyproject.toml b/pyproject.toml index 496dcb0..a09e616 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,6 +76,9 @@ ignore_missing_imports = true asyncio_mode = "auto" testpaths = ["tests"] addopts = "-ra" +markers = [ + "real_host_probe: exercises the host-setting probes themselves, so the hermetic preflight fixture stands aside", +] [tool.uv.sources] sandloom-runtime = { path = "runtime", editable = true } diff --git a/scripts/install-polyglot-launchers.py b/scripts/install-polyglot-launchers.py index a64b860..8a7cd8e 100644 --- a/scripts/install-polyglot-launchers.py +++ b/scripts/install-polyglot-launchers.py @@ -40,6 +40,25 @@ def install_java(java_home: Path, destination: Path) -> None: ) +def _lld_can_be_opted_out(tool: Path) -> bool: + """Whether this compiler understands `-C linker-features=-lld`. + + Rust 1.90 started linking with its bundled `rust-lld` by default on + x86_64 Linux, through a `gcc-ld/ld.lld` wrapper that finds the real linker + by reading `/proc/self/exe`. In the procfs-free `basic` level there is no + `/proc`, so every link ends in `lld-wrapper: could not get the path of the + current executable`. Older compilers link with the system linker and do not + know the option, so it is only passed to a tool that accepts it. `rustdoc` + needs it as much as `rustc`: it links every doctest itself.""" + result = subprocess.run( + [str(tool), "-C", "linker-features=-lld", "--version"], + capture_output=True, + text=True, + check=False, + ) + return result.returncode == 0 + + def install_rust(sysroot: Path, destination: Path) -> None: """Call pinned native tools directly; rustup's proxies need self-discovery.""" sysroot = sysroot.resolve(strict=True) @@ -62,6 +81,10 @@ def install_rust(sysroot: Path, destination: Path) -> None: f"export RUSTDOC={shlex.quote(str(destination / 'rustdoc'))}\n" ) else: + if _lld_can_be_opted_out(binary): + # Link with the system linker, which needs no procfs. This only + # changes which linker runs; the output is the same program. + body += 'set -- -C linker-features=-lld "$@"\n' # rustc/rustdoc normally derive the sysroot from current_exe(), # which is /proc/self/exe on Linux. Explicit sysroots still work. body += ( diff --git a/scripts/integration-test.sh b/scripts/integration-test.sh index a04c6d2..528a587 100755 --- a/scripts/integration-test.sh +++ b/scripts/integration-test.sh @@ -104,6 +104,7 @@ run_polyglot_tests() { echo "==> verifying every language at basic without procfs" docker run --rm \ --security-opt seccomp=unconfined \ + --security-opt apparmor=unconfined \ -v "$(pwd)/scripts:/verification:ro" \ --entrypoint /app/.venv/bin/python \ "${POLYGLOT_IMAGE}" /verification/verify-polyglot-runtime.py --levels basic @@ -114,6 +115,7 @@ run_polyglot_tests() { docker rm -f agent-sandbox-polyglot >/dev/null 2>&1 || true docker run -d --name agent-sandbox-polyglot \ --security-opt seccomp=unconfined \ + --security-opt apparmor=unconfined \ --security-opt systempaths=unconfined \ -e SANDBOX_INTERNAL_TOKEN="${POLYGLOT_TOKEN}" \ -e SANDBOX_ADVERTISE_HOST=localhost \ diff --git a/scripts/verify-deployment.py b/scripts/verify-deployment.py index 53bc73d..935c6b1 100755 --- a/scripts/verify-deployment.py +++ b/scripts/verify-deployment.py @@ -308,6 +308,8 @@ def shell_output(self, script: str, *, login: bool = False) -> str: "sandbox_write_file", "sandbox_read_file", "sandbox_release", + "sandbox_suspend", + "sandbox_resume", } diff --git a/src/agent_sandbox/app.py b/src/agent_sandbox/app.py index db984a3..57c0c4f 100644 --- a/src/agent_sandbox/app.py +++ b/src/agent_sandbox/app.py @@ -44,8 +44,12 @@ MakeDirectoryRequest, MovePathRequest, ResolveRequest, + ResumeLocalRequest, + ResumeResponse, RouteResponse, SandboxAuditResponse, + SuspendResponse, + SuspendSandboxRequest, TemplateAttachRequest, TemplateBuildRequest, TemplateListResponse, @@ -80,6 +84,7 @@ async def require_internal_auth( ExecRequest, FileWriteRequest, ResolveRequest, + SuspendSandboxRequest, ) ) ) @@ -161,6 +166,16 @@ def relay_worker_response(response: httpx.Response) -> Response: "SANDBOX_EXEC_SCOPE_BUSY": 409, "SANDBOX_EXEC_SCOPE_LOCKED": 409, "SANDBOX_FILE_PATH_LOCKED": 409, + # Suspend and resume. A suspended sandbox refuses work until it is resumed; + # a suspend refuses while anything runs in it. Both are the caller's move. + "SANDBOX_NOT_FOUND": 404, + "SANDBOX_SUSPENDED": 409, + "SANDBOX_SUSPEND_BUSY": 409, + "SANDBOX_SUSPEND_IN_PROGRESS": 409, + "SANDBOX_NOT_READY": 409, + "SANDBOX_SNAPSHOT_FAILED": 409, + "SANDBOX_SUSPEND_UNSUPPORTED": 501, + "SANDBOX_SNAPSHOT_UNAVAILABLE": 501, # The deployment's object store, not the request: templates cannot be shared # until an operator fixes the bucket, and a retry after that is the same call. "OBJECT_STORE_UNAVAILABLE": 503, @@ -185,6 +200,8 @@ def reclamation_periods(config: Settings) -> dict[str, float]: "idle_ttl_seconds": config.idle_ttl_seconds, "orphan_running_grace_seconds": config.orphan_running_grace_seconds, "orphan_release_grace_seconds": config.orphan_release_grace_seconds, + "suspended_retention_seconds": config.suspended_retention_seconds, + "orphan_dormant_dir_ttl_seconds": config.orphan_dormant_dir_ttl_seconds, } @@ -528,6 +545,25 @@ async def delete_template(name: str, _: InternalAuth) -> dict[str, str]: removed = service.remove_template(name) return {"status": "REMOVED" if removed else "NOT_FOUND", "name": name} + @app.post("/api/v1/sandboxes/{sandbox_id}/suspend", response_model=SuspendResponse) + async def suspend_sandbox( + sandbox_id: str, request: SuspendSandboxRequest, _: InternalAuth + ) -> SuspendResponse: + """Release the sandbox's capacity slot and keep its workspace. + + Refused with 409 `SANDBOX_SUSPEND_BUSY` while any command or file call + is running. Repeating it on a suspended sandbox changes nothing. + """ + return await service.suspend(sandbox_id, request.generation) + + @app.post("/api/v1/sandboxes/{sandbox_id}/resume", response_model=ResumeResponse) + async def resume_sandbox(sandbox_id: str, _: InternalAuth) -> ResumeResponse: + """Take a slot again and return the route to use. A no-op when awake. + + The generation changes only if the sandbox moved to another worker. + """ + return await service.resume(sandbox_id) + @app.delete("/api/v1/sandboxes/{sandbox_id}") async def release(sandbox_id: str, _: InternalAuth) -> dict[str, str]: route = await database.find_route(sandbox_id) @@ -919,6 +955,64 @@ async def attach_templates( attached = await complete_in_thread(service.attach_templates, sandbox_id, request) return {"status": "ATTACHED", "templates": attached} + @app.post("/internal/v1/sandboxes/{sandbox_id}/suspend") + async def suspend_local( + sandbox_id: str, + request: SuspendSandboxRequest, + _: InternalAuth, + worker_id: Annotated[str | None, Header(alias="X-Sandbox-Worker-ID")] = None, + ) -> dict[str, object]: + if worker_id != service.worker_id: + raise HTTPException(status_code=409, detail="STALE_SANDBOX_ROUTE") + route = await service.validate_local_route( + sandbox_id, request.generation, allow_dormant=True + ) + if route.status not in {"SUSPENDING", "SUSPENDED"}: + # Only a claimed route is suspended; the claim is what fences exec. + raise HTTPException(status_code=409, detail="STALE_SANDBOX_ROUTE") + result = await service.dormant_runtime().suspend( + sandbox_id, request.generation, snapshot=service.snapshot_on_suspend(route) + ) + await service.refresh_capacity() + return {"sandbox_id": sandbox_id, "generation": request.generation, **result} + + @app.post("/internal/v1/sandboxes/{sandbox_id}/resume") + async def resume_local( + sandbox_id: str, + request: ResumeLocalRequest, + _: InternalAuth, + worker_id: Annotated[str | None, Header(alias="X-Sandbox-Worker-ID")] = None, + ) -> dict[str, object]: + if worker_id != service.worker_id or request.worker_epoch != service.worker_epoch: + raise HTTPException(status_code=409, detail="STALE_SANDBOX_ROUTE") + route = await service.validate_local_route(sandbox_id, request.generation) + if route.sandbox_uid != request.sandbox_uid or route.profile_id != request.profile: + raise HTTPException(status_code=409, detail="STALE_SANDBOX_ROUTE") + runtime = service.dormant_runtime() + if ( + sandbox_id not in service.runtime.sandboxes + and len(service.runtime.sandboxes) >= config.worker_capacity + ): + # Resuming takes a slot like creating does; this worker has none. + raise RuntimeError("NO_SANDBOX_WORKER_AVAILABLE") + sandbox, source = await runtime.resume( + sandbox_id, request.generation, request.sandbox_uid, restore=request.restore + ) + # Before READY, so a suspend that follows cannot have its fresh + # snapshot deleted by this one's cleanup. + await service.delete_dormant_snapshot(sandbox_id) + await database.mark_route_ready( + sandbox_id, request.generation, service.worker_id, service.worker_epoch + ) + await service.refresh_capacity() + return { + "sandbox_id": sandbox_id, + "generation": sandbox.generation, + "uid": sandbox.uid, + "status": "READY", + "workspace_source": source, + } + @app.delete("/internal/v1/sandboxes/{sandbox_id}") async def destroy_local( sandbox_id: str, diff --git a/src/agent_sandbox/backends.py b/src/agent_sandbox/backends.py index 21bbccf..c3804ee 100644 --- a/src/agent_sandbox/backends.py +++ b/src/agent_sandbox/backends.py @@ -85,6 +85,88 @@ def as_directory_operations(backend: object) -> DirectoryOperations | None: return cast("DirectoryOperations", backend) +class DormantLifecycle(Protocol): + """Suspending a sandbox to release its slot, and resuming it later. + + Optional for the same reason `DirectoryOperations` is: a backend from an + earlier release keeps loading, and the API answers 501 for it instead. + Narrow with `as_dormant_lifecycle()`. + """ + + dormant: dict[str, Any] + + async def suspend( + self, sandbox_id: str, generation: int, *, snapshot: bool = False + ) -> dict[str, object]: ... + async def resume( + self, sandbox_id: str, generation: int, uid: int, *, restore: str = "reuse" + ) -> tuple[LocalSandbox, str]: ... + async def reclaim_dormant_disk(self) -> list[str]: ... + + +def as_dormant_lifecycle(backend: object) -> DormantLifecycle | None: + """Narrow a backend to suspend/resume, or report that it cannot.""" + if backend is None: + return None + required = ("suspend", "resume", "reclaim_dormant_disk") + if not all(callable(getattr(backend, name, None)) for name in required): + return None + return cast("DormantLifecycle", backend) + + +class OrphanWorkspaceReclaim(Protocol): + """Deleting workspace directories this worker no longer owns. + + The worker lists its local directories; the control plane decides which of + them are orphaned from the metadata store; the worker deletes those that + have stayed orphaned for the configured period. Optional, like the other + capabilities here. Narrow with `as_orphan_workspace_reclaim()`. + """ + + async def local_workspace_ids(self) -> list[str]: ... + async def reclaim_orphan_workspaces( + self, orphaned: list[str], *, ttl_seconds: int + ) -> list[str]: ... + + +def as_orphan_workspace_reclaim(backend: object) -> OrphanWorkspaceReclaim | None: + """Narrow a backend to orphaned-directory reclamation, or report that it cannot.""" + if backend is None: + return None + required = ("local_workspace_ids", "reclaim_orphan_workspaces") + if not all(callable(getattr(backend, name, None)) for name in required): + return None + return cast("OrphanWorkspaceReclaim", backend) + + +class DormantAdoption(Protocol): + """Taking back a dormant sandbox this process no longer remembers. + + Which dormant directories may be evicted under disk pressure is kept in the + worker's memory, so a restart forgets them and they would never be evicted. + The control plane still knows which suspended routes name this worker; it + hands those back through `adopt_dormant`. Optional. Narrow with + `as_dormant_adoption()`. + """ + + async def adopt_dormant( + self, + sandbox_id: str, + *, + generation: int, + uid: int, + suspended_at: float, + snapshot: bool, + ) -> bool: ... + + +def as_dormant_adoption(backend: object) -> DormantAdoption | None: + """Narrow a backend to dormant adoption, or report that it cannot.""" + if backend is None or not callable(getattr(backend, "adopt_dormant", None)): + return None + return cast("DormantAdoption", backend) + + def create_execution_backend(settings: Settings) -> ExecutionBackend: """Load the explicitly configured backend. @@ -103,9 +185,11 @@ def create_execution_backend(settings: Settings) -> ExecutionBackend: __all__ = [ "DirectoryOperations", + "DormantLifecycle", "ExecutionBackend", "TemplateCachePruning", "as_directory_operations", + "as_dormant_lifecycle", "as_template_cache_pruning", "create_execution_backend", ] diff --git a/src/agent_sandbox/config.py b/src/agent_sandbox/config.py index ef194e6..3d5337c 100644 --- a/src/agent_sandbox/config.py +++ b/src/agent_sandbox/config.py @@ -75,6 +75,21 @@ class Settings(BaseSettings): orphan_running_grace_seconds: int = 60 orphan_reaper_batch_size: int = 100 maintenance_interval_seconds: float = 60.0 + # How long a suspended (dormant) sandbox keeps its workspace before the + # maintenance reaper releases it and reclaims the disk. 0 keeps it until a + # client releases it. + suspended_retention_seconds: int = Field(default=30 * 24 * 3600, ge=0) + # Whether a suspend also archives the workspace to the object store, so a + # resume can land on another worker: `auto` does so for local storage when a + # store is configured, `always` refuses to suspend without one, `never` keeps + # the dormant workspace on its worker only. + suspend_snapshot: str = "auto" + # How long a workspace directory this worker no longer owns (its sandbox + # resumed on another worker, or was released while this worker was away) + # stays on disk before the maintenance cycle deletes it. Counted from when + # the worker first saw it orphaned, by a marker on disk, so a restart does + # not reset it. 0 never deletes such directories. + orphan_dormant_dir_ttl_seconds: int = Field(default=24 * 3600, ge=0) disk_high_watermark_percent: int = 90 min_free_bytes: int = 1024 * 1024 * 1024 @@ -252,6 +267,14 @@ def validate_network_mode(cls, value: str) -> str: raise ValueError("network_mode must be 'host' or 'isolated'") return normalized + @field_validator("suspend_snapshot") + @classmethod + def validate_suspend_snapshot(cls, value: str) -> str: + normalized = value.strip().lower() + if normalized not in {"auto", "always", "never"}: + raise ValueError("suspend_snapshot must be 'auto', 'always' or 'never'") + return normalized + @field_validator("isolation_required_features", "isolation_optional_features") @classmethod def validate_isolation_features(cls, value: list[str]) -> list[str]: diff --git a/src/agent_sandbox/console.py b/src/agent_sandbox/console.py index cf5818a..2f38644 100644 --- a/src/agent_sandbox/console.py +++ b/src/agent_sandbox/console.py @@ -336,7 +336,7 @@ }; const STATUS_TONE = { READY: "ok", ACTIVE: "ok", RUNNING: "run", ASSIGNED: "warn", - RELEASING: "warn", RELEASED: "off", SUCCEEDED: "ok", + RELEASING: "warn", RELEASED: "off", SUSPENDING: "warn", SUSPENDED: "off", SUCCEEDED: "ok", FAILED: "bad", TIMED_OUT: "bad", CANCELLED: "off", LOST: "bad", DRAINING: "warn", }; const status = (value) => el("span", { class: "st " + (STATUS_TONE[value] || "") }, value || "—"); @@ -371,7 +371,7 @@ const disk = d.disk || {}; const rec = d.reclamation || {}; const byStatus = d.sandboxes_by_status || {}; - const order = ["READY", "RUNNING", "ASSIGNED", "RELEASING", "RELEASED"]; + const order = ["READY", "RUNNING", "ASSIGNED", "SUSPENDED", "RELEASING", "RELEASED"]; const seen = Object.keys(byStatus); const statuses = order.filter((k) => k in byStatus).concat(seen.filter((k) => !order.includes(k))); return [ @@ -481,7 +481,7 @@ }), el("select", { onchange: (e) => { S.filters.status = e.target.value; S.page.offset = 0; load(); }, - }, ["", "READY", "RUNNING", "ASSIGNED", "RELEASING", "RELEASED"].map((v) => + }, ["", "READY", "RUNNING", "ASSIGNED", "SUSPENDED", "RELEASING", "RELEASED"].map((v) => el("option", { value: v, selected: S.filters.status === v }, v || "All statuses"))), ), rows.length diff --git a/src/agent_sandbox/dormant.py b/src/agent_sandbox/dormant.py new file mode 100644 index 0000000..306a894 --- /dev/null +++ b/src/agent_sandbox/dormant.py @@ -0,0 +1,233 @@ +"""Dormant workspace snapshots: what lets a suspended sandbox wake elsewhere. + +A suspended sandbox keeps its directory on the worker that hosted it, and that +is enough to resume it there. With local storage, resuming anywhere else needs +the bytes, so a suspend can also archive the directories that hold state — +`/workspace`, `/home` and `/envs` — to the object store. `/cache` and the logs +are throwaway and are recreated empty. + +Layout, under the store's base prefix: + + checkpoints//dormant/manifest.json + checkpoints//dormant//.tar.gz + +The manifest is written last and names the nonce, so a snapshot is visible only +once every part is uploaded, and a failed snapshot never replaces a good one +with a half-written one. Parts are verified against the manifest's digests +before they are extracted. + +This is a filesystem snapshot, not a process checkpoint: nothing running is +captured, and nothing is running when it is taken. +""" + +from __future__ import annotations + +import contextlib +import json +import logging +import os +import shutil +import tempfile +import time +import uuid +from pathlib import Path +from typing import Any + +from .archive import ArchiveError, ArchiveLimits, snapshot_directory +from .templates import TemplateError, TemplateObjectStore, digest_of_file, extract_archive +from .workspace import open_directory + +logger = logging.getLogger(__name__) + +# The sandbox directories that carry state across a suspend. +SNAPSHOT_PARTS = ("workspace", "home", "envs") +MANIFEST_VERSION = 1 + + +def snapshot_prefix(sandbox_id: str) -> str: + return f"checkpoints/{sandbox_id}/dormant" + + +def manifest_key(sandbox_id: str) -> str: + return f"{snapshot_prefix(sandbox_id)}/manifest.json" + + +def part_key(sandbox_id: str, nonce: str, part: str) -> str: + return f"{snapshot_prefix(sandbox_id)}/{nonce}/{part}.tar.gz" + + +def read_manifest(store: TemplateObjectStore, sandbox_id: str) -> dict[str, Any] | None: + """The current snapshot's manifest, or None when there is no snapshot.""" + + with tempfile.TemporaryDirectory() as scratch: + destination = Path(scratch) / "manifest.json" + try: + store.download_to(manifest_key(sandbox_id), destination) + except FileNotFoundError: + return None + try: + data = json.loads(destination.read_text(encoding="utf-8")) + except (OSError, ValueError): + logger.warning("dormant snapshot manifest is unreadable: sandbox_id=%s", sandbox_id) + return None + if not isinstance(data, dict) or data.get("sandbox_id") != sandbox_id: + return None + if not isinstance(data.get("nonce"), str) or not isinstance(data.get("parts"), dict): + return None + return data + + +def write_snapshot( + store: TemplateObjectStore, + *, + sandbox_id: str, + generation: int, + root: Path, + staging_root: Path, + limits: ArchiveLimits, +) -> dict[str, Any]: + """Archive the stateful directories of `root` and publish a new manifest. + + The caller fences every cooperating writer first; see `snapshot_directory` + for why this is not an atomic filesystem snapshot. + """ + + previous = read_manifest(store, sandbox_id) + nonce = uuid.uuid4().hex + staging_root.mkdir(mode=0o700, parents=True, exist_ok=True) + scratch = Path(tempfile.mkdtemp(prefix="dormant-", dir=staging_root)) + parts: dict[str, dict[str, Any]] = {} + try: + for part in SNAPSHOT_PARTS: + if not (root / part).is_dir(): + continue + archive = scratch / f"{part}.tar.gz" + try: + with open_directory(root, (part,)) as directory: + digest, size = snapshot_directory(directory, archive, limits) + except (ArchiveError, OSError, ValueError) as exc: + logger.warning( + "dormant snapshot failed: sandbox_id=%s part=%s error=%s", + sandbox_id, + part, + exc, + ) + raise RuntimeError("SANDBOX_SNAPSHOT_FAILED") from exc + store.upload_file(part_key(sandbox_id, nonce, part), archive) + parts[part] = {"digest": digest, "size_bytes": size} + archive.unlink(missing_ok=True) + manifest = { + "version": MANIFEST_VERSION, + "sandbox_id": sandbox_id, + "generation": generation, + "nonce": nonce, + "created_at": time.time(), + "parts": parts, + } + document = scratch / "manifest.json" + document.write_text(json.dumps(manifest, sort_keys=True), encoding="utf-8") + store.upload_file(manifest_key(sandbox_id), document) + except BaseException: + # Parts of a manifest nobody will ever publish. + for part in parts: + with contextlib.suppress(Exception): + store.delete(part_key(sandbox_id, nonce, part)) + raise + finally: + shutil.rmtree(scratch, ignore_errors=True) + if previous is not None: + _delete_parts(store, sandbox_id, previous) + return manifest + + +def restore_snapshot( + store: TemplateObjectStore, + *, + sandbox_id: str, + destination: Path, + uid: int, + max_extract_bytes: int | None, +) -> dict[str, Any] | None: + """Rebuild a sandbox tree at `destination` from the current snapshot. + + Returns the manifest, or None when there is no snapshot. `destination` must + not exist; the caller renames it into place once this returns, so a failed + restore never leaves a half-populated sandbox root behind. + """ + + manifest = read_manifest(store, sandbox_id) + if manifest is None: + return None + destination.mkdir(mode=0o700, parents=True) + with tempfile.TemporaryDirectory(dir=destination.parent) as scratch: + for part, details in sorted(manifest["parts"].items()): + if part not in SNAPSHOT_PARTS or not isinstance(details, dict): + raise RuntimeError("SANDBOX_SNAPSHOT_FAILED") + archive = Path(scratch) / f"{part}.tar.gz" + store.download_to(part_key(sandbox_id, manifest["nonce"], part), archive) + if digest_of_file(archive) != details.get("digest"): + logger.warning( + "dormant snapshot digest mismatch: sandbox_id=%s part=%s", sandbox_id, part + ) + raise RuntimeError("SANDBOX_SNAPSHOT_FAILED") + try: + # The template extractor already refuses traversal, escaping + # links, and device nodes, and strips setuid bits. + extract_archive(archive, destination / part, max_bytes=max_extract_bytes) + except TemplateError as exc: + raise RuntimeError("SANDBOX_SNAPSHOT_FAILED") from exc + archive.unlink(missing_ok=True) + for name in ("workspace", "home", "cache", "envs", "logs"): + (destination / name).mkdir(exist_ok=True) + chown_tree(destination, uid) + return manifest + + +def chown_tree(root: Path, uid: int) -> None: + """Hand a restored tree to the sandbox UID without following any link. + + Extraction runs as the service user, and a symlink in agent-controlled + content may point anywhere on the host: `lchown` changes the link itself. + """ + + os.chown(root, uid, uid, follow_symlinks=False) + for current, directories, files in os.walk(root, followlinks=False): + for name in (*directories, *files): + os.chown(os.path.join(current, name), uid, uid, follow_symlinks=False) + + +def delete_snapshot(store: TemplateObjectStore, sandbox_id: str) -> bool: + """Remove the current snapshot, manifest first. Returns whether one existed.""" + + manifest = read_manifest(store, sandbox_id) + if manifest is None: + return False + store.delete(manifest_key(sandbox_id)) + _delete_parts(store, sandbox_id, manifest) + return True + + +def _delete_parts(store: TemplateObjectStore, sandbox_id: str, manifest: dict[str, Any]) -> None: + for part in manifest.get("parts", {}): + try: + store.delete(part_key(sandbox_id, str(manifest["nonce"]), str(part))) + except Exception: + logger.warning( + "could not delete a superseded dormant snapshot part: sandbox_id=%s part=%s", + sandbox_id, + part, + exc_info=True, + ) + + +__all__ = [ + "SNAPSHOT_PARTS", + "chown_tree", + "delete_snapshot", + "manifest_key", + "part_key", + "read_manifest", + "restore_snapshot", + "snapshot_prefix", + "write_snapshot", +] diff --git a/src/agent_sandbox/mcp_server.py b/src/agent_sandbox/mcp_server.py index 1aad2d2..934dffe 100644 --- a/src/agent_sandbox/mcp_server.py +++ b/src/agent_sandbox/mcp_server.py @@ -16,7 +16,14 @@ from .config import Settings from .models import Route -from .schemas import ExecRequest, ExecResponse, ResolveRequest, SandboxAuditResponse +from .schemas import ( + ExecRequest, + ExecResponse, + ResolveRequest, + ResumeResponse, + SandboxAuditResponse, + SuspendResponse, +) from .service import SandboxService, route_audit from .storage import MetadataStore @@ -128,7 +135,9 @@ def build_mcp_server( "agent-sandbox", instructions=( "Operate Sandboxes through the unified Control Plane. Always keep the generation " - "returned by sandbox_resolve and call sandbox_release in a finally step." + "returned by sandbox_resolve and call sandbox_release in a finally step. " + "While waiting on something slow, call sandbox_suspend to give the slot back " + "and sandbox_resume (or sandbox_resolve) to continue with the same workspace." ), ) @@ -360,6 +369,29 @@ async def sandbox_read_file( ) return FileContent.model_validate(_worker_payload(response)) + @mcp.tool( + name="sandbox_suspend", + description=( + "Release the Sandbox capacity slot and keep its workspace until resumed. " + "Refused while a command runs; repeating it is a no-op. Running processes " + "are not preserved." + ), + annotations=IDEMPOTENT_WRITE, + ) + async def sandbox_suspend(sandbox_id: str, generation: int) -> SuspendResponse: + return await service.suspend(sandbox_id, generation) + + @mcp.tool( + name="sandbox_resume", + description=( + "Wake a suspended Sandbox with its workspace and return the route to use; " + "a no-op when it is awake. Use the returned generation from then on." + ), + annotations=IDEMPOTENT_WRITE, + ) + async def sandbox_resume(sandbox_id: str) -> ResumeResponse: + return await service.resume(sandbox_id) + @mcp.tool( name="sandbox_release", description="Release the current Sandbox route and destroy its runtime.", diff --git a/src/agent_sandbox/preflight.py b/src/agent_sandbox/preflight.py index f21f6bb..2370fad 100644 --- a/src/agent_sandbox/preflight.py +++ b/src/agent_sandbox/preflight.py @@ -140,6 +140,9 @@ def validate_runtime_environment(settings: Settings) -> list[str]: return warnings +_APPARMOR_USERNS_SYSCTL = Path("/proc/sys/kernel/apparmor_restrict_unprivileged_userns") + + def check_user_namespace_support() -> list[str]: """Report host settings that break Bubblewrap in ways that are hard to read. @@ -164,9 +167,8 @@ def check_user_namespace_support() -> list[str]: return [] warnings: list[str] = [] - restrict = Path("/proc/sys/kernel/apparmor_restrict_unprivileged_userns") try: - if restrict.read_text(encoding="utf-8").strip() == "1": + if _APPARMOR_USERNS_SYSCTL.read_text(encoding="utf-8").strip() == "1": warnings.append( "kernel.apparmor_restrict_unprivileged_userns=1 strips capabilities from new " "user namespaces; Bubblewrap probes may fail with 'Operation not permitted'" diff --git a/src/agent_sandbox/runtime.py b/src/agent_sandbox/runtime.py index 9f73b1b..aeb4dc3 100644 --- a/src/agent_sandbox/runtime.py +++ b/src/agent_sandbox/runtime.py @@ -40,6 +40,7 @@ from .archive import ArchiveLimits from .config import Settings +from .dormant import restore_snapshot, write_snapshot from .isolation_policy import negotiate_features from .schemas import ExecRequest, ExecResponse from .templates import ( @@ -80,6 +81,24 @@ def workspace(self) -> Path: return self.root / "workspace" +@dataclass(slots=True) +class DormantSandbox: + """A suspended sandbox: its directory is kept, its capacity slot is not. + + Only what a resume on this worker needs. `snapshot` records whether the + object store holds a copy, which is what makes the local directory safe to + evict under disk pressure. + """ + + sandbox_id: str + generation: int + uid: int + suspended_at: float + snapshot: bool + templates: tuple[TemplateRecord, ...] = () + evicted: bool = False + + @dataclass(frozen=True, slots=True) class RuntimeExtension: """Generic customization supplied by an optional execution plugin. @@ -236,6 +255,8 @@ def __init__( self._attached_templates: dict[str, tuple[TemplateRecord, ...]] = {} self._destroying: set[str] = set() self._destroy_tasks: dict[str, asyncio.Task[None]] = {} + # Suspended sandboxes. Not in `sandboxes`, so they hold no capacity slot. + self.dormant: dict[str, DormantSandbox] = {} async def probe(self) -> dict[str, object]: runtime_report = await asyncio.to_thread(self.builder.probe) @@ -521,6 +542,10 @@ def _path_lock( stream.fileno(), (fcntl.LOCK_SH if shared else fcntl.LOCK_EX) | fcntl.LOCK_NB, ) + # A suspend unregisters the sandbox while it holds the global lock + # exclusively, so a call that obtained its handle earlier finds out + # here instead of writing into a dormant workspace. + self._check_not_dormant(sandbox) yield except BlockingIOError as exc: raise RuntimeError("SANDBOX_FILE_PATH_LOCKED") from exc @@ -701,6 +726,13 @@ def build_template( mount_target=f"{TEMPLATE_MOUNT_ROOT}/{name}", ) + def _check_not_dormant(self, sandbox: LocalSandbox) -> None: + if sandbox.sandbox_id in self.dormant: + raise RuntimeError("SANDBOX_SUSPENDED") + current = self.sandboxes.get(sandbox.sandbox_id) + if current is not None and current.generation != sandbox.generation: + raise RuntimeError("STALE_SANDBOX_GENERATION") + def _check_not_destroying(self, sandbox_id: str) -> None: if sandbox_id in self._destroying: raise RuntimeError("SANDBOX_RELEASING") @@ -764,6 +796,7 @@ def _destroy_files(self, sandbox_id: str, delete_files: bool) -> None: def _remove_sandbox(self, sandbox_id: str, *, delete_files: bool) -> None: sandbox = self.sandboxes.pop(sandbox_id, None) self.last_active_at.pop(sandbox_id, None) + self.dormant.pop(sandbox_id, None) # Unpin the revisions this sandbox held. Leaving them here would keep # every revision ever attached to this worker in the cache for as long # as the process lives, because pruning never evicts a pinned revision. @@ -784,6 +817,348 @@ def _remove_sandbox(self, sandbox_id: str, *, delete_files: bool) -> None: # the same id is not removed later by the background task. shutil.rmtree(root, True) + # ── suspend and resume ── + # + # A suspended sandbox keeps its directory and gives up its slot: it leaves + # `sandboxes`, which is what the heartbeat counts as running sessions. No + # process is frozen; suspend refuses while anything is running. + + async def suspend( + self, sandbox_id: str, generation: int, *, snapshot: bool = False + ) -> dict[str, object]: + """Release the capacity slot of an idle sandbox, keeping its files. + + Idempotent: suspending a sandbox that is already dormant at this + generation reports success and changes nothing. + """ + return await complete_in_thread(self._suspend_sandbox, sandbox_id, generation, snapshot) + + def _suspend_sandbox( + self, sandbox_id: str, generation: int, snapshot: bool + ) -> dict[str, object]: + self._check_not_destroying(sandbox_id) + with self._template_lock(sandbox_id): + self._check_not_destroying(sandbox_id) + sandbox = self.sandboxes.get(sandbox_id) + if sandbox is None: + dormant = self.dormant.get(sandbox_id) + if dormant is not None and dormant.generation == generation: + return {"status": "SUSPENDED", "snapshot": dormant.snapshot} + raise RuntimeError("STALE_SANDBOX_GENERATION") + if sandbox.generation != generation: + raise RuntimeError("STALE_SANDBOX_GENERATION") + if self._has_processes(sandbox_id): + raise RuntimeError("SANDBOX_SUSPEND_BUSY") + try: + with self._execution_lock(sandbox, None): + # Exclusive and nonblocking: any command or file call still + # inside this sandbox makes the suspend fail cleanly, and + # none can start until the sandbox is unregistered below. + if self._has_processes(sandbox_id): + raise RuntimeError("SANDBOX_SUSPEND_BUSY") + store = self.templates.object_store + if snapshot: + if store is None: + raise RuntimeError("SANDBOX_SNAPSHOT_UNAVAILABLE") + write_snapshot( + store, + sandbox_id=sandbox_id, + generation=generation, + root=sandbox.root, + staging_root=self._staging_root(), + limits=self.templates.snapshot_limits, + ) + self.sandboxes.pop(sandbox_id, None) + self.last_active_at.pop(sandbox_id, None) + self.dormant[sandbox_id] = DormantSandbox( + sandbox_id=sandbox_id, + generation=generation, + uid=sandbox.uid, + suspended_at=time.time(), + snapshot=snapshot, + # Kept, but not pinned: a dormant sandbox has nothing + # mounted, and a resume materializes them again. + templates=self._attached_templates.pop(sandbox_id, ()), + ) + except RuntimeError as exc: + if str(exc) in {"SANDBOX_SHARED_WORKSPACE_LOCKED", "SANDBOX_EXEC_SCOPE_LOCKED"}: + raise RuntimeError("SANDBOX_SUSPEND_BUSY") from exc + raise + return {"status": "SUSPENDED", "snapshot": snapshot} + + async def resume( + self, sandbox_id: str, generation: int, uid: int, *, restore: str = "reuse" + ) -> tuple[LocalSandbox, str]: + """Register a dormant sandbox again and return it with where its files came from. + + `restore="reuse"` keeps the local directory and falls back to the + snapshot only when the directory is gone; `restore="snapshot"` replaces + whatever is local with the snapshot, which is what a sandbox moving to a + different worker needs. Idempotent for an already registered sandbox. + """ + if restore not in {"reuse", "snapshot"}: + raise ValueError("restore must be 'reuse' or 'snapshot'") + return await complete_in_thread(self._resume_sandbox, sandbox_id, generation, uid, restore) + + def _resume_sandbox( + self, sandbox_id: str, generation: int, uid: int, restore: str + ) -> tuple[LocalSandbox, str]: + self._check_not_destroying(sandbox_id) + with self._template_lock(sandbox_id): + self._check_not_destroying(sandbox_id) + existing = self.sandboxes.get(sandbox_id) + if existing and existing.generation == generation and existing.uid == uid: + self.touch(sandbox_id) + return existing, "active" + root = _sandbox_root(self.settings.workspace_root, sandbox_id) + source = "reused" + if restore == "snapshot" or not root.exists(): + self._restore_locked(sandbox_id, uid, root) + source = "restored" + dormant = self.dormant.pop(sandbox_id, None) + self.sandboxes.pop(sandbox_id, None) + sandbox = self._create_sandbox_locked(sandbox_id, generation, uid) + if dormant is not None and dormant.templates: + self._attached_templates[sandbox_id] = dormant.templates + return sandbox, source + + def _restore_locked(self, sandbox_id: str, uid: int, root: Path) -> None: + store = self.templates.object_store + if store is None: + raise RuntimeError("SANDBOX_WORKSPACE_LOST") + if not self.disk_status()["available"]: + raise RuntimeError("SANDBOX_WORKER_DISK_PRESSURE") + staging = self._staging_root() / f"restore-{uuid.uuid4().hex}" + try: + manifest = restore_snapshot( + store, + sandbox_id=sandbox_id, + destination=staging, + uid=uid, + max_extract_bytes=self.settings.template_max_extract_bytes, + ) + if manifest is None: + raise RuntimeError("SANDBOX_WORKSPACE_LOST") + if root.exists(): + # A stale local copy loses to the snapshot it predates. + self._move_to_trash(sandbox_id, root) + os.replace(staging, root) + finally: + shutil.rmtree(staging, ignore_errors=True) + + def _staging_root(self) -> Path: + # Beside the sandboxes, so the final rename stays on one filesystem, and + # under a reserved name no sandbox id can take. + path = self.settings.workspace_root / ".staging" + path.mkdir(mode=0o700, parents=True, exist_ok=True) + return path + + def _move_to_trash(self, sandbox_id: str, root: Path) -> Path | None: + trash_root = self.settings.workspace_root / ".trash" + trash_root.mkdir(mode=0o700, parents=True, exist_ok=True) + target = trash_root / f"{sandbox_id}-{uuid.uuid4().hex}" + try: + os.replace(root, target) + except OSError: + shutil.rmtree(root, True) + return None + return target + + # ── orphaned workspace directories ── + # + # A sandbox that resumed on another worker from its snapshot, or that was + # released while this worker was down, leaves its directory here. Only the + # control plane can tell that from the metadata store, so it passes in the + # ids it found orphaned; this side keeps the clock and does the deletion. + # The clock is a marker file per directory, created the first time the + # directory is reported orphaned, so a restarted worker keeps counting + # instead of starting over, and a directory that becomes owned again (or + # disappears) loses its marker. + + def _orphan_marker_root(self) -> Path: + # Beside the lifecycle locks, outside every sandbox mount, and skipped + # by the template cache because of the leading dot. + return self.settings.template_cache_root / ".orphan-workspaces" + + def _owned_here(self, sandbox_id: str) -> bool: + # Only what runs here, or is being released here, is protected by + # memory. A dormant entry is not: whether a sandbox is still suspended + # on this worker is the route's to say, and a route released while this + # worker was unreachable leaves a dormant entry behind that nothing + # else clears. The control plane never reports a sandbox whose route + # still names this worker, so a sandbox suspended here is never + # reported orphaned. + return sandbox_id in self.sandboxes or sandbox_id in self._destroying + + async def local_workspace_ids(self) -> list[str]: + """Directories under the workspace root that no live sandbox holds.""" + return await asyncio.to_thread(self._local_workspace_ids) + + def _local_workspace_ids(self) -> list[str]: + root = self.settings.workspace_root + if not root.is_dir(): + return [] + found: list[str] = [] + with os.scandir(root) as entries: + for entry in entries: + if entry.name in _RESERVED_ROOT_NAMES or entry.name.startswith("."): + continue + if not entry.is_dir(follow_symlinks=False): + continue + if self._owned_here(entry.name): + continue + found.append(entry.name) + return sorted(found) + + async def reclaim_orphan_workspaces( + self, orphaned: list[str], *, ttl_seconds: int + ) -> list[str]: + """Delete the directories that have been orphaned for at least `ttl_seconds`. + + `orphaned` is the full set the control plane found this cycle: a marker + for any id not in it is dropped, so ownership coming back resets the + clock. Returns the ids whose directory was deleted. + """ + return await asyncio.to_thread( + self._reclaim_orphan_workspaces, set(orphaned), ttl_seconds, time.time() + ) + + def _reclaim_orphan_workspaces( + self, orphaned: set[str], ttl_seconds: int, now: float + ) -> list[str]: + markers = self._orphan_marker_root() + markers.mkdir(mode=0o700, parents=True, exist_ok=True) + for marker in list(markers.iterdir()): + if marker.name not in orphaned: + marker.unlink(missing_ok=True) + removed: list[str] = [] + for sandbox_id in sorted(orphaned): + try: + root = _sandbox_root(self.settings.workspace_root, sandbox_id) + except ValueError: + continue + marker = markers / sandbox_id + if root.is_symlink() or not root.is_dir() or self._owned_here(sandbox_id): + marker.unlink(missing_ok=True) + continue + try: + first_seen = marker.stat().st_mtime + except FileNotFoundError: + marker.touch(mode=0o600) + continue + if now - first_seen < ttl_seconds: + continue + try: + with self._template_lock(sandbox_id): + # Re-checked under the lock every create, resume and + # release takes: a sandbox that came back here meanwhile + # keeps its directory. + if self._owned_here(sandbox_id): + marker.unlink(missing_ok=True) + continue + # A dormant entry for a route released or moved elsewhere + # is stale; it goes with the directory. + self.dormant.pop(sandbox_id, None) + self._attached_templates.pop(sandbox_id, None) + if root.exists() and (trashed := self._move_to_trash(sandbox_id, root)): + shutil.rmtree(trashed, True) + marker.unlink(missing_ok=True) + removed.append(sandbox_id) + except RuntimeError: + # Busy: a create, resume or release holds the lock. Next cycle. + continue + return removed + + def _has_processes(self, sandbox_id: str) -> bool: + return any(sid == sandbox_id for sid, _ in self.processes) + + async def reclaim_dormant_disk(self) -> list[str]: + """Evict local copies of snapshotted dormant sandboxes under disk pressure. + + Oldest first, and only while the disk is over its watermark. A sandbox + with no snapshot is never evicted here: its directory is the only copy, + and only retention expiry (a release) may delete it. An evicted sandbox + stays suspended; its resume restores from the snapshot. + """ + evicted: list[str] = [] + if self.disk_status()["available"]: + return evicted + candidates = sorted( + (item for item in self.dormant.values() if item.snapshot and not item.evicted), + key=lambda item: item.suspended_at, + ) + for item in candidates: + if await asyncio.to_thread(self._evict_dormant, item.sandbox_id): + evicted.append(item.sandbox_id) + logger.info("evicted dormant workspace under disk pressure: %s", item.sandbox_id) + if self.disk_status()["available"]: + break + return evicted + + async def adopt_dormant( + self, + sandbox_id: str, + *, + generation: int, + uid: int, + suspended_at: float, + snapshot: bool, + ) -> bool: + """Remember a dormant directory this process has forgotten, so disk pressure can evict it. + + A restart empties `dormant`, and the directories it described stay on + disk. The control plane knows which suspended routes still name this + worker and whether their snapshot exists; this takes that back. It + changes no file. Returns whether the sandbox is newly remembered. + """ + return await asyncio.to_thread( + self._adopt_dormant, sandbox_id, generation, uid, suspended_at, snapshot + ) + + def _adopt_dormant( + self, sandbox_id: str, generation: int, uid: int, suspended_at: float, snapshot: bool + ) -> bool: + try: + root = _sandbox_root(self.settings.workspace_root, sandbox_id) + with self._template_lock(sandbox_id): + if ( + sandbox_id in self.sandboxes + or sandbox_id in self.dormant + or sandbox_id in self._destroying + or root.is_symlink() + or not root.is_dir() + ): + return False + self.dormant[sandbox_id] = DormantSandbox( + sandbox_id=sandbox_id, + generation=generation, + uid=uid, + suspended_at=suspended_at, + snapshot=snapshot, + ) + return True + except (RuntimeError, ValueError): + # A create, resume or release holds the lock, or the id is not a + # directory name: leave it to the next cycle. + return False + + def _evict_dormant(self, sandbox_id: str) -> bool: + try: + with self._template_lock(sandbox_id): + item = self.dormant.get(sandbox_id) + if item is None or not item.snapshot or sandbox_id in self.sandboxes: + return False + root = _sandbox_root(self.settings.workspace_root, sandbox_id) + if root.exists() and (trashed := self._move_to_trash(sandbox_id, root)): + # Removed now rather than by the next sweep: the point is + # to free the disk this cycle. + shutil.rmtree(trashed, True) + item.evicted = True + return True + except RuntimeError: + # A resume or release holds the lock; it decides this sandbox. + return False + async def shutdown(self) -> None: for process in list(self.processes.values()): await self._terminate(process) @@ -1019,8 +1394,11 @@ def _resolve_sandbox_path(sandbox: LocalSandbox, virtual_path: str) -> Path: raise ValueError(f"template source path must live under {allowed}") +_RESERVED_ROOT_NAMES = frozenset({".trash", ".staging"}) + + def _sandbox_root(workspace_root: Path, sandbox_id: str) -> Path: - if sandbox_id == ".trash": + if sandbox_id in _RESERVED_ROOT_NAMES: raise ValueError("sandbox id uses a reserved name") root = workspace_root.resolve(strict=False) candidate = (root / sandbox_id).resolve(strict=False) diff --git a/src/agent_sandbox/schemas.py b/src/agent_sandbox/schemas.py index e57740e..f1e9adf 100644 --- a/src/agent_sandbox/schemas.py +++ b/src/agent_sandbox/schemas.py @@ -90,6 +90,38 @@ class ConnectSandboxRequest(BaseModel): generation: int = Field(ge=1) +class SuspendSandboxRequest(BaseModel): + """Release a sandbox's capacity slot and keep its workspace until resumed.""" + + generation: int = Field(ge=1) + + +class SuspendResponse(BaseModel): + sandbox_id: str + generation: int + status: str + # Whether the object store holds a copy, so a resume can land on any worker. + snapshot: bool = False + # False when the sandbox was already suspended: the call changed nothing. + suspended: bool = True + # When the reaper will release it; None when retention is disabled. + retained_until: UtcDatetime | None = None + + +class ResumeLocalRequest(CreateSandboxRequest): + """Worker protocol: register a dormant sandbox again.""" + + restore: Literal["reuse", "snapshot"] = "reuse" + + +class ResumeResponse(RouteResponse): + # False when the sandbox was already awake: the call changed nothing. + resumed: bool = True + # `active`: it was not suspended. `reused`: its directory was still on the + # worker. `restored`: it was rebuilt from the snapshot. + workspace_source: Literal["active", "reused", "restored"] = "reused" + + class ExecRequest(BaseModel): exec_id: str = Field(min_length=3, max_length=128, pattern=r"^[A-Za-z0-9_.:-]+$") generation: int = Field(ge=1) diff --git a/src/agent_sandbox/service.py b/src/agent_sandbox/service.py index f47249f..eb47932 100644 --- a/src/agent_sandbox/service.py +++ b/src/agent_sandbox/service.py @@ -6,23 +6,29 @@ import logging import time import uuid -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from typing import Any, Literal, cast import httpx from .backends import ( + DormantLifecycle, ExecutionBackend, + as_dormant_adoption, + as_dormant_lifecycle, + as_orphan_workspace_reclaim, as_template_cache_pruning, create_execution_backend, ) from .config import Settings from .credentials import CredentialBroker, create_credential_broker +from .dormant import delete_snapshot, read_manifest from .models import ( AUDIT_HISTORY_MAX_BYTES, AUDIT_HISTORY_MAX_ENTRIES, AUDIT_HISTORY_RETENTION_DAYS, Route, + as_utc, ) from .observability import error_code from .preflight import ( @@ -37,12 +43,19 @@ ExecResponse, LifecycleAuditEntry, ResolveRequest, + ResumeResponse, RouteResponse, SandboxAuditResponse, + SuspendResponse, TemplateAttachRequest, TemplateBuildRequest, ) -from .storage import MetadataStore, as_exec_history_pruning +from .storage import ( + DormantRouteStore, + MetadataStore, + as_dormant_route_store, + as_exec_history_pruning, +) from .templates import ( TemplateCatalog, TemplateRecord, @@ -65,7 +78,11 @@ "PROFILE_UPGRADE", "WORKER_REASSIGNED", "RELEASE_RETRY", + "SUSPEND_EXPIRED", ] +# A suspended route is parked: no exec, file call, or proxy reaches it until a +# resume, which is the only way back to ASSIGNED/READY. +DORMANT_STATUSES = frozenset({"SUSPENDING", "SUSPENDED"}) def workspace_actor(route: Route) -> str: @@ -322,6 +339,10 @@ async def _maintenance(self) -> None: cast(int, self.reaper_status["failures_total"]) + 1 ) logger.exception("failed to reclaim orphaned sandbox: %s", route.sandbox_id) + dormant_released, dormant_gone = await self._reclaim_dormant() + released += dormant_released + gone_already += dormant_gone + await self._reclaim_orphan_workspaces() self.reaper_status.update( { "last_run_at": datetime.now(UTC).isoformat(), @@ -373,7 +394,12 @@ def healthy(self) -> bool: return time.monotonic() - self.last_heartbeat_at < self.settings.heartbeat_ttl_seconds async def validate_local_route( - self, sandbox_id: str, generation: int, *, allow_releasing: bool = False + self, + sandbox_id: str, + generation: int, + *, + allow_releasing: bool = False, + allow_dormant: bool = False, ) -> Route: route = await self.database.find_route(sandbox_id) invalid_statuses = {"RELEASED"} if allow_releasing else {"RELEASED", "RELEASING"} @@ -385,6 +411,8 @@ async def validate_local_route( or route.status in invalid_statuses ): raise RuntimeError("STALE_SANDBOX_ROUTE") + if route.status in DORMANT_STATUSES and not allow_dormant: + raise RuntimeError("SANDBOX_SUSPENDED") return route async def proxy_to_worker( @@ -411,6 +439,9 @@ async def proxy_to_worker( or route.status in {"RELEASED", "RELEASING"} ): raise RuntimeError("STALE_SANDBOX_ROUTE") + if route.status in DORMANT_STATUSES: + # Answered here, so a suspended sandbox costs its worker nothing. + raise RuntimeError("SANDBOX_SUSPENDED") worker = await self.registry.get(route.worker_id) if ( not worker @@ -463,6 +494,12 @@ async def resolve(self, request: ResolveRequest) -> RouteResponse: raise RuntimeError("SANDBOX_SCOPE_MISMATCH") if route.profile_id != request.profile: raise RuntimeError("SANDBOX_PROFILE_MISMATCH") + if route.status in DORMANT_STATUSES: + # Resolving is how an agent comes back, so it wakes the sandbox + # rather than reassigning it, which would discard the workspace. + return await self.resume( + request.sandbox_id, created_by=f"workspace:{request.workspace_scope_id}" + ) if route.profile_hash != self.settings.profile_hash: # profile_hash is only a runtime capability fingerprint, so an old route # necessarily lags after an upgrade. Rebuild the route (new worker, bumped @@ -691,6 +728,439 @@ async def _finish_release( reason=reason, released_by=released_by, ) + if route.status in DORMANT_STATUSES or route.status == "RELEASING": + # A RELEASING route may be the retry of a suspended one whose first + # attempt failed, and its snapshot is only deleted here. Deleting + # is a no-op when no manifest exists. + await self._delete_dormant_snapshot(route.sandbox_id) + + # ── suspend and resume ── + # + # A suspended sandbox keeps its workspace and gives back its capacity slot. + # Nothing is frozen: suspend is refused while a command runs, and a resume + # starts with no processes. See docs/SUSPEND_RESUME.md. + + def _dormant_store(self) -> DormantRouteStore: + store = as_dormant_route_store(self.database) + if store is None: + raise RuntimeError("SANDBOX_SUSPEND_UNSUPPORTED") + return store + + def dormant_runtime(self) -> DormantLifecycle: + runtime = as_dormant_lifecycle(self.runtime) + if runtime is None: + raise RuntimeError("SANDBOX_SUSPEND_UNSUPPORTED") + return runtime + + def snapshot_on_suspend(self, route: Route) -> bool: + """Whether a suspend of this route also archives it to the object store.""" + mode = self.settings.suspend_snapshot + if route.storage_mode == "shared" or mode == "never": + # A shared workspace is reachable from every worker already. + return False + has_store = self.template_catalog.object_store is not None + if mode == "always" and not has_store: + raise RuntimeError("SANDBOX_SNAPSHOT_UNAVAILABLE") + return has_store + + def _retained_until(self, route: Route) -> datetime | None: + retention = self.settings.suspended_retention_seconds + if retention <= 0: + return None + start = route.last_active_at or datetime.now(UTC) + return as_utc(start) + timedelta(seconds=retention) + + async def suspend( + self, sandbox_id: str, generation: int, *, actor: str | None = None + ) -> SuspendResponse: + """Release the sandbox's slot, keep its workspace. Idempotent. + + The route is claimed (READY -> SUSPENDING) by a conditional update on + the same row exec admission locks, so a suspend and an exec cannot both + win: a running command makes this answer `SANDBOX_SUSPEND_BUSY`, and an + exec after the claim answers `SANDBOX_SUSPENDED`. Other sandboxes, and + other scopes, are never touched. + """ + store = self._dormant_store() + route = await self.database.find_route(sandbox_id) + if route is None: + raise RuntimeError("SANDBOX_NOT_FOUND") + if route.generation != generation: + raise RuntimeError("STALE_SANDBOX_GENERATION") + if not await store.begin_suspend(sandbox_id, generation): + return self._suspend_response(route, suspended=False) + snapshot = await self._finish_suspend(route) + current = await self.database.find_route(sandbox_id) + logger.info( + "suspended sandbox: sandbox_id=%s generation=%s snapshot=%s actor=%s", + sandbox_id, + generation, + snapshot, + actor or workspace_actor(route), + ) + return self._suspend_response(current or route, snapshot=snapshot) + + def _suspend_response( + self, route: Route, *, suspended: bool = True, snapshot: bool | None = None + ) -> SuspendResponse: + return SuspendResponse( + sandbox_id=route.sandbox_id, + generation=route.generation, + status=route.status, + snapshot=bool(snapshot), + suspended=suspended, + retained_until=self._retained_until(route), + ) + + async def _finish_suspend(self, route: Route) -> bool: + """Have the owning worker release the slot, then park the route.""" + store = self._dormant_store() + worker = await self.registry.get(route.worker_id) + snapshot = False + if worker and worker.get("worker_epoch") == route.worker_epoch: + try: + payload = await self._call_worker( + worker, + "POST", + f"/internal/v1/sandboxes/{route.sandbox_id}/suspend", + generation=route.generation, + json_body={"generation": route.generation}, + # A snapshot archives and uploads the workspace. + timeout_seconds=self.settings.template_snapshot_timeout_seconds + 60, + ) + except RuntimeError as exc: + if str(exc) != "SANDBOX_WORKER_UNREACHABLE": + # The worker refused (busy, snapshot failed): nothing was + # released, so the sandbox goes back to READY untouched. + await store.abort_suspend(route.sandbox_id, route.generation) + # Unreachable stays SUSPENDING; a retry or the reaper finishes it. + raise + snapshot = bool(payload.get("snapshot")) + # A worker that is gone holds no slot and runs nothing, so there is + # nothing to release; resume decides where the workspace comes from. + await store.finish_suspend(route.sandbox_id, route.generation) + return snapshot + + async def resume(self, sandbox_id: str, *, created_by: str | None = None) -> ResumeResponse: + """Wake a suspended sandbox; a no-op for one that is awake. + + Prefers the worker that kept the directory. Otherwise a shared + workspace moves to any worker, and a local one moves only if the object + store holds its snapshot. Moving bumps the generation. + """ + store = self._dormant_store() + for _ in range(3): + route = await self.database.find_route(sandbox_id) + if route is None: + raise RuntimeError("SANDBOX_NOT_FOUND") + if route.status in {"RELEASED", "RELEASING"}: + raise RuntimeError("STALE_SANDBOX_ROUTE") + if route.status == "SUSPENDING": + raise RuntimeError("SANDBOX_SUSPEND_IN_PROGRESS") + if route.status != "SUSPENDED": + return await self._awake_response(route) + actor = created_by or workspace_actor(route) + worker, bump, restore = await self._resume_plan(route) + claimed = await store.begin_resume( + route, + worker, + bump_generation=bump, + profile_hash=self.settings.profile_hash, + created_by=actor, + ) + if claimed is None: + # Another resume, or a release, got there first; look again. + continue + try: + payload = await self._call_worker( + worker, + "POST", + f"/internal/v1/sandboxes/{sandbox_id}/resume", + generation=claimed.generation, + json_body={ + "generation": claimed.generation, + "sandbox_uid": claimed.sandbox_uid, + "profile": claimed.profile_id, + "worker_epoch": claimed.worker_epoch, + "restore": restore, + }, + timeout_seconds=self.settings.template_snapshot_timeout_seconds + 60, + ) + except BaseException: + # Back to sleep, where a retry finds it. The workspace was not + # touched, or a failed restore was staged and thrown away. + await store.abort_resume(sandbox_id, claimed.generation) + raise + source = str(payload.get("workspace_source") or "reused") + current = await self.database.find_route(sandbox_id) or claimed + logger.info( + "resumed sandbox: sandbox_id=%s generation=%s worker_id=%s source=%s", + sandbox_id, + current.generation, + current.worker_id, + source, + ) + return ResumeResponse( + **self.route_response(current, worker).model_dump(), + resumed=True, + workspace_source=cast("Literal['active', 'reused', 'restored']", source), + ) + raise RuntimeError("SANDBOX_BUSY_OR_STALE") + + async def _awake_response(self, route: Route) -> ResumeResponse: + worker = await self.registry.get(route.worker_id) + if not worker or worker.get("worker_epoch") != route.worker_epoch: + # Awake but orphaned: the same recovery resolve performs. + resolved = await self.resolve( + ResolveRequest( + sandbox_id=route.sandbox_id, + workspace_scope_id=route.workspace_scope_id, + profile=route.profile_id, + ) + ) + return ResumeResponse( + **resolved.model_dump(exclude={"resumed", "workspace_source"}), + resumed=False, + workspace_source="active", + ) + return ResumeResponse( + **self.route_response(route, worker).model_dump(), + resumed=False, + workspace_source="active", + ) + + async def _resume_plan(self, route: Route) -> tuple[dict[str, Any], bool, str]: + """Pick (worker, bump_generation, restore) for a suspended route.""" + profile_hash = self.settings.profile_hash + original = await self.registry.get(route.worker_id) + if ( + original + and original.get("worker_id") == route.worker_id + and original.get("status") == "ACTIVE" + and original.get("profile_hash") == profile_hash + and int(original.get("running_sessions", 0)) < int(original.get("capacity", 0)) + ): + # Same worker: its directory is reused. Only a restarted worker (new + # epoch) or a new profile needs a new generation. + bump = ( + original.get("worker_epoch") != route.worker_epoch + or route.profile_hash != profile_hash + ) + return original, bump, "reuse" + if route.storage_mode == "shared": + return await self.registry.select(profile_hash=profile_hash), True, "reuse" + if await self._dormant_snapshot_exists(route.sandbox_id): + candidate = await self.registry.select( + profile_hash=profile_hash, + exclude={route.worker_id} if route.worker_id else None, + ) + return candidate, True, "snapshot" + if original: + # Its worker is alive but full or draining: the slot comes back. + raise RuntimeError("NO_SANDBOX_WORKER_AVAILABLE") + # The only copy is on a worker that is not running. It may return; + # until then, or until retention expires, there is nothing to resume. + raise RuntimeError("SANDBOX_DORMANT_WORKSPACE_UNAVAILABLE") + + async def _dormant_snapshot_exists(self, sandbox_id: str) -> bool: + object_store = self.template_catalog.object_store + if object_store is None: + return False + manifest = await asyncio.to_thread(read_manifest, object_store, sandbox_id) + return manifest is not None + + async def _delete_dormant_snapshot(self, sandbox_id: str) -> None: + object_store = self.template_catalog.object_store + if object_store is None: + return + try: + await asyncio.to_thread(delete_snapshot, object_store, sandbox_id) + except Exception: + logger.warning( + "could not delete a dormant snapshot: sandbox_id=%s", sandbox_id, exc_info=True + ) + + async def delete_dormant_snapshot(self, sandbox_id: str) -> None: + """Best-effort: a resumed sandbox's snapshot is stale the moment it runs.""" + await self._delete_dormant_snapshot(sandbox_id) + + async def _reclaim_dormant(self) -> tuple[int, int]: + """Retention expiry and disk pressure, for suspended sandboxes. + + Expired suspended routes are released, which deletes the directory on + their worker and the snapshot in the object store. Stalled suspends are + finished. Under disk pressure this worker also drops local copies of + dormant sandboxes whose snapshot exists; they stay suspended and resume + from the snapshot. + """ + released = 0 + gone_already = 0 + runtime = as_dormant_lifecycle(self.runtime) + if runtime is not None: + try: + await self._adopt_forgotten_dormant(runtime) + evicted = await runtime.reclaim_dormant_disk() + if evicted: + self.reaper_status["dormant_evicted_total"] = cast( + int, self.reaper_status.get("dormant_evicted_total", 0) + ) + len(evicted) + except Exception: + logger.exception("dormant disk reclamation failed; retrying next cycle") + store = as_dormant_route_store(self.database) + if store is None: + return released, gone_already + routes = await store.list_dormant_routes_to_reclaim( + retention_seconds=self.settings.suspended_retention_seconds, + suspending_grace_seconds=self.settings.orphan_release_grace_seconds, + limit=self.settings.orphan_reaper_batch_size, + ) + for route in routes: + try: + if route.status == "SUSPENDING": + await self._finish_suspend(route) + continue + await self.release(route, reason="SUSPEND_EXPIRED", released_by=SYSTEM_REAPER_ACTOR) + released += 1 + logger.info( + "released expired suspended sandbox: sandbox_id=%s generation=%s", + route.sandbox_id, + route.generation, + ) + except Exception as exc: + if error_code(exc) in _ROUTE_GONE_CODES: + # Counted like the idle sweep counts it: a client released + # or resumed it between the listing and this call. + gone_already += 1 + continue + self.reaper_status["failures_total"] = ( + cast(int, self.reaper_status["failures_total"]) + 1 + ) + logger.exception("failed to reclaim suspended sandbox: %s", route.sandbox_id) + return released, gone_already + + async def _adopt_forgotten_dormant(self, runtime: DormantLifecycle) -> None: + """Hand a restarted worker back the dormant copies it may evict. + + Only under disk pressure, so a healthy disk costs nothing. A directory + qualifies when its route is suspended on this worker and its snapshot is + in the object store; without one it is the only copy and stays. + """ + adoption = as_dormant_adoption(self.runtime) + local = as_orphan_workspace_reclaim(self.runtime) + if ( + adoption is None + or local is None + or self.settings.shared_root is not None + or self.template_catalog.object_store is None + or self.runtime.disk_status()["available"] + ): + return + for sandbox_id in await local.local_workspace_ids(): + if sandbox_id in runtime.dormant: + continue + route = await self.database.find_route(sandbox_id) + if ( + route is None + or route.status != "SUSPENDED" + or route.worker_id != self.worker_id + or route.storage_mode == "shared" + or not await self._dormant_snapshot_exists(sandbox_id) + ): + continue + since = as_utc(route.updated_at).timestamp() if route.updated_at else time.time() + if await adoption.adopt_dormant( + sandbox_id, + generation=route.generation, + uid=route.sandbox_uid, + suspended_at=since, + snapshot=True, + ): + logger.info("took back a dormant workspace after a restart: %s", sandbox_id) + + async def _reclaim_orphan_workspaces(self) -> int: + """Delete local workspace directories this worker no longer owns. + + A directory is orphaned when its route is released, or points at + another worker: the sandbox resumed elsewhere from its snapshot, or was + reassigned while this worker was away. It is deleted once it has stayed + orphaned for `SANDBOX_ORPHAN_DORMANT_DIR_TTL_SECONDS`. A directory whose + route still names this worker (active, suspended, or being created) is + never a candidate, and neither is one with no route at all, which this + metadata store cannot vouch for. Shared workspaces are not this + worker's to delete. + """ + ttl = self.settings.orphan_dormant_dir_ttl_seconds + if ttl <= 0 or self.settings.shared_root is not None: + return 0 + runtime = as_orphan_workspace_reclaim(self.runtime) + if runtime is None: + return 0 + try: + orphaned: list[str] = [] + for sandbox_id in await runtime.local_workspace_ids(): + route = await self.database.find_route(sandbox_id) + if route is None or route.storage_mode == "shared": + continue + if route.status == "RELEASED" or route.worker_id != self.worker_id: + orphaned.append(sandbox_id) + removed = await runtime.reclaim_orphan_workspaces(orphaned, ttl_seconds=ttl) + except Exception: + logger.exception("orphaned workspace reclamation failed; retrying next cycle") + return 0 + for sandbox_id in removed: + logger.info("deleted orphaned workspace directory: sandbox_id=%s", sandbox_id) + if removed: + self.reaper_status["orphan_workspaces_deleted_total"] = cast( + int, self.reaper_status.get("orphan_workspaces_deleted_total", 0) + ) + len(removed) + return len(removed) + + async def refresh_capacity(self) -> None: + """Publish a heartbeat now, so a released or retaken slot is visible at once.""" + try: + await self._heartbeat() + except Exception: + logger.warning("capacity heartbeat failed; the next cycle publishes it") + + async def _call_worker( + self, + worker: dict[str, Any], + method: str, + path: str, + *, + generation: int, + json_body: dict[str, Any] | None = None, + timeout_seconds: float = 30.0, + ) -> dict[str, Any]: + headers = { + "Authorization": f"Bearer {self.settings.internal_token}", + "X-Sandbox-Worker-ID": str(worker["worker_id"]), + "X-Sandbox-Generation": str(generation), + } + try: + async with httpx.AsyncClient( + timeout=httpx.Timeout(timeout_seconds, connect=10.0), headers=headers + ) as client: + response = await client.request( + method, f"{str(worker['endpoint']).rstrip('/')}{path}", json=json_body + ) + except httpx.RequestError as exc: + raise RuntimeError("SANDBOX_WORKER_UNREACHABLE") from exc + if not response.is_success: + raise RuntimeError(_worker_error_code(response)) + payload = response.json() + return payload if isinstance(payload, dict) else {} + + +def _worker_error_code(response: httpx.Response) -> str: + """The code a worker refused with: plain text, or FastAPI's `detail`.""" + try: + data = response.json() + except ValueError: + data = None + if isinstance(data, dict) and isinstance(data.get("detail"), str): + return str(data["detail"]) + return response.text.strip() or f"SANDBOX_WORKER_HTTP_{response.status_code}" def _exec_from_row(row: dict[str, Any]) -> ExecResponse: diff --git a/src/agent_sandbox/sql_database.py b/src/agent_sandbox/sql_database.py index 30a58bb..e7dcc04 100644 --- a/src/agent_sandbox/sql_database.py +++ b/src/agent_sandbox/sql_database.py @@ -328,39 +328,14 @@ async def reassign_route( return current row = dict(row_mapping) now = utc_now_naive() - history = bounded_lifecycle_history(row.get("lifecycle_history_json"), now=now) - values: dict[str, Any] = { - "worker_id": worker["worker_id"], - "worker_epoch": worker["worker_epoch"], - "profile_hash": new_hash, - "generation": route.generation + 1, - "status": "ASSIGNED", - "active_exec_id": None, - "last_active_at": now, - "generation_started_at": now, - "generation_created_by": created_by, - "ready_at": None, - "lifecycle_count": int(row.get("lifecycle_count") or 1) + 1, - "updated_at": now, - } - if row["status"] != "RELEASED": - event = completed_lifecycle( - row, - released_at=now, - reason=reason, - released_by="system:route-reassign", - ) - history = bounded_lifecycle_history(history, event=event, now=now) - lifetime_ms = int(event["lifetime_ms"]) - values.update( - last_released_generation=route.generation, - last_released_at=now, - last_release_reason=reason, - last_released_by="system:route-reassign", - last_lifetime_ms=lifetime_ms, - total_lifetime_ms=int(row.get("total_lifetime_ms") or 0) + lifetime_ms, - ) - values["lifecycle_history_json"] = history + values = _reassignment_values( + row, + worker, + profile_hash=new_hash, + reason=reason, + created_by=created_by, + now=now, + ) updated = await connection.execute( update(route_table) .where( @@ -420,7 +395,8 @@ async def touch_route(self, sandbox_id: str, generation: int) -> None: .where( route_table.c.sandbox_id == sandbox_id, route_table.c.generation == generation, - route_table.c.status.not_in(("RELEASED", "RELEASING")), + # A dormant route's last_active_at is its retention clock. + route_table.c.status.not_in(_INACTIVE_STATUSES), ) .values(last_active_at=now, updated_at=now) ) @@ -683,6 +659,14 @@ async def begin_exec( .with_for_update() ) route = route_result.mappings().first() + if ( + route is not None + and int(route["generation"]) == generation + and str(route["status"]) in _DORMANT_STATUSES + ): + # Admission and suspend both decide on this locked row, so an + # exec either got in before the suspend or is refused here. + raise RuntimeError("SANDBOX_SUSPENDED") if ( route is None or int(route["generation"]) != generation @@ -722,11 +706,15 @@ async def begin_exec( ) # active_exec_id keeps one representative execution for operators; # it is no longer an admission gate. - await connection.execute( + admitted = await connection.execute( update(route_table) .where( route_table.c.sandbox_id == sandbox_id, route_table.c.generation == generation, + # Re-checked in the write itself: SQLite has no FOR UPDATE, + # so the read above can predate a suspend or release that + # committed since. Zero rows rolls the insert back. + route_table.c.status.in_(("ASSIGNED", "READY", "RUNNING")), ) .values( active_exec_id=func.coalesce(route_table.c.active_exec_id, exec_id), @@ -735,6 +723,20 @@ async def begin_exec( updated_at=now, ) ) + if admitted.rowcount != 1: + status_now = ( + await connection.execute( + select(route_table.c.status).where( + route_table.c.sandbox_id == sandbox_id, + route_table.c.generation == generation, + ) + ) + ).scalar() + raise RuntimeError( + "SANDBOX_SUSPENDED" + if status_now in _DORMANT_STATUSES + else "STALE_SANDBOX_GENERATION" + ) return None async def finish_exec( @@ -825,6 +827,187 @@ async def get_exec(self, sandbox_id: str, exec_id: str) -> dict[str, Any] | None row = result.mappings().first() return _exec_row(dict(row)) if row else None + # ── suspend and resume ── + # + # READY -> SUSPENDING -> SUSPENDED -> ASSIGNED -> READY. Each step is a + # conditional update on (sandbox_id, generation, status), so it either + # happens exactly once or reports what happened instead. + + async def begin_suspend(self, sandbox_id: str, generation: int) -> bool: + """Claim an idle route for suspension. + + True when this call claimed it, or a previous suspend is still in + flight and should be finished; False when it is already suspended. + Raises `SANDBOX_SUSPEND_BUSY` while a command runs, and + `STALE_SANDBOX_GENERATION` when the caller's generation is not current. + """ + now = utc_now_naive() + async with self._engine().begin() as connection: + result = await connection.execute( + update(route_table) + .where( + route_table.c.sandbox_id == sandbox_id, + route_table.c.generation == generation, + route_table.c.active_exec_id.is_(None), + route_table.c.status == "READY", + ) + .values(status="SUSPENDING", updated_at=now) + ) + if result.rowcount == 1: + return True + current = await self.find_route(sandbox_id) + if current is None or current.generation != generation: + raise RuntimeError("STALE_SANDBOX_GENERATION") + if current.status == "SUSPENDED": + return False + if current.status == "SUSPENDING": + return True + if current.status == "RUNNING": + raise RuntimeError("SANDBOX_SUSPEND_BUSY") + if current.status == "ASSIGNED": + # Resolved but never created: there is no slot to release yet. + raise RuntimeError("SANDBOX_NOT_READY") + raise RuntimeError("STALE_SANDBOX_GENERATION") + + async def finish_suspend(self, sandbox_id: str, generation: int) -> None: + now = utc_now_naive() + async with self._engine().begin() as connection: + await connection.execute( + update(route_table) + .where( + route_table.c.sandbox_id == sandbox_id, + route_table.c.generation == generation, + route_table.c.status == "SUSPENDING", + ) + # last_active_at starts the retention clock. + .values(status="SUSPENDED", active_exec_id=None, last_active_at=now, updated_at=now) + ) + current = await self.find_route(sandbox_id) + if current is None or current.generation != generation or current.status != "SUSPENDED": + raise RuntimeError("STALE_SANDBOX_GENERATION") + + async def abort_suspend(self, sandbox_id: str, generation: int) -> None: + now = utc_now_naive() + async with self._engine().begin() as connection: + await connection.execute( + update(route_table) + .where( + route_table.c.sandbox_id == sandbox_id, + route_table.c.generation == generation, + route_table.c.status == "SUSPENDING", + ) + .values(status="READY", updated_at=now) + ) + + async def begin_resume( + self, + route: Route, + worker: dict[str, Any], + *, + bump_generation: bool, + profile_hash: str, + created_by: str, + ) -> Route | None: + """Claim a suspended route for one worker, or return None if another call won. + + Without `bump_generation` the route stays on its worker at its + generation, so a client holding that generation keeps using it. Moving + to another worker, or to a restarted one, is a reassignment and bumps + the generation exactly like any other, which fences the old owner. + """ + async with self._engine().begin() as connection: + result = await connection.execute( + select(route_table) + .where( + route_table.c.sandbox_id == route.sandbox_id, + route_table.c.generation == route.generation, + route_table.c.status == "SUSPENDED", + ) + .with_for_update() + ) + row_mapping = result.mappings().first() + if row_mapping is None: + return None + now = utc_now_naive() + if bump_generation: + values = _reassignment_values( + dict(row_mapping), + worker, + profile_hash=profile_hash, + reason="RESUME_REASSIGNED", + created_by=created_by, + now=now, + ) + else: + values = { + "status": "ASSIGNED", + "worker_epoch": worker["worker_epoch"], + "last_active_at": now, + "updated_at": now, + } + updated = await connection.execute( + update(route_table) + .where( + route_table.c.sandbox_id == route.sandbox_id, + route_table.c.generation == route.generation, + route_table.c.status == "SUSPENDED", + ) + .values(**values) + ) + if updated.rowcount != 1: + return None + return await self.find_route(route.sandbox_id) + + async def abort_resume(self, sandbox_id: str, generation: int) -> None: + """Put a route whose resume failed back to sleep, where a retry can find it.""" + now = utc_now_naive() + async with self._engine().begin() as connection: + await connection.execute( + update(route_table) + .where( + route_table.c.sandbox_id == sandbox_id, + route_table.c.generation == generation, + route_table.c.status == "ASSIGNED", + ) + .values(status="SUSPENDED", updated_at=now) + ) + + async def list_dormant_routes_to_reclaim( + self, + *, + retention_seconds: int, + suspending_grace_seconds: int, + limit: int, + ) -> list[Route]: + """Suspended routes past retention, and suspends that stalled midway. + + `retention_seconds=0` keeps suspended routes indefinitely; stalled + suspends are still returned so they can be finished. + """ + now = utc_now_naive() + conditions = [ + and_( + route_table.c.status == "SUSPENDING", + route_table.c.updated_at < now - timedelta(seconds=suspending_grace_seconds), + ) + ] + if retention_seconds > 0: + conditions.append( + and_( + route_table.c.status == "SUSPENDED", + route_table.c.last_active_at < now - timedelta(seconds=retention_seconds), + ) + ) + async with self._engine().connect() as connection: + result = await connection.execute( + select(route_table) + .where(or_(*conditions)) + .order_by(route_table.c.last_active_at.asc()) + .limit(limit) + ) + rows = result.mappings().all() + return [route_from_mapping(dict(row)) for row in rows] + # ── fleet queries ── # # Every other read is keyed by a sandbox_id the caller already knows. An @@ -938,6 +1121,57 @@ async def list_workers(self) -> list[dict[str, Any]]: return [dict(row) for row in rows] +_DORMANT_STATUSES = frozenset({"SUSPENDING", "SUSPENDED"}) +_INACTIVE_STATUSES = ("RELEASED", "RELEASING", "SUSPENDING", "SUSPENDED") + + +def _reassignment_values( + row: dict[str, Any], + worker: dict[str, Any], + *, + profile_hash: str, + reason: str, + created_by: str, + now: Any, +) -> dict[str, Any]: + """The columns a reassignment writes: next generation, new owner, audit event.""" + generation = int(row["generation"]) + history = bounded_lifecycle_history(row.get("lifecycle_history_json"), now=now) + values: dict[str, Any] = { + "worker_id": worker["worker_id"], + "worker_epoch": worker["worker_epoch"], + "profile_hash": profile_hash, + "generation": generation + 1, + "status": "ASSIGNED", + "active_exec_id": None, + "last_active_at": now, + "generation_started_at": now, + "generation_created_by": created_by, + "ready_at": None, + "lifecycle_count": int(row.get("lifecycle_count") or 1) + 1, + "updated_at": now, + } + if row["status"] != "RELEASED": + event = completed_lifecycle( + row, + released_at=now, + reason=reason, + released_by="system:route-reassign", + ) + history = bounded_lifecycle_history(history, event=event, now=now) + lifetime_ms = int(event["lifetime_ms"]) + values.update( + last_released_generation=generation, + last_released_at=now, + last_release_reason=reason, + last_released_by="system:route-reassign", + last_lifetime_ms=lifetime_ms, + total_lifetime_ms=int(row.get("total_lifetime_ms") or 0) + lifetime_ms, + ) + values["lifecycle_history_json"] = history + return values + + def database_backend_name(settings: Settings) -> str: return normalize_async_database_url(settings.resolved_database_url()).get_backend_name() diff --git a/src/agent_sandbox/storage.py b/src/agent_sandbox/storage.py index 6a04365..e958b54 100644 --- a/src/agent_sandbox/storage.py +++ b/src/agent_sandbox/storage.py @@ -103,6 +103,56 @@ def as_exec_history_pruning(store: object) -> ExecHistoryPruning | None: return cast("ExecHistoryPruning", store) +class DormantRouteStore(Protocol): + """Route transitions for suspend and resume. + + A separate capability, like `ExecHistoryPruning`: a metadata plugin written + against an earlier release keeps working, and suspend/resume answer 501 on + it. Every method is a compare-and-set on (sandbox_id, generation, status), + which is what makes the transitions safe against concurrent exec admission, + release, and a second suspend or resume. See `as_dormant_route_store`. + """ + + async def begin_suspend(self, sandbox_id: str, generation: int) -> bool: ... + async def finish_suspend(self, sandbox_id: str, generation: int) -> None: ... + async def abort_suspend(self, sandbox_id: str, generation: int) -> None: ... + async def begin_resume( + self, + route: Route, + worker: dict[str, Any], + *, + bump_generation: bool, + profile_hash: str, + created_by: str, + ) -> Route | None: ... + async def abort_resume(self, sandbox_id: str, generation: int) -> None: ... + async def list_dormant_routes_to_reclaim( + self, + *, + retention_seconds: int, + suspending_grace_seconds: int, + limit: int, + ) -> list[Route]: ... + + +def as_dormant_route_store(store: object) -> DormantRouteStore | None: + """Narrow a metadata store to suspend/resume transitions, or report it cannot.""" + + if store is None: + return None + required = ( + "begin_suspend", + "finish_suspend", + "abort_suspend", + "begin_resume", + "abort_resume", + "list_dormant_routes_to_reclaim", + ) + if not all(callable(getattr(store, name, None)) for name in required): + return None + return cast("DormantRouteStore", store) + + class FleetQueries(Protocol): """Fleet-wide reads used by the admin console. @@ -163,4 +213,11 @@ def create_metadata_store(settings: Settings) -> MetadataStore: return SqlAlchemyDatabase(settings) -__all__ = ["FleetQueries", "MetadataStore", "as_fleet_queries", "create_metadata_store"] +__all__ = [ + "DormantRouteStore", + "FleetQueries", + "MetadataStore", + "as_dormant_route_store", + "as_fleet_queries", + "create_metadata_store", +] diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index fb68ace..4f1e14e 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -92,6 +92,12 @@ def minio_endpoint() -> str: return MINIO_ENDPOINT +@pytest.fixture(scope="session") +def s3_credentials() -> tuple[str, str]: + """Access and secret key for the S3-compatible store (LocalStack, RustFS, MinIO, ...).""" + return MINIO_ACCESS_KEY, MINIO_SECRET_KEY + + @pytest.fixture def settings_factory(tmp_path: Path) -> Any: """Build Settings pointed at a real database with a per-test workspace.""" diff --git a/tests/integration/test_dormant_adapters.py b/tests/integration/test_dormant_adapters.py new file mode 100644 index 0000000..a619485 --- /dev/null +++ b/tests/integration/test_dormant_adapters.py @@ -0,0 +1,368 @@ +"""Suspend and resume on the middleware the project ships adapters for. + +The unit suite proves the `DormantRouteStore` transitions on SQLite and the +snapshot archive against a directory. These run the same contract on real +MySQL and PostgreSQL servers, where exec admission and suspend race on a +`SELECT ... FOR UPDATE` row lock instead of SQLite's single writer, and the +snapshot against a real S3 API, where deletes and listings behave like the +service they stand in for. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import time +from collections.abc import AsyncIterator, Iterator +from datetime import timedelta +from pathlib import Path +from typing import Any + +import pytest +from sqlalchemy import update + +from agent_sandbox.archive import ArchiveLimits +from agent_sandbox.blobstore import BlobStore +from agent_sandbox.dormant import ( + delete_snapshot, + manifest_key, + read_manifest, + restore_snapshot, + write_snapshot, +) +from agent_sandbox.models import Route, utc_now_naive +from agent_sandbox.sql_database import SqlAlchemyDatabase, metadata, route_table +from agent_sandbox.storage import as_dormant_route_store + +BUCKET = "agent-sandbox-integration" +WORKER_A = {"worker_id": "worker-a", "worker_epoch": "epoch-a"} +WORKER_B = {"worker_id": "worker-b", "worker_epoch": "epoch-b"} + + +@pytest.fixture(params=["mysql_url", "postgres_url"]) +def database_url(request: pytest.FixtureRequest) -> str: + return str(request.getfixturevalue(request.param)) + + +@pytest.fixture +async def store(database_url: str, settings_factory: Any) -> AsyncIterator[SqlAlchemyDatabase]: + database = SqlAlchemyDatabase(settings_factory(database_url)) + await database.connect() + engine = database.engine + assert engine is not None + async with engine.begin() as connection: + await connection.run_sync(metadata.drop_all) + await connection.run_sync(metadata.create_all) + try: + yield database + finally: + async with engine.begin() as connection: + await connection.run_sync(metadata.drop_all) + await database.close() + + +async def _ready(store: SqlAlchemyDatabase, sandbox_id: str = "sandbox-a") -> Route: + for worker in (WORKER_A, WORKER_B): + await store.upsert_worker( + worker_id=worker["worker_id"], + epoch=worker["worker_epoch"], + endpoint=f"http://{worker['worker_id']}:8080", + status="ACTIVE", + running=0, + ) + route = await store.create_route( + sandbox_id=sandbox_id, workspace_scope_id="tenant-a", worker=WORKER_A + ) + await store.mark_route_ready(route.sandbox_id, route.generation, "worker-a", "epoch-a") + current = await store.find_route(sandbox_id) + assert current is not None + return current + + +async def _exec(store: SqlAlchemyDatabase, route: Route, exec_id: str) -> Any: + return await store.begin_exec( + sandbox_id=route.sandbox_id, + exec_id=exec_id, + generation=route.generation, + worker_id="worker-a", + command=json.dumps({"argv": ["/bin/true"], "exec_scope": "thread-a"}), + exec_scope="thread-a", + ) + + +async def _finish(store: SqlAlchemyDatabase, route: Route, exec_id: str) -> None: + await store.finish_exec( + sandbox_id=route.sandbox_id, + exec_id=exec_id, + status="SUCCEEDED", + exit_code=0, + stdout="", + stderr="", + truncated=False, + ) + + +def test_the_sql_adapter_implements_the_dormant_route_store(store: SqlAlchemyDatabase) -> None: + # An adapter that does not would answer 501 SANDBOX_SUSPEND_UNSUPPORTED. + assert as_dormant_route_store(store) is not None + + +async def test_suspend_is_a_claim_then_a_finish_and_repeats_as_a_no_op( + store: SqlAlchemyDatabase, +) -> None: + route = await _ready(store) + + assert await store.begin_suspend(route.sandbox_id, route.generation) is True + assert await store.begin_suspend(route.sandbox_id, route.generation) is True + await store.finish_suspend(route.sandbox_id, route.generation) + suspended = await store.find_route(route.sandbox_id) + assert suspended is not None and suspended.status == "SUSPENDED" + assert await store.begin_suspend(route.sandbox_id, route.generation) is False + with pytest.raises(RuntimeError, match="STALE_SANDBOX_GENERATION"): + await store.begin_suspend(route.sandbox_id, route.generation + 1) + + +async def test_a_running_command_refuses_the_suspend_and_a_suspend_refuses_commands( + store: SqlAlchemyDatabase, +) -> None: + route = await _ready(store) + + assert await _exec(store, route, "exec-1") is None + with pytest.raises(RuntimeError, match="SANDBOX_SUSPEND_BUSY"): + await store.begin_suspend(route.sandbox_id, route.generation) + await _finish(store, route, "exec-1") + assert await store.begin_suspend(route.sandbox_id, route.generation) is True + with pytest.raises(RuntimeError, match="SANDBOX_SUSPENDED"): + await _exec(store, route, "exec-2") + await store.finish_suspend(route.sandbox_id, route.generation) + with pytest.raises(RuntimeError, match="SANDBOX_SUSPENDED"): + await _exec(store, route, "exec-3") + + +async def test_exec_admission_and_suspend_have_exactly_one_winner_every_time( + store: SqlAlchemyDatabase, +) -> None: + """The race the feature turns on, under the server's own row locks.""" + outcomes: set[str] = set() + for round_ in range(12): + route = await _ready(store, f"sandbox-race-{round_}") + results = await asyncio.gather( + _exec(store, route, f"exec-race-{round_}"), + store.begin_suspend(route.sandbox_id, route.generation), + return_exceptions=True, + ) + exec_won = results[0] is None + suspend_won = results[1] is True + assert exec_won != suspend_won, results + current = await store.find_route(route.sandbox_id) + assert current is not None + assert current.status == ("RUNNING" if exec_won else "SUSPENDING") + outcomes.add("exec" if exec_won else "suspend") + assert outcomes # both may win over time; neither may ever both win + + +async def test_one_resume_wins_and_a_move_bumps_and_audits_the_generation( + store: SqlAlchemyDatabase, +) -> None: + route = await _ready(store) + await store.begin_suspend(route.sandbox_id, route.generation) + await store.finish_suspend(route.sandbox_id, route.generation) + suspended = await store.find_route(route.sandbox_id) + assert suspended is not None + + first, second = await asyncio.gather( + store.begin_resume( + suspended, + WORKER_A, + bump_generation=False, + profile_hash="profile-integration", + created_by="a", + ), + store.begin_resume( + suspended, + WORKER_A, + bump_generation=False, + profile_hash="profile-integration", + created_by="b", + ), + ) + winners = [item for item in (first, second) if item is not None] + assert len(winners) == 1 and winners[0].status == "ASSIGNED" + assert winners[0].generation == route.generation + + # Back to sleep, then a resume on another worker: the generation moves. + await store.abort_resume(route.sandbox_id, route.generation) + again = await store.find_route(route.sandbox_id) + assert again is not None and again.status == "SUSPENDED" + moved = await store.begin_resume( + again, + WORKER_B, + bump_generation=True, + profile_hash="profile-integration", + created_by="workspace:tenant-a", + ) + assert moved is not None + assert (moved.worker_id, moved.generation, moved.status) == ( + "worker-b", + route.generation + 1, + "ASSIGNED", + ) + assert moved.last_release_reason == "RESUME_REASSIGNED" + with pytest.raises(RuntimeError, match="STALE_SANDBOX_GENERATION"): + await _exec(store, route, "exec-old-generation") + + +async def test_retention_expiry_and_stalled_suspends_are_listed( + store: SqlAlchemyDatabase, +) -> None: + expired = await _ready(store, "sandbox-expired") + stalled = await _ready(store, "sandbox-stalled") + await store.begin_suspend(expired.sandbox_id, expired.generation) + await store.finish_suspend(expired.sandbox_id, expired.generation) + await store.begin_suspend(stalled.sandbox_id, stalled.generation) + engine = store.engine + assert engine is not None + async with engine.begin() as connection: + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == expired.sandbox_id) + .values(last_active_at=utc_now_naive() - timedelta(hours=2)) + ) + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == stalled.sandbox_id) + .values(updated_at=utc_now_naive() - timedelta(minutes=10)) + ) + + listed = await store.list_dormant_routes_to_reclaim( + retention_seconds=3600, suspending_grace_seconds=300, limit=10 + ) + assert sorted((item.sandbox_id, item.status) for item in listed) == [ + ("sandbox-expired", "SUSPENDED"), + ("sandbox-stalled", "SUSPENDING"), + ] + assert [ + item.sandbox_id + for item in await store.list_dormant_routes_to_reclaim( + retention_seconds=3 * 3600, suspending_grace_seconds=300, limit=10 + ) + ] == ["sandbox-stalled"] + # Zero disables retention but not the completion of a stalled suspend. + assert [ + item.sandbox_id + for item in await store.list_dormant_routes_to_reclaim( + retention_seconds=0, suspending_grace_seconds=300, limit=10 + ) + ] == ["sandbox-stalled"] + # An expired suspended route goes through the ordinary release. + assert await store.begin_release(expired.sandbox_id, expired.generation) is True + await store.release_route( + expired.sandbox_id, expired.generation, reason="SUSPEND_EXPIRED", released_by="system" + ) + released = await store.find_route(expired.sandbox_id) + assert released is not None and released.status == "RELEASED" + assert released.last_release_reason == "SUSPEND_EXPIRED" + + +# ── the snapshot on a real S3 API ── + + +@pytest.fixture +def blob_store(minio_endpoint: str, s3_credentials: tuple[str, str]) -> Iterator[BlobStore]: + import boto3 + from botocore.config import Config + + from agent_sandbox.config import Settings + + client = boto3.client( + "s3", + endpoint_url=minio_endpoint, + aws_access_key_id=s3_credentials[0], + aws_secret_access_key=s3_credentials[1], + region_name="us-east-1", + config=Config(signature_version="s3v4", s3={"addressing_style": "path"}), + ) + try: + client.create_bucket(Bucket=BUCKET) + except (client.exceptions.BucketAlreadyOwnedByYou, client.exceptions.BucketAlreadyExists): + pass + prefix = f"dormant-{time.time_ns()}/" + settings = Settings( + internal_token="integration-token", + blobstore_endpoint=minio_endpoint, + blobstore_bucket=BUCKET, + blobstore_base_prefix=prefix, + blobstore_region="us-east-1", + ) + instance = BlobStore(settings, client=client) + try: + yield instance + finally: + listing = client.list_objects_v2(Bucket=BUCKET, Prefix=prefix) + contents = listing.get("Contents", []) + if contents: + client.delete_objects( + Bucket=BUCKET, Delete={"Objects": [{"Key": item["Key"]} for item in contents]} + ) + + +def _keys(blob_store: BlobStore) -> list[str]: + client = blob_store._s3() + listing = client.list_objects_v2(Bucket=BUCKET, Prefix=blob_store.base_prefix) + return sorted(item["Key"] for item in listing.get("Contents", [])) + + +def test_a_snapshot_round_trips_through_s3_and_is_deleted_completely( + blob_store: BlobStore, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(os, "chown", lambda *_args, **_kwargs: None) + root = tmp_path / "sandbox" + for name in ("workspace", "home", "envs", "cache", "logs"): + (root / name).mkdir(parents=True) + (root / "workspace" / "notes.txt").write_text("state before the wait") + (root / "workspace" / "link").symlink_to("/etc/hostname") + (root / "envs" / "marker").write_text("installed") + limits = ArchiveLimits() + + write_snapshot( + blob_store, + sandbox_id="sb-1", + generation=1, + root=root, + staging_root=tmp_path / "staging", + limits=limits, + ) + + manifest = read_manifest(blob_store, "sb-1") + assert manifest is not None and set(manifest["parts"]) == {"workspace", "home", "envs"} + assert any(key.endswith(manifest_key("sb-1")) for key in _keys(blob_store)) + + restored = tmp_path / "restored" + result = restore_snapshot( + blob_store, + sandbox_id="sb-1", + destination=restored, + uid=os.getuid(), + max_extract_bytes=64 * 1024 * 1024, + ) + assert result is not None + assert (restored / "workspace" / "notes.txt").read_text() == "state before the wait" + assert os.readlink(restored / "workspace" / "link") == "/etc/hostname" + assert (restored / "envs" / "marker").read_text() == "installed" + + # A newer snapshot replaces the older one's parts instead of piling up. + (root / "workspace" / "notes.txt").write_text("second") + write_snapshot( + blob_store, + sandbox_id="sb-1", + generation=1, + root=root, + staging_root=tmp_path / "staging", + limits=limits, + ) + assert len([key for key in _keys(blob_store) if key.endswith(".tar.gz")]) == 3 + + assert delete_snapshot(blob_store, "sb-1") is True + assert _keys(blob_store) == [] + assert delete_snapshot(blob_store, "sb-1") is False diff --git a/tests/integration/test_s3_blobstore.py b/tests/integration/test_s3_blobstore.py index 5021581..4441839 100644 --- a/tests/integration/test_s3_blobstore.py +++ b/tests/integration/test_s3_blobstore.py @@ -24,15 +24,15 @@ @pytest.fixture -def s3_client(minio_endpoint: str) -> Any: +def s3_client(minio_endpoint: str, s3_credentials: tuple[str, str]) -> Any: import boto3 from botocore.config import Config return boto3.client( "s3", endpoint_url=minio_endpoint, - aws_access_key_id="integration", - aws_secret_access_key="integration-secret", + aws_access_key_id=s3_credentials[0], + aws_secret_access_key=s3_credentials[1], region_name="us-east-1", config=Config(signature_version="s3v4", s3={"addressing_style": "path"}), ) diff --git a/tests/test_admin_console.py b/tests/test_admin_console.py index 1365cb1..d7fba90 100644 --- a/tests/test_admin_console.py +++ b/tests/test_admin_console.py @@ -668,6 +668,8 @@ async def test_overview_reports_the_periods_this_replica_reclaims_on(tmp_path: P "idle_ttl_seconds": 45, "orphan_running_grace_seconds": 7, "orphan_release_grace_seconds": 10, + "suspended_retention_seconds": 2592000, + "orphan_dormant_dir_ttl_seconds": 86400, } diff --git a/tests/test_fleet_verifier_contract.py b/tests/test_fleet_verifier_contract.py index 69ed2af..cc9d0a7 100644 --- a/tests/test_fleet_verifier_contract.py +++ b/tests/test_fleet_verifier_contract.py @@ -75,6 +75,8 @@ async def test_health_reports_the_periods_this_deployment_runs_with(tmp_path: Pa "idle_ttl_seconds": 45, "orphan_running_grace_seconds": 7, "orphan_release_grace_seconds": 10, + "suspended_retention_seconds": 2592000, + "orphan_dormant_dir_ttl_seconds": 86400, } @@ -128,6 +130,8 @@ def _periods(**overrides: Any) -> dict[str, float]: "heartbeat_ttl_seconds": 10, "idle_ttl_seconds": 45, "orphan_release_grace_seconds": 10, + "suspended_retention_seconds": 2592000, + "orphan_dormant_dir_ttl_seconds": 86400, } values.update(overrides) return values diff --git a/tests/test_internal_api_fence.py b/tests/test_internal_api_fence.py index 01b8cf2..179976a 100644 --- a/tests/test_internal_api_fence.py +++ b/tests/test_internal_api_fence.py @@ -107,6 +107,23 @@ "/internal/v1/sandboxes/sb-1/templates", {"json": {"generation": 2, "templates": ["env-a"]}}, ), + ( + "POST", + "/internal/v1/sandboxes/sb-1/suspend", + {"json": {"generation": 2}}, + ), + ( + "POST", + "/internal/v1/sandboxes/sb-1/resume", + { + "json": { + "generation": 2, + "sandbox_uid": 20001, + "profile": "coding-default", + "worker_epoch": "epoch-1", + } + }, + ), ( "DELETE", "/internal/v1/sandboxes/sb-1", diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 760d99a..628edfe 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,5 +1,6 @@ from __future__ import annotations +import importlib.util from collections.abc import AsyncIterator from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta @@ -58,6 +59,15 @@ async def connector() -> AsyncIterator[object]: yield client +def _deployment_verifier_tools() -> set[str]: + path = Path(__file__).resolve().parents[1] / "scripts" / "verify-deployment.py" + spec = importlib.util.spec_from_file_location("verify_deployment_tools", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return set(module.MCP_TOOLS) + + async def test_http_mcp_uses_2026_07_28_and_exposes_sandbox_toolset(tmp_path: Path) -> None: app = create_app(_settings(tmp_path)) mcp_server = app.state.mcp_server @@ -77,8 +87,13 @@ async def test_http_mcp_uses_2026_07_28_and_exposes_sandbox_toolset(tmp_path: Pa "sandbox_cancel", "sandbox_write_file", "sandbox_read_file", + "sandbox_suspend", + "sandbox_resume", "sandbox_release", } + # The deployment verifier pins the same toolset and fails a real + # deployment on any difference, so the two lists move together. + assert {tool.name for tool in tools.tools} == _deployment_verifier_tools() result = await client.call_tool("sandbox_profile", {}) assert result.is_error is False diff --git a/tests/test_polyglot_launchers.py b/tests/test_polyglot_launchers.py index 2938b6a..b402cd2 100644 --- a/tests/test_polyglot_launchers.py +++ b/tests/test_polyglot_launchers.py @@ -25,6 +25,13 @@ def executable(path: Path, body: str) -> None: path.chmod(0o755) +# A compiler older than the one that links with lld by default does not know the +# opt-out, and says so with a failing exit. The launcher tests below that are not +# about the linker use tools like that, so the arguments they print are only the +# ones the caller passed. +OLDER_COMPILER = 'case "$*" in *linker-features*) exit 1 ;; esac\n' + + def test_java_launcher_preserves_arguments_and_scopes_library_path( tmp_path: Path, installer: ModuleType ) -> None: @@ -59,7 +66,7 @@ def test_rust_launcher_supplies_sysroot_unless_explicit( ) -> None: home = tmp_path / "rust with spaces" for name in ("rustc", "cargo", "rustdoc"): - executable(home / "bin" / name, 'printf "%s\\n" "$LD_LIBRARY_PATH" "$@"\n') + executable(home / "bin" / name, OLDER_COMPILER + 'printf "%s\\n" "$LD_LIBRARY_PATH" "$@"\n') destination = tmp_path / "rust launchers" installer.install_rust(home, destination) result = subprocess.run( @@ -76,7 +83,7 @@ def test_rust_launcher_supplies_sysroot_unless_explicit( def test_cargo_invokes_adapted_compiler_and_rustdoc(tmp_path: Path, installer: ModuleType) -> None: home = tmp_path / "rust" for name in ("rustc", "rustdoc"): - executable(home / "bin" / name, 'printf "%s\\n" "$@"\n') + executable(home / "bin" / name, OLDER_COMPILER + 'printf "%s\\n" "$@"\n') executable(home / "bin" / "cargo", '"$RUSTC" --version\n"$RUSTDOC" --version\n') destination = tmp_path / "launchers" installer.install_rust(home, destination) @@ -86,6 +93,37 @@ def test_cargo_invokes_adapted_compiler_and_rustdoc(tmp_path: Path, installer: M assert result.stdout.splitlines() == ["--sysroot", str(home), "--version"] * 2 +@pytest.mark.parametrize("tool", ["rustc", "rustdoc"]) +def test_rust_is_linked_with_the_system_linker_when_the_compiler_can_be_told_to( + tmp_path: Path, installer: ModuleType, tool: str +) -> None: + """Rust 1.90's default linker finds itself through /proc/self/exe. + + The `basic` level has no procfs, so the launcher must opt out of it — for + rustdoc as well, which links every doctest itself. A real compiler would + refuse an option it does not know; the fake one here accepts everything, so + the probe finds the opt-out supported.""" + home = tmp_path / "rust" + for name in ("rustc", "rustdoc", "cargo"): + executable(home / "bin" / name, 'printf "%s\\n" "$@"\n') + destination = tmp_path / "launchers" + installer.install_rust(home, destination) + + result = subprocess.run( + [str(destination / tool), "input.rs"], capture_output=True, text=True, check=True + ) + + assert result.stdout.splitlines() == [ + "--sysroot", + str(home), + "-C", + "linker-features=-lld", + "input.rs", + ] + cargo = (destination / "cargo").read_text() + assert "linker-features" not in cargo # cargo reaches the compiler through RUSTC + + def test_rust_override_is_refused_explicitly(tmp_path: Path, installer: ModuleType) -> None: home = tmp_path / "rust" for name in ("rustc", "rustdoc", "cargo"): diff --git a/tests/test_preflight.py b/tests/test_preflight.py index 4fa7e56..fe61ee6 100644 --- a/tests/test_preflight.py +++ b/tests/test_preflight.py @@ -6,13 +6,35 @@ import pytest +from agent_sandbox import preflight from agent_sandbox.config import Settings from agent_sandbox.preflight import ( + check_user_namespace_support, toolchain_isolation_warnings, validate_runtime_environment, ) +@pytest.fixture(autouse=True) +def host_allows_user_namespaces( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Make the preflight tests independent of the kernel they run on. + + `validate_runtime_environment` also reports the host's user-namespace + settings, so on a machine with `kernel.apparmor_restrict_unprivileged_userns=1` + (GitHub's hosted runners, Ubuntu 24.04) every exact-list assertion below + gained a warning that has nothing to do with what it tests. The probe has its + own tests, which opt out of this fixture. + """ + if request.node.get_closest_marker("real_host_probe"): + return + sysctl = tmp_path / "apparmor_restrict_unprivileged_userns" + sysctl.write_text("0\n", encoding="utf-8") + monkeypatch.setattr(preflight, "_APPARMOR_USERNS_SYSCTL", sysctl) + monkeypatch.setattr(preflight, "_has_bounding_capability", lambda _capability: True) + + def _executable(path: Path) -> Path: path.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8") path.chmod(0o755) @@ -264,3 +286,31 @@ def test_an_empty_token_in_production_stops_startup( with pytest.raises(RuntimeError, match="SANDBOX_INTERNAL_TOKEN is not configured"): validate_runtime_environment(settings) + + +@pytest.mark.real_host_probe +@pytest.mark.parametrize(("content", "warns"), [("1\n", True), ("0\n", False), (None, False)]) +def test_the_apparmor_userns_restriction_is_reported_only_when_set( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, content: str | None, warns: bool +) -> None: + sysctl = tmp_path / "apparmor_restrict_unprivileged_userns" + if content is not None: + sysctl.write_text(content, encoding="utf-8") + monkeypatch.setattr(preflight, "_APPARMOR_USERNS_SYSCTL", sysctl) + monkeypatch.setattr(preflight, "_has_bounding_capability", lambda _capability: True) + monkeypatch.setattr(preflight.sys, "platform", "linux") + + found = check_user_namespace_support() + + assert bool(found) is warns + if warns: + assert "apparmor_restrict_unprivileged_userns=1" in found[0] + + +@pytest.mark.real_host_probe +def test_a_root_caller_without_cap_setfcap_is_reported(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(preflight.sys, "platform", "linux") + monkeypatch.setattr(preflight.os, "geteuid", lambda: 0) + monkeypatch.setattr(preflight, "_has_bounding_capability", lambda _capability: False) + + assert any("CAP_SETFCAP" in warning for warning in check_user_namespace_support()) diff --git a/tests/test_suspend_resume.py b/tests/test_suspend_resume.py new file mode 100644 index 0000000..0d08f2f --- /dev/null +++ b/tests/test_suspend_resume.py @@ -0,0 +1,1230 @@ +"""Suspend and resume: release a sandbox's slot, keep its workspace, wake it later. + +Three layers, each tested where its guarantee lives: the metadata store's +compare-and-set transitions (which make suspend and exec admission mutually +exclusive), the worker runtime (slot, directory, snapshot, disk eviction), and +the control plane's choice of where a resume lands. The end-to-end tests at the +bottom drive the HTTP API through two in-process workers. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import os +import shutil +import time +from dataclasses import replace +from datetime import timedelta +from pathlib import Path +from typing import Any + +import httpx +import pytest +from sqlalchemy import update + +from agent_sandbox.config import Settings +from agent_sandbox.dormant import manifest_key, read_manifest +from agent_sandbox.models import Route, utc_now_naive +from agent_sandbox.runtime import LocalSandbox, SandboxRuntime +from agent_sandbox.schemas import ExecRequest, ExecResponse +from agent_sandbox.service import SandboxService +from agent_sandbox.sql_database import SqlAlchemyDatabase, route_table + + +@pytest.fixture(autouse=True) +def _allow_chown_without_root(monkeypatch: pytest.MonkeyPatch) -> None: + """Workers chown to the sandbox UID, which needs root on a real worker.""" + monkeypatch.setattr(os, "chown", lambda *_args, **_kwargs: None) + + +class DirectoryStore: + """An object store on a local directory, with the template store's contract.""" + + def __init__(self, root: Path) -> None: + self.root = root + self.root.mkdir(parents=True, exist_ok=True) + + def _path(self, key: str) -> Path: + return self.root / key + + def upload_file(self, key: str, path: str | Path) -> str: + target = self._path(key) + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(path, target) + return f"blob://{key}" + + def download_to(self, key_or_uri: str, destination: str | Path) -> None: + source = self._path(key_or_uri.removeprefix("blob://")) + if not source.exists(): + raise FileNotFoundError(key_or_uri) + shutil.copyfile(source, destination) + + def delete(self, key_or_uri: str) -> None: + self._path(key_or_uri.removeprefix("blob://")).unlink(missing_ok=True) + + def keys(self) -> list[str]: + return sorted( + str(path.relative_to(self.root)) for path in self.root.rglob("*") if path.is_file() + ) + + +# ── metadata store transitions ── + + +async def _store(tmp_path: Path, **overrides: Any) -> tuple[SqlAlchemyDatabase, Route]: + pytest.importorskip("aiosqlite") + settings = Settings( + internal_token="token", + local_root=tmp_path / "sandboxes", + database_url=f"sqlite+aiosqlite:///{tmp_path / 'control.db'}", + database_auto_ddl=True, + profile_hash="profile-a", + **overrides, + ) + database = SqlAlchemyDatabase(settings) + await database.connect() + route = await database.create_route( + sandbox_id="sandbox-a", + workspace_scope_id="tenant-a", + worker={"worker_id": "worker-a", "worker_epoch": "epoch-a"}, + ) + await database.mark_route_ready(route.sandbox_id, route.generation, "worker-a", "epoch-a") + current = await database.find_route(route.sandbox_id) + assert current is not None + return database, current + + +async def _begin_exec(database: SqlAlchemyDatabase, route: Route, exec_id: str) -> Any: + return await database.begin_exec( + sandbox_id=route.sandbox_id, + exec_id=exec_id, + generation=route.generation, + worker_id="worker-a", + command=json.dumps({"argv": ["/bin/true"], "exec_scope": "thread-a"}), + exec_scope="thread-a", + ) + + +async def test_suspend_is_a_claim_then_a_finish_and_repeats_as_a_no_op(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + assert await database.begin_suspend(route.sandbox_id, route.generation) is True + # A retry of a suspend that stalled midway is told to finish it. + assert await database.begin_suspend(route.sandbox_id, route.generation) is True + await database.finish_suspend(route.sandbox_id, route.generation) + suspended = await database.find_route(route.sandbox_id) + assert suspended is not None + assert suspended.status == "SUSPENDED" + assert suspended.generation == route.generation + # Already suspended: nothing to do, and nothing is done. + assert await database.begin_suspend(route.sandbox_id, route.generation) is False + with pytest.raises(RuntimeError, match="STALE_SANDBOX_GENERATION"): + await database.begin_suspend(route.sandbox_id, route.generation + 1) + finally: + await database.close() + + +async def test_a_running_command_refuses_the_suspend_and_a_suspend_refuses_commands( + tmp_path: Path, +) -> None: + database, route = await _store(tmp_path) + try: + assert await _begin_exec(database, route, "exec-1") is None + with pytest.raises(RuntimeError, match="SANDBOX_SUSPEND_BUSY"): + await database.begin_suspend(route.sandbox_id, route.generation) + await database.finish_exec( + sandbox_id=route.sandbox_id, + exec_id="exec-1", + status="SUCCEEDED", + exit_code=0, + stdout="", + stderr="", + truncated=False, + ) + assert await database.begin_suspend(route.sandbox_id, route.generation) is True + # Claimed: admission is refused with a code that says why. + with pytest.raises(RuntimeError, match="SANDBOX_SUSPENDED"): + await _begin_exec(database, route, "exec-2") + await database.finish_suspend(route.sandbox_id, route.generation) + with pytest.raises(RuntimeError, match="SANDBOX_SUSPENDED"): + await _begin_exec(database, route, "exec-3") + # A replayed exec_id still answers from its record, suspended or not. + replay = await _begin_exec(database, route, "exec-1") + assert replay is not None and replay["status"] == "SUCCEEDED" + finally: + await database.close() + + +async def test_exec_admission_and_suspend_race_to_exactly_one_winner(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + outcomes = await asyncio.gather( + _begin_exec(database, route, "exec-race"), + database.begin_suspend(route.sandbox_id, route.generation), + return_exceptions=True, + ) + exec_won = outcomes[0] is None + suspend_won = outcomes[1] is True + assert exec_won != suspend_won, outcomes + current = await database.find_route(route.sandbox_id) + assert current is not None + assert current.status == ("RUNNING" if exec_won else "SUSPENDING") + finally: + await database.close() + + +async def test_abort_suspend_returns_the_route_to_ready(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + assert await database.begin_suspend(route.sandbox_id, route.generation) + await database.abort_suspend(route.sandbox_id, route.generation) + current = await database.find_route(route.sandbox_id) + assert current is not None and current.status == "READY" + assert await _begin_exec(database, route, "exec-after-abort") is None + finally: + await database.close() + + +async def test_resume_in_place_keeps_the_generation_and_only_one_resume_wins( + tmp_path: Path, +) -> None: + database, route = await _store(tmp_path) + worker = {"worker_id": "worker-a", "worker_epoch": "epoch-a"} + try: + await database.begin_suspend(route.sandbox_id, route.generation) + await database.finish_suspend(route.sandbox_id, route.generation) + suspended = await database.find_route(route.sandbox_id) + assert suspended is not None + first, second = await asyncio.gather( + database.begin_resume( + suspended, worker, bump_generation=False, profile_hash="profile-a", created_by="a" + ), + database.begin_resume( + suspended, worker, bump_generation=False, profile_hash="profile-a", created_by="b" + ), + ) + winners = [item for item in (first, second) if item is not None] + assert len(winners) == 1 + assert winners[0].status == "ASSIGNED" + assert winners[0].generation == route.generation + await database.mark_route_ready(route.sandbox_id, route.generation, "worker-a", "epoch-a") + assert await _begin_exec(database, route, "exec-after-resume") is None + finally: + await database.close() + + +async def test_resume_elsewhere_bumps_the_generation_and_audits_it(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + await database.begin_suspend(route.sandbox_id, route.generation) + await database.finish_suspend(route.sandbox_id, route.generation) + suspended = await database.find_route(route.sandbox_id) + assert suspended is not None + moved = await database.begin_resume( + suspended, + {"worker_id": "worker-b", "worker_epoch": "epoch-b"}, + bump_generation=True, + profile_hash="profile-a", + created_by="workspace:tenant-a", + ) + assert moved is not None + assert (moved.worker_id, moved.generation, moved.status) == ( + "worker-b", + route.generation + 1, + "ASSIGNED", + ) + assert moved.last_release_reason == "RESUME_REASSIGNED" + # The old generation is fenced: nothing can run against it any more. + with pytest.raises(RuntimeError, match="STALE_SANDBOX_GENERATION"): + await _begin_exec(database, route, "exec-old-generation") + # A failed resume goes back to sleep where a retry finds it. + await database.abort_resume(moved.sandbox_id, moved.generation) + current = await database.find_route(route.sandbox_id) + assert current is not None and current.status == "SUSPENDED" + finally: + await database.close() + + +async def test_the_retention_clock_is_not_reset_and_expiry_is_listed(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + await database.begin_suspend(route.sandbox_id, route.generation) + await database.finish_suspend(route.sandbox_id, route.generation) + async with database._engine().begin() as connection: + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == route.sandbox_id) + .values(last_active_at=utc_now_naive() - timedelta(hours=2)) + ) + before = await database.find_route(route.sandbox_id) + await database.touch_route(route.sandbox_id, route.generation) + after = await database.find_route(route.sandbox_id) + assert before is not None and after is not None + assert after.last_active_at == before.last_active_at + + assert ( + await database.list_dormant_routes_to_reclaim( + retention_seconds=3 * 3600, suspending_grace_seconds=300, limit=10 + ) + == [] + ) + expired = await database.list_dormant_routes_to_reclaim( + retention_seconds=3600, suspending_grace_seconds=300, limit=10 + ) + assert [item.sandbox_id for item in expired] == [route.sandbox_id] + # Zero disables expiry. + assert ( + await database.list_dormant_routes_to_reclaim( + retention_seconds=0, suspending_grace_seconds=300, limit=10 + ) + == [] + ) + # A suspended route can still be released by its client. + assert await database.begin_release(route.sandbox_id, route.generation) is True + finally: + await database.close() + + +async def test_a_stalled_suspend_is_listed_for_completion(tmp_path: Path) -> None: + database, route = await _store(tmp_path) + try: + await database.begin_suspend(route.sandbox_id, route.generation) + async with database._engine().begin() as connection: + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == route.sandbox_id) + .values(updated_at=utc_now_naive() - timedelta(minutes=10)) + ) + stalled = await database.list_dormant_routes_to_reclaim( + retention_seconds=0, suspending_grace_seconds=300, limit=10 + ) + assert [(item.sandbox_id, item.status) for item in stalled] == [ + (route.sandbox_id, "SUSPENDING") + ] + finally: + await database.close() + + +# ── worker runtime ── + + +def _runtime(tmp_path: Path, name: str = "worker-a", **overrides: Any) -> SandboxRuntime: + settings = Settings( + internal_token="token", + local_root=tmp_path / name / "sandboxes", + template_root=tmp_path / name / "templates", + min_free_bytes=0, + **overrides, + ) + return SandboxRuntime(settings) + + +async def _live(runtime: SandboxRuntime, sandbox_id: str = "sb-1") -> LocalSandbox: + sandbox = await runtime.create(sandbox_id, 1, os.getuid()) + (sandbox.workspace / "notes.txt").write_text("state before the wait") + (sandbox.root / "envs" / "venv").mkdir() + (sandbox.root / "envs" / "venv" / "marker").write_text("installed") + return sandbox + + +async def test_suspend_releases_the_slot_and_keeps_the_directory(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + sandbox = await _live(runtime) + + result = await runtime.suspend("sb-1", 1) + + assert result == {"status": "SUSPENDED", "snapshot": False} + assert "sb-1" not in runtime.sandboxes # the slot the heartbeat counts + assert runtime.dormant["sb-1"].generation == 1 + assert (sandbox.workspace / "notes.txt").read_text() == "state before the wait" + # Repeating is a no-op. + assert await runtime.suspend("sb-1", 1) == {"status": "SUSPENDED", "snapshot": False} + with pytest.raises(RuntimeError, match="STALE_SANDBOX_GENERATION"): + runtime.get("sb-1", 1) + + +async def test_a_stale_handle_cannot_touch_a_suspended_workspace(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + sandbox = await _live(runtime) + await runtime.suspend("sb-1", 1) + + with pytest.raises(RuntimeError, match="SANDBOX_SUSPENDED"): + await runtime.write_file(sandbox, "/workspace/late.txt", base64.b64encode(b"x").decode()) + assert not (sandbox.workspace / "late.txt").exists() + + +async def test_suspend_refuses_while_a_command_holds_the_sandbox(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + sandbox = await _live(runtime) + other = await _live(runtime, "sb-2") + entered = asyncio.Event() + release = asyncio.Event() + + async def execute_locked(_sandbox: LocalSandbox, request: ExecRequest) -> ExecResponse: + entered.set() + await release.wait() + return ExecResponse(exec_id=request.exec_id, status="SUCCEEDED") + + runtime._execute_locked = execute_locked # type: ignore[method-assign] + running = asyncio.create_task( + runtime.execute( + sandbox, + ExecRequest(exec_id="exec-1", generation=1, argv=["true"], exec_scope="thread-a"), + ) + ) + await entered.wait() + try: + with pytest.raises(RuntimeError, match="SANDBOX_SUSPEND_BUSY"): + await runtime.suspend("sb-1", 1) + assert "sb-1" in runtime.sandboxes + # Another sandbox is not affected by this one being busy. + await runtime.suspend("sb-2", 1) + assert "sb-2" in runtime.dormant + finally: + release.set() + await running + assert other.root.exists() + await runtime.suspend("sb-1", 1) + assert "sb-1" in runtime.dormant + + +async def test_resume_on_the_same_worker_reuses_the_directory(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + sandbox = await _live(runtime) + await runtime.suspend("sb-1", 1) + + resumed, source = await runtime.resume("sb-1", 1, os.getuid()) + + assert source == "reused" + assert resumed.root == sandbox.root + assert "sb-1" in runtime.sandboxes and "sb-1" not in runtime.dormant + assert (resumed.workspace / "notes.txt").read_text() == "state before the wait" + # Idempotent: a second resume finds it awake. + again, source = await runtime.resume("sb-1", 1, os.getuid()) + assert (again, source) == (resumed, "active") + + +async def test_resume_without_a_directory_or_snapshot_reports_the_loss(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + sandbox = await _live(runtime) + await runtime.suspend("sb-1", 1) + shutil.rmtree(sandbox.root) + + with pytest.raises(RuntimeError, match="SANDBOX_WORKSPACE_LOST"): + await runtime.resume("sb-1", 2, os.getuid()) + assert "sb-1" not in runtime.sandboxes + + +async def test_a_snapshot_restores_the_workspace_on_another_worker(tmp_path: Path) -> None: + store = DirectoryStore(tmp_path / "store") + source_worker = _runtime(tmp_path, "worker-a") + source_worker.templates.object_store = store + sandbox = await _live(source_worker) + (sandbox.workspace / "link").symlink_to("/etc/hostname") + + assert await source_worker.suspend("sb-1", 1, snapshot=True) == { + "status": "SUSPENDED", + "snapshot": True, + } + manifest = read_manifest(store, "sb-1") + assert manifest is not None and set(manifest["parts"]) == {"workspace", "home", "envs"} + + target_worker = _runtime(tmp_path, "worker-b") + target_worker.templates.object_store = store + resumed, source = await target_worker.resume("sb-1", 2, os.getuid(), restore="snapshot") + + assert source == "restored" + assert (resumed.workspace / "notes.txt").read_text() == "state before the wait" + assert (resumed.root / "envs" / "venv" / "marker").read_text() == "installed" + # Links come back as links, never as the file they point at. + assert os.readlink(resumed.workspace / "link") == "/etc/hostname" + for name in ("cache", "logs"): + assert (resumed.root / name).is_dir() + + +async def test_a_new_snapshot_replaces_the_previous_one(tmp_path: Path) -> None: + store = DirectoryStore(tmp_path / "store") + runtime = _runtime(tmp_path) + runtime.templates.object_store = store + sandbox = await _live(runtime) + await runtime.suspend("sb-1", 1, snapshot=True) + first = read_manifest(store, "sb-1") + await runtime.resume("sb-1", 1, os.getuid()) + (sandbox.workspace / "notes.txt").write_text("second") + await runtime.suspend("sb-1", 1, snapshot=True) + second = read_manifest(store, "sb-1") + + assert first is not None and second is not None + assert first["nonce"] != second["nonce"] + assert not any(first["nonce"] in key for key in store.keys()) + + +async def test_snapshot_always_without_a_store_refuses_and_keeps_the_sandbox( + tmp_path: Path, +) -> None: + runtime = _runtime(tmp_path) + runtime.templates.object_store = None + await _live(runtime) + + with pytest.raises(RuntimeError, match="SANDBOX_SNAPSHOT_UNAVAILABLE"): + await runtime.suspend("sb-1", 1, snapshot=True) + assert "sb-1" in runtime.sandboxes and "sb-1" not in runtime.dormant + + +async def test_disk_pressure_evicts_only_snapshotted_dormant_copies( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + store = DirectoryStore(tmp_path / "store") + runtime = _runtime(tmp_path) + runtime.templates.object_store = store + backed_up = await _live(runtime, "sb-backed-up") + only_copy = await _live(runtime, "sb-only-copy") + await runtime.suspend("sb-backed-up", 1, snapshot=True) + await runtime.suspend("sb-only-copy", 1, snapshot=False) + pressure = {"available": False} + real_status = runtime.disk_status + + def disk_status() -> dict[str, int | bool]: + return {**real_status(), "available": pressure["available"]} + + monkeypatch.setattr(runtime, "disk_status", disk_status) + + assert await runtime.reclaim_dormant_disk() == ["sb-backed-up"] + assert not backed_up.root.exists() + assert only_copy.root.exists() + assert runtime.dormant["sb-backed-up"].evicted + + pressure["available"] = True + resumed, source = await runtime.resume("sb-backed-up", 1, os.getuid()) + assert source == "restored" + assert (resumed.workspace / "notes.txt").read_text() == "state before the wait" + + +# ── where a resume lands ── + + +def _route(**overrides: Any) -> Route: + base = Route( + sandbox_id="sb-1", + workspace_scope_id="tenant-a", + worker_id="worker-a", + worker_epoch="epoch-a", + generation=4, + sandbox_uid=20001, + profile_id="coding-default", + profile_hash="profile-a", + status="SUSPENDED", + storage_mode="local", + ) + return replace(base, **overrides) + + +def _worker(worker_id: str, epoch: str, **overrides: Any) -> dict[str, Any]: + return { + "worker_id": worker_id, + "worker_epoch": epoch, + "endpoint": f"http://{worker_id}:8080", + "status": "ACTIVE", + "profile_hash": "profile-a", + "capacity": 4, + "running_sessions": 1, + **overrides, + } + + +class _Registry: + def __init__(self, *workers: dict[str, Any]) -> None: + self.workers = {item["worker_id"]: item for item in workers} + + async def get(self, worker_id: str | None) -> dict[str, Any] | None: + return self.workers.get(worker_id or "") + + async def select(self, *, profile_hash: str, exclude: set[str] | None = None) -> dict[str, Any]: + for item in self.workers.values(): + if item["worker_id"] not in (exclude or set()): + return item + raise RuntimeError("NO_SANDBOX_WORKER_AVAILABLE") + + +def _planner(tmp_path: Path, registry: _Registry, *, store: Any = None) -> SandboxService: + service = SandboxService( + Settings(internal_token="token", profile_hash="profile-a", local_root=tmp_path), + object(), # type: ignore[arg-type] + registry, # type: ignore[arg-type] + ) + service.template_catalog.object_store = store + return service + + +async def test_a_live_original_worker_reuses_the_directory_in_place(tmp_path: Path) -> None: + service = _planner(tmp_path, _Registry(_worker("worker-a", "epoch-a"))) + worker, bump, restore = await service._resume_plan(_route()) + assert (worker["worker_id"], bump, restore) == ("worker-a", False, "reuse") + + +async def test_a_restarted_original_worker_reuses_with_a_new_generation(tmp_path: Path) -> None: + service = _planner(tmp_path, _Registry(_worker("worker-a", "epoch-restarted"))) + worker, bump, restore = await service._resume_plan(_route()) + assert (worker["worker_id"], bump, restore) == ("worker-a", True, "reuse") + + +async def test_a_shared_workspace_moves_to_any_worker(tmp_path: Path) -> None: + service = _planner(tmp_path, _Registry(_worker("worker-b", "epoch-b"))) + worker, bump, restore = await service._resume_plan(_route(storage_mode="shared")) + assert (worker["worker_id"], bump, restore) == ("worker-b", True, "reuse") + + +async def test_a_local_workspace_moves_only_with_a_snapshot(tmp_path: Path) -> None: + store = DirectoryStore(tmp_path / "store") + registry = _Registry(_worker("worker-b", "epoch-b")) + service = _planner(tmp_path, registry, store=store) + + with pytest.raises(RuntimeError, match="SANDBOX_DORMANT_WORKSPACE_UNAVAILABLE"): + await service._resume_plan(_route()) + + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps({"sandbox_id": "sb-1", "nonce": "n", "parts": {}})) + store.upload_file(manifest_key("sb-1"), manifest) + worker, bump, restore = await service._resume_plan(_route()) + assert (worker["worker_id"], bump, restore) == ("worker-b", True, "snapshot") + + +async def test_a_full_original_worker_without_a_snapshot_asks_for_a_retry( + tmp_path: Path, +) -> None: + full = _worker("worker-a", "epoch-a", running_sessions=4) + service = _planner(tmp_path, _Registry(full)) + with pytest.raises(RuntimeError, match="NO_SANDBOX_WORKER_AVAILABLE"): + await service._resume_plan(_route()) + + +# ── end to end, over HTTP, across two workers ── + +TOKEN = "test-token" +AUTH = {"Authorization": f"Bearer {TOKEN}"} +_RealAsyncClient = httpx.AsyncClient + + +class _Fleet(httpx.AsyncBaseTransport): + """Routes worker-to-worker calls to in-process apps by host name.""" + + def __init__(self) -> None: + self.apps: dict[str, httpx.ASGITransport] = {} + + def add(self, host: str, app: Any) -> None: + self.apps[host] = httpx.ASGITransport(app=app) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + transport = self.apps.get(request.url.host) + if transport is None: + raise httpx.ConnectError("worker is down", request=request) + return await transport.handle_async_request(request) + + +async def _worker_app( + tmp_path: Path, name: str, fleet: _Fleet, store: DirectoryStore, registry: Any = None +) -> Any: + from agent_sandbox.app import create_app + + settings = Settings( + internal_token=TOKEN, + advertise_host=name, + port=8080, + local_root=tmp_path / name / "sandboxes", + template_root=tmp_path / name / "templates", + database_url=f"sqlite+aiosqlite:///{tmp_path / 'control.db'}", + database_auto_ddl=True, + min_free_bytes=0, + profile_hash="profile-a", + # One UID the test process owns, so a reused directory validates. + uid_start=os.getuid(), + uid_end=os.getuid(), + suspend_snapshot="always", + ) + app = create_app(settings) + service = app.state.sandbox_service + if registry is not None: + service.registry = registry + service.template_catalog.object_store = store + service.runtime.templates.object_store = store + settings.workspace_root.mkdir(parents=True, exist_ok=True) + await service.database.connect() + await service._heartbeat() + fleet.add(name, app) + return app + + +@pytest.fixture +async def fleet(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Any: + pytest.importorskip("aiosqlite") + transport = _Fleet() + + def client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + kwargs["transport"] = transport + return _RealAsyncClient(*args, **kwargs) + + monkeypatch.setattr(httpx, "AsyncClient", client) + store = DirectoryStore(tmp_path / "store") + app_a = await _worker_app(tmp_path, "worker-a", transport, store) + apps = [app_a] + yield transport, store, apps + for app in apps: + await app.state.sandbox_service.database.close() + + +async def _api(app: Any, method: str, path: str, **kwargs: Any) -> httpx.Response: + async with _RealAsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as client: + return await client.request(method, path, headers=AUTH, **kwargs) + + +async def _ready_sandbox(app: Any) -> int: + resolved = await _api( + app, + "POST", + "/api/v1/sandboxes/resolve", + json={"sandbox_id": "agent-1", "workspace_scope_id": "tenant-a"}, + ) + assert resolved.status_code == 200, resolved.text + generation = int(resolved.json()["generation"]) + created = await _api(app, "POST", "/api/v1/sandboxes/agent-1", json={"generation": generation}) + assert created.status_code == 200, created.text + written = await _api( + app, + "PUT", + "/api/v1/sandboxes/agent-1/files", + json={ + "generation": generation, + "path": "/workspace/progress.txt", + "content_base64": base64.b64encode(b"step 41 of 90").decode(), + }, + ) + assert written.status_code == 200, written.text + # What the heartbeat loop would publish within its interval. + await app.state.sandbox_service._heartbeat() + return generation + + +async def _read(app: Any, generation: int) -> httpx.Response: + return await _api( + app, + "GET", + "/api/v1/sandboxes/agent-1/files", + params={"path": "/workspace/progress.txt", "generation": generation}, + ) + + +async def _running_sessions(app: Any) -> int: + service = app.state.sandbox_service + worker = await service.registry.get(service.worker_id) + return int(worker["running_sessions"]) + + +async def test_suspend_then_resume_in_place_over_http(fleet: Any) -> None: + _, store, (app,) = fleet + generation = await _ready_sandbox(app) + assert await _running_sessions(app) == 1 + + suspended = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + assert suspended.status_code == 200, suspended.text + body = suspended.json() + assert (body["status"], body["snapshot"], body["suspended"]) == ("SUSPENDED", True, True) + assert body["retained_until"] is not None + # The slot is back, and published at once. + assert await _running_sessions(app) == 0 + + exec_response = await _api( + app, + "POST", + "/api/v1/sandboxes/agent-1/exec", + json={"exec_id": "exec-1", "generation": generation, "argv": ["true"]}, + ) + assert (exec_response.status_code, exec_response.text) == (409, "SANDBOX_SUSPENDED") + assert (await _read(app, generation)).status_code == 409 + + again = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + assert again.status_code == 200 and again.json()["suspended"] is False + + resumed = await _api(app, "POST", "/api/v1/sandboxes/agent-1/resume") + assert resumed.status_code == 200, resumed.text + body = resumed.json() + assert (body["status"], body["generation"], body["resumed"], body["workspace_source"]) == ( + "READY", + generation, + True, + "reused", + ) + assert await _running_sessions(app) == 1 + # A snapshot of a sandbox that is running again is stale; it is removed. + assert store.keys() == [] + content = await _read(app, generation) + assert content.status_code == 200, content.text + assert base64.b64decode(content.json()["content_base64"]) == b"step 41 of 90" + + repeat = await _api(app, "POST", "/api/v1/sandboxes/agent-1/resume") + assert repeat.status_code == 200 + assert (repeat.json()["resumed"], repeat.json()["workspace_source"]) == (False, "active") + + +async def test_resolve_wakes_a_suspended_sandbox(fleet: Any) -> None: + _, _, (app,) = fleet + generation = await _ready_sandbox(app) + await _api(app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + + resolved = await _api( + app, + "POST", + "/api/v1/sandboxes/resolve", + json={"sandbox_id": "agent-1", "workspace_scope_id": "tenant-a"}, + ) + + assert resolved.status_code == 200, resolved.text + assert (resolved.json()["status"], resolved.json()["generation"]) == ("READY", generation) + assert (await _read(app, generation)).status_code == 200 + + +async def test_a_busy_sandbox_refuses_to_suspend_over_http(fleet: Any) -> None: + _, _, (app,) = fleet + generation = await _ready_sandbox(app) + database = app.state.sandbox_service.database + await database.begin_exec( + sandbox_id="agent-1", + exec_id="exec-long", + generation=generation, + worker_id="worker-a-8080", + command=json.dumps({"argv": ["sleep", "600"], "exec_scope": "thread-a"}), + exec_scope="thread-a", + ) + + response = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + + assert (response.status_code, response.text) == (409, "SANDBOX_SUSPEND_BUSY") + route = await database.find_route("agent-1") + assert route is not None and route.status == "RUNNING" + assert await _running_sessions(app) == 1 + + +async def test_resume_moves_to_another_worker_from_the_snapshot( + fleet: Any, tmp_path: Path +) -> None: + transport, store, apps = fleet + app_a = apps[0] + generation = await _ready_sandbox(app_a) + suspended = await _api( + app_a, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + assert suspended.json()["snapshot"] is True + + registry = app_a.state.sandbox_service.registry + app_b = await _worker_app(tmp_path, "worker-b", transport, store, registry=registry) + apps.append(app_b) + # Worker A goes away with the only local copy. + await registry.unregister(app_a.state.sandbox_service.worker_id) + del transport.apps["worker-a"] + + resumed = await _api(app_b, "POST", "/api/v1/sandboxes/agent-1/resume") + + assert resumed.status_code == 200, resumed.text + body = resumed.json() + assert (body["worker_id"], body["generation"], body["workspace_source"]) == ( + "worker-b-8080", + generation + 1, + "restored", + ) + content = await _read(app_b, generation + 1) + assert content.status_code == 200, content.text + assert base64.b64decode(content.json()["content_base64"]) == b"step 41 of 90" + # The old generation is fenced. + assert (await _read(app_b, generation)).status_code == 409 + + +async def test_resume_without_a_copy_anywhere_is_retryable_not_destructive( + fleet: Any, tmp_path: Path +) -> None: + transport, store, apps = fleet + app_a = apps[0] + service_a = app_a.state.sandbox_service + service_a.settings.suspend_snapshot = "never" + generation = await _ready_sandbox(app_a) + await _api(app_a, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + app_b = await _worker_app(tmp_path, "worker-b", transport, store, registry=service_a.registry) + apps.append(app_b) + await service_a.registry.unregister(service_a.worker_id) + + response = await _api(app_b, "POST", "/api/v1/sandboxes/agent-1/resume") + + assert (response.status_code, response.text) == (503, "SANDBOX_DORMANT_WORKSPACE_UNAVAILABLE") + route = await service_a.database.find_route("agent-1") + assert route is not None + assert (route.status, route.worker_id, route.generation) == ( + "SUSPENDED", + service_a.worker_id, + generation, + ) + + +async def test_retention_expiry_releases_and_reclaims_the_disk(fleet: Any) -> None: + _, store, (app,) = fleet + service = app.state.sandbox_service + generation = await _ready_sandbox(app) + await _api(app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + root = service.settings.workspace_root / "agent-1" + assert root.exists() and store.keys() + async with service.database._engine().begin() as connection: + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == "agent-1") + .values( + last_active_at=utc_now_naive() + - timedelta(seconds=service.settings.suspended_retention_seconds + 60) + ) + ) + + await service._maintenance() + + route = await service.database.find_route("agent-1") + assert route is not None + assert (route.status, route.last_release_reason, route.last_released_by) == ( + "RELEASED", + "SUSPEND_EXPIRED", + "system:orphan-reaper", + ) + assert not root.exists() + assert "agent-1" not in service.runtime.dormant + assert store.keys() == [] + # Counted by the same reaper counters as an idle-timeout release. + assert service.reaper_status["last_released"] == 1 + assert service.reaper_status["released_total"] == 1 + + +# ── workspace directories left behind on a worker that no longer owns them ── + + +def _age(marker: Path, seconds: float) -> None: + then = time.time() - seconds + os.utime(marker, (then, then)) + + +async def test_an_orphaned_directory_is_deleted_only_after_its_ttl(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + active = await _live(runtime, "sb-active") + stale = await _live(runtime, "sb-stale") + # What a resume elsewhere leaves: a directory no sandbox here holds. + runtime.sandboxes.pop("sb-stale") + + # A live sandbox is never even listed. + assert await runtime.local_workspace_ids() == ["sb-stale"] + + # The first sighting starts the clock and deletes nothing. + assert await runtime.reclaim_orphan_workspaces(["sb-stale"], ttl_seconds=3600) == [] + marker = runtime._orphan_marker_root() / "sb-stale" + assert marker.exists() and stale.root.exists() + _age(marker, 3000) + assert await runtime.reclaim_orphan_workspaces(["sb-stale"], ttl_seconds=3600) == [] + assert stale.root.exists() + + # A restarted worker reads the clock from disk instead of starting over. + restarted = _runtime(tmp_path) + _age(marker, 3601) + assert await restarted.reclaim_orphan_workspaces(["sb-stale"], ttl_seconds=3600) == [ + "sb-stale" + ] + assert not stale.root.exists() + assert not marker.exists() + assert (active.workspace / "notes.txt").exists() + + +async def test_a_live_sandbox_is_never_deleted_even_if_reported_orphaned( + tmp_path: Path, +) -> None: + runtime = _runtime(tmp_path) + active = await _live(runtime, "sb-active") + markers = runtime._orphan_marker_root() + markers.mkdir(parents=True, exist_ok=True) + (markers / "sb-active").touch() + _age(markers / "sb-active", 10_000) + + assert await runtime.reclaim_orphan_workspaces(["sb-active"], ttl_seconds=60) == [] + assert active.root.exists() + assert list(markers.iterdir()) == [] + + +async def test_a_dormant_entry_whose_route_was_released_elsewhere_is_reclaimed( + tmp_path: Path, +) -> None: + """The route decides, not memory: a release that could not reach this + worker leaves a dormant entry here that nothing else would ever clear.""" + runtime = _runtime(tmp_path) + dormant = await _live(runtime, "sb-dormant") + await runtime.suspend("sb-dormant", 1) + + assert await runtime.local_workspace_ids() == ["sb-dormant"] + await runtime.reclaim_orphan_workspaces(["sb-dormant"], ttl_seconds=60) + _age(runtime._orphan_marker_root() / "sb-dormant", 61) + + assert await runtime.reclaim_orphan_workspaces(["sb-dormant"], ttl_seconds=60) == [ + "sb-dormant" + ] + assert not dormant.root.exists() + assert "sb-dormant" not in runtime.dormant + + +async def test_an_orphan_that_is_owned_again_loses_its_clock(tmp_path: Path) -> None: + runtime = _runtime(tmp_path) + stale = await _live(runtime, "sb-stale") + runtime.sandboxes.pop("sb-stale") + await runtime.reclaim_orphan_workspaces(["sb-stale"], ttl_seconds=60) + marker = runtime._orphan_marker_root() / "sb-stale" + _age(marker, 120) + + # The control plane stopped reporting it: its route points here again. + assert await runtime.reclaim_orphan_workspaces([], ttl_seconds=60) == [] + assert not marker.exists() + # Reported again later, the clock starts from zero. + assert await runtime.reclaim_orphan_workspaces(["sb-stale"], ttl_seconds=60) == [] + assert stale.root.exists() + + +async def test_the_old_worker_reclaims_the_copy_a_cross_worker_resume_left( + fleet: Any, tmp_path: Path +) -> None: + transport, store, apps = fleet + app_a = apps[0] + service_a = app_a.state.sandbox_service + service_a.settings.orphan_dormant_dir_ttl_seconds = 60 + generation = await _ready_sandbox(app_a) + await _api(app_a, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + old_root = service_a.settings.workspace_root / "agent-1" + markers = service_a.runtime._orphan_marker_root() + + # Suspended here: listed (memory does not protect a dormant entry), but + # its route names worker A, so no clock starts. + await service_a._maintenance() + assert old_root.exists() + assert not (markers / "agent-1").exists() + + # Suspended here, and worker A restarted (its memory of the dormant + # sandbox is gone): the route still names worker A, so it is kept. + service_a.runtime.dormant.clear() + await service_a._maintenance() + assert old_root.exists() + assert not (markers / "agent-1").exists() + + registry = service_a.registry + app_b = await _worker_app(tmp_path, "worker-b", transport, store, registry=registry) + apps.append(app_b) + service_b = app_b.state.sandbox_service + service_b.settings.orphan_dormant_dir_ttl_seconds = 60 + await registry.unregister(service_a.worker_id) + resumed = await _api(app_b, "POST", "/api/v1/sandboxes/agent-1/resume") + assert resumed.json()["workspace_source"] == "restored" + + # Worker A sees the copy is no longer its own and starts the clock. + await service_a._maintenance() + assert old_root.exists() + assert (markers / "agent-1").exists() + _age(markers / "agent-1", 61) + await service_a._maintenance() + + assert not old_root.exists() + assert service_a.reaper_status["orphan_workspaces_deleted_total"] == 1 + # The active copy on worker B is its own and stays. + await service_b._maintenance() + assert not (service_b.runtime._orphan_marker_root() / "agent-1").exists() + content = await _read(app_b, generation + 1) + assert base64.b64decode(content.json()["content_base64"]) == b"step 41 of 90" + + +async def test_a_ttl_of_zero_keeps_orphaned_directories(fleet: Any) -> None: + _, _, (app,) = fleet + service = app.state.sandbox_service + service.settings.orphan_dormant_dir_ttl_seconds = 0 + service.runtime.local_workspace_ids = None # would fail if it were consulted + + assert await service._reclaim_orphan_workspaces() == 0 + + +async def test_an_expiry_that_fails_midway_is_finished_by_the_normal_retry(fleet: Any) -> None: + """Same release path as an idle timeout: claim, worker delete, audit, retry. + + The first attempt cannot reach the worker and leaves the route RELEASING; + the reaper's ordinary RELEASE_RETRY finishes it, and that retry has to + delete the snapshot too, or it is left in the object store for good. + """ + transport, store, (app,) = fleet + service = app.state.sandbox_service + generation = await _ready_sandbox(app) + await _api(app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + root = service.settings.workspace_root / "agent-1" + assert root.exists() and store.keys() + async with service.database._engine().begin() as connection: + await connection.execute( + update(route_table) + .where(route_table.c.sandbox_id == "agent-1") + .values( + last_active_at=utc_now_naive() + - timedelta(seconds=service.settings.suspended_retention_seconds + 60) + ) + ) + worker_app = transport.apps.pop("worker-a") + + await service._maintenance() + + route = await service.database.find_route("agent-1") + assert route is not None and route.status == "RELEASING" + assert service.reaper_status["failures_total"] == 1 + assert root.exists() and store.keys() + + transport.apps["worker-a"] = worker_app + service.settings.orphan_release_grace_seconds = 0 + await service._maintenance() + + route = await service.database.find_route("agent-1") + assert route is not None + assert (route.status, route.last_release_reason, route.last_released_by) == ( + "RELEASED", + "RELEASE_RETRY", + "system:orphan-reaper", + ) + assert not root.exists() + assert store.keys() == [] + assert service.reaper_status["released_total"] == 1 + + +# ── adapters written before suspend existed ── + + +class _WithoutDormantRoutes: + """A metadata store that predates `DormantRouteStore`: everything else delegates.""" + + _HIDDEN = frozenset( + { + "begin_suspend", + "finish_suspend", + "abort_suspend", + "begin_resume", + "abort_resume", + "list_dormant_routes_to_reclaim", + } + ) + + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __getattr__(self, name: str) -> Any: + if name in self._HIDDEN: + raise AttributeError(name) + return getattr(self._inner, name) + + +async def test_a_metadata_store_without_the_protocol_answers_501_and_keeps_working( + fleet: Any, +) -> None: + _, _, (app,) = fleet + service = app.state.sandbox_service + generation = await _ready_sandbox(app) + service.database = _WithoutDormantRoutes(service.database) + + suspended = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + resumed = await _api(app, "POST", "/api/v1/sandboxes/agent-1/resume") + + assert (suspended.status_code, suspended.text) == (501, "SANDBOX_SUSPEND_UNSUPPORTED") + assert (resumed.status_code, resumed.text) == (501, "SANDBOX_SUSPEND_UNSUPPORTED") + assert (await _read(app, generation)).status_code == 200 + # The maintenance cycle skips the dormant sweep instead of failing. + await service._maintenance() + assert service.reaper_status["failures_total"] == 0 + + +async def test_an_execution_backend_without_the_protocol_answers_501_and_stays_ready( + fleet: Any, +) -> None: + _, _, (app,) = fleet + service = app.state.sandbox_service + generation = await _ready_sandbox(app) + service.runtime.suspend = None + service.runtime.resume = None + service.runtime.reclaim_dormant_disk = None + + suspended = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + + assert (suspended.status_code, suspended.text) == (501, "SANDBOX_SUSPEND_UNSUPPORTED") + route = await service.database.find_route("agent-1") + assert route is not None and route.status == "READY" + assert (await _read(app, generation)).status_code == 200 + await service._maintenance() + assert service.reaper_status["failures_total"] == 0 + + +# ── disk pressure after a restart ── + + +async def test_a_restarted_worker_still_evicts_its_snapshotted_dormant_copies( + fleet: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + _, _, (app,) = fleet + service = app.state.sandbox_service + generation = await _ready_sandbox(app) + suspended = await _api( + app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation} + ) + assert suspended.json()["snapshot"] is True + root = service.settings.workspace_root / "agent-1" + # A restart: the directory stays, the worker's memory of it does not. + service.runtime.dormant.clear() + pressure = {"available": True} + real_status = service.runtime.disk_status + monkeypatch.setattr( + service.runtime, + "disk_status", + lambda: {**real_status(), "available": pressure["available"]}, + ) + + await service._maintenance() + assert root.exists() and "agent-1" not in service.runtime.dormant # no pressure: nothing to do + + pressure["available"] = False + await service._maintenance() + assert not root.exists() + assert service.reaper_status["dormant_evicted_total"] == 1 + route = await service.database.find_route("agent-1") + assert route is not None and route.status == "SUSPENDED" + + pressure["available"] = True + resumed = await _api(app, "POST", "/api/v1/sandboxes/agent-1/resume") + assert resumed.json()["workspace_source"] == "restored" + assert (await _read(app, resumed.json()["generation"])).status_code == 200 + + +async def test_a_restarted_worker_keeps_a_dormant_copy_that_has_no_snapshot( + fleet: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + _, _, (app,) = fleet + service = app.state.sandbox_service + service.settings.suspend_snapshot = "never" + generation = await _ready_sandbox(app) + await _api(app, "POST", "/api/v1/sandboxes/agent-1/suspend", json={"generation": generation}) + root = service.settings.workspace_root / "agent-1" + service.runtime.dormant.clear() + real_status = service.runtime.disk_status + monkeypatch.setattr( + service.runtime, "disk_status", lambda: {**real_status(), "available": False} + ) + + await service._maintenance() + + assert root.exists() + assert "agent-1" not in service.runtime.dormant diff --git a/tests/test_template_security.py b/tests/test_template_security.py index ae6d0dc..343521f 100644 --- a/tests/test_template_security.py +++ b/tests/test_template_security.py @@ -80,6 +80,7 @@ def replace_before_open(path: Any, flags: int, *args: Any, **kwargs: Any) -> int if path == "data" and kwargs.get("dir_fd") is not None and not changed: changed = True target = source / "data" + original = target.stat() if not target.is_dir() else None if target.is_dir(): target.rename(source / "old") target.symlink_to(tmp_path, target_is_directory=True) @@ -89,6 +90,17 @@ def replace_before_open(path: Any, flags: int, *args: Any, **kwargs: Any) -> int os.mkfifo(target) elif kind == "file-replacement": target.write_text("host secret") + # The replacement is the same size as the original and a + # filesystem hands a just-freed inode straight back, so + # without this the two files can agree on inode, size and + # (the clock ticks in milliseconds) both timestamps, and + # nothing could tell them apart. A real replacement is + # written at another time; say so, so the test is about the + # check and not about the timer. + assert original is not None + os.utime( + target, ns=(original.st_atime_ns, original.st_mtime_ns + 1_000_000_000) + ) else: target.symlink_to(host) return original_open(path, flags, *args, **kwargs)