diff --git a/docs/architecture.md b/docs/architecture.md index fec9688ca..634411541 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -99,13 +99,13 @@ Every lifecycle the framework owns, as ordered phases. **Rollout lifecycle.** `setup` (resolve config, build the environment object) → `start` (sandbox up) → `provision environment` (Environment plane starts services) → `readiness gate` (framework-guaranteed; the agent never runs before the world is healthy) → `connect agent` (ACP) → `execute` (the tree grows: Steps and Branches) → `verify` (Reward plane scores) → `teardown`. -**Branch lifecycle.** `quiesce` (disconnect the active agent) → `checkpoint` (snapshot Environment state) → `fork` (N children) → for each child, `restore` the Environment and start a fresh agent session → `score / aggregate` (per-child return → `V(parent)`) → restore the parent's linear rollout state. Agent-session snapshotting is not implemented. Callers that set `require_sandbox_snapshot=True` also require a sandbox with snapshot capability, but the current branch engine still uses the Environment snapshot as its restore point. +**Branch lifecycle.** `quiesce` (disconnect the active agent) → `checkpoint` (snapshot the requested layers — Environment state by default; adding `"sandbox"` to `snapshot_layers` composes a container snapshot with it, environment first) → `fork` (N children) → for each child, `restore` the composed checkpoint (container first, then environment state) and start a fresh agent session (an `env-ready` child runs as a fresh rollout that re-installs the agent for itself) → `score / aggregate` (per-child return → `V(parent)`) → restore the parent's linear rollout state. Agent-session snapshotting is not implemented. `require_sandbox_snapshot=True` keeps its original check-only semantics — it gates on sandbox snapshot capability without composing the layer; requesting the layer through `snapshot_layers` gates the same way and then actually composes the container snapshot into the checkpoint. **Environment lifecycle** (Han's roll-out / roll-back). `provision` → `readiness` → `query` (expose state to the verifier) → `snapshot` → `restore` → `reset` → `teardown`. `snapshot`/`restore` are definitional — the substrate every `Branch` runs on. **Sandbox lifecycle.** configure `expose_ports` → `start` → `exec` / `upload_file` / `upload_dir` / `download_file` / `download_dir` → optional `snapshot` / `restore` → `stop`. -A Rollout branch currently rolls back Environment state and starts a fresh agent session for every child. Container and agent-session checkpoint composition remain future work. The one store that deliberately does **not** roll back with a `Branch` is the continual-learning learner store (capability 5). +A Rollout branch rolls back the composed checkpoint — Environment state plus, when the fork requested it (`snapshot_layers`), the container layer — and starts a fresh agent session for every child. Container + environment checkpoint composition is implemented; agent-session checkpoint composition remains future work. The one store that deliberately does **not** roll back with a `Branch` is the continual-learning learner store (capability 5). ## The four planes diff --git a/docs/continue-runs.md b/docs/continue-runs.md index e32b35ab6..f6fbdc069 100644 --- a/docs/continue-runs.md +++ b/docs/continue-runs.md @@ -13,10 +13,13 @@ environment and continues its own loop with **no injected prompt**. ## The problem it solves -A finished run keeps nothing of the container — cleanup tears the sandbox down. -What survives on disk is the run folder: `config.json`, `result.json`, -`prompts.json`, and `trajectory/llm_trajectory.jsonl`. So a historical timeout -has only its *trajectory* + the *task*; there is no saved container to restore. +A finished run keeps nothing of the container unless it opted into snapshot +retention (`bench eval run --keep-snapshots`, rollout-branching RFC §3.6) — +cleanup tears the sandbox down and marks any recorded stage refs `ephemeral` +in `stage_snapshots.json`. What survives on disk otherwise is the run folder: +`config.json`, `result.json`, `prompts.json`, and +`trajectory/llm_trajectory.jsonl`. So a historical timeout has only its +*trajectory* + the *task*; there is no saved container to restore. `bench eval continue` reconstructs the missing state from the trajectory. @@ -64,6 +67,8 @@ optional. | `--require-timeout` | off | Refuse runs whose recorded status isn't a timeout. | | `--strict-divergence` | off | Abort if replay leaves the original rails. | | `--replay-only` | off | Rebuild via replay and stop at the cut-point (no live model needed). | +| `--max-exchanges K` | all recorded | Replay only the first K recorded exchanges, then go live ([Cut-points](#cut-points)). | +| `--cut-stage STAGE` | — | Cut at a recorded stage boundary by name, resolving K from the run's `stage_snapshots.json` ([Cut-points](#cut-points)). | ### Models and credentials @@ -75,14 +80,76 @@ optional. host needs that provider's credentials (e.g. `GEMINI_API_KEY`) in its environment. `--replay-only` skips the live leg entirely. +## Cut-points + +By default the proxy replays the **entire** recorded prefix before going live. +`--max-exchanges K` cuts the replay short: the first K recorded exchanges are +replayed, then the proxy switches to the live model exactly as if the recording +had ended there. This is the replay cut-point API from the +[rollout-branching RFC §3.5](./rollout-branching-rfc.md) — replay a trajectory +verbatim up to a stage boundary, then go live, to localize which stage a run +went wrong in. + +- `K` must satisfy `1 <= K <= n_recorded`; anything else fails closed before a + sandbox boots. +- The continued run's `source_provenance` gains a `cut_point` block: + `n_replayed_exchanges` plus two request digests named by what they hash — + `served_request_digest` (the request the agent *actually* sent at the cut) + and `recorded_request_digest` (the recorded request it answered for), both + sha256 over the canonical JSON of the comparable projection + `{messages, tools}` (`request_digest_basis` states this in the artifact). + Divergence is checked per replayed exchange on the same basis — a + same-message-count prompt/content/tool change is detected, not only a count + mismatch — and every event (exchange index + both digests) is recorded in + the block's `divergences` list. A divergence annotates rather than aborts + (fidelity caveats are recorded, not hidden — RFC §3.5); `--strict-divergence` + remains the opt-in abort. The block also carries `workspace_digest`: a + deterministic digest of the continuation workspace (`/app` — file contents, + tree and modes) taken as the run crosses the cut into the live leg; when no + live sandbox is reachable at that moment (sandbox proxy mode, a run that + never crossed the cut, a digest failure) the field is `null` and + `workspace_digest_reason` says why — it is never fabricated. The block's + `accounting` field names its basis: in **host** proxy mode the orchestrator + reconciles the block after the run with what the live replay proxy + *actually served* (`accounting: "served"`, plus `configured_max_exchanges` + when a cut was requested — so a run that went live before reaching the + requested cut is visible in artifacts); in **sandbox** proxy mode the + uploaded recording is truncated to the configured prefix and the block + records that basis (`accounting: "configured"`, with the live-only fields + null). A natural-end continuation records the same block, documenting the + end of the recording. +- The stitched `llm_trajectory.jsonl` contains only the replayed prefix (the + first K *parsed* recorded exchanges — a malformed recorded line is never + replayed and never stitched) plus the live suffix. +- Cut-points can be named by **stage** instead of by number: + `--cut-stage ` (e.g. `--cut-stage post-research`) resolves the + exchange index the original run recorded when that stage boundary closed. + A run that captures stage boundaries (`RolloutConfig.snapshot_stages`, or + `Rollout.mark_stage()` for `post-research`) records + `exchanges_completed` per stage in its `stage_snapshots.json`; a cut at + that stage replays exactly that many exchanges. The resolved stage is + recorded as `branch_stage` in the `cut_point` block. Every miss fails + closed with a typed `ReplayCutPointError`: a run with no recorded stages, + an unrecorded stage (the error lists the stages the run *did* record), a + stage recorded without an index (`exchanges_completed: null` — the usage + gateway could not count at capture time), or a stage that closed before + the first exchange. `--cut-stage` and `--max-exchanges` are mutually + exclusive. Through the Python API + (`benchflow.continue_run.orchestrator.continue_run`), an explicit + `stage_tags` mapping (`stage -> 1-based completed-exchange count`) + overrides the recorded registry. + ## Limitations and caveats - **`openhands` only** for now (the proxy seam relies on `LLM_BASE_URL`). - **Replay fidelity is best-effort.** Replay re-runs the original shell commands for real; if a command's output diverges from the original (network, timestamps, nondeterminism), the agent may see a different - observation than recorded. A message-count check warns on divergence - (`--strict-divergence` aborts instead). + observation than recorded. Divergence warns rather than aborts + (`--strict-divergence` aborts instead): through the host replay proxy the + per-exchange check compares content digests of the comparable + `{messages, tools}` projection ([Cut-points](#cut-points)); the in-sandbox + proxy checks message counts only. - **"Identical output" means a faithful continuation**, not a bit-identical result — the model samples, and no "original full run" exists past the timeout. The bar is: the stitched trajectory reads as one continuous run, as diff --git a/docs/reference/cli.md b/docs/reference/cli.md index 01c0c2018..c635adae2 100644 --- a/docs/reference/cli.md +++ b/docs/reference/cli.md @@ -286,6 +286,7 @@ bench eval run --tasks-dir ./tasks --matrix matrix.yaml --trials 3 | `--eval-results-task` | — | Benchmark `task_id`, as defined in the dataset's `eval.yaml` | | `--matrix` | — | YAML model matrix for repeated evals; currently requires `--tasks-dir` | | `--trials` | `1` | Number of trials for `--matrix` | +| `--keep-snapshots` | `false` | Export each captured stage-snapshot image (`docker save`) to `/snapshots/.tar` before cleanup destroys it; `stage_snapshots.json` records the tar's path, sha256 and image id per stage with `ephemeral: false`. Without the flag the committed `bf-snap-*` images die with the run, and cleanup marks each recorded ref `ephemeral: true` with `exported: null` so it never reads as restorable. Only meaningful when the run captures stage snapshots (a task-declared `branch_execution: forked-snapshot`, or an SDK `snapshot_stages` request). Load the exported images back later with [`bench eval import-snapshots`](#bench-eval-import-snapshots) | `--publish-hf`/`--publish-bucket` also write a `README.md` run summary (agent, model, per-task reward and any error/verifier issue, deduplicated @@ -361,6 +362,174 @@ model endpoint; use a model id available to that provider. Provider-specific sampling options are not inferred; pass them explicitly with `--source-env-sampling-arg`. +### bench eval ablate + +Run one task once, snapshot a lifecycle stage boundary as the run passes it, +then fork that one world into a child per **arm** and compare the rewards. The +mechanism is the rollout-branching design — see +[the RFC](../rollout-branching-rfc.md) (§3.2 stage boundaries, §3.3 per-child +deltas, §3.4 lineage artifacts). Every arm starts from a byte-identical world +and differs by exactly one recorded delta, so the arms are comparable in a way +two independent runs are not. + +```bash +# the skills ablation: same env, one child installs the pack, one does not +bench eval ablate --tasks-dir tasks/surface-ion-trap-shuttling \ + --agent claude-agent-acp --model claude-sonnet-4-5 --sandbox docker + +# plan injection at the end of the agent's run, against the parent's own reward +bench eval ablate --tasks-dir tasks/citation-check --at-stage pre-verify \ + --arms inject:oracle-plan.md,inject:decoy-plan.md --json +``` + +| Flag | Default | Description | +|---|---|---| +| `--tasks-dir` | — | Task directory to ablate (or a collection holding exactly one task); several tasks fail closed — the arms are the axis, the task is fixed | +| `--agent` | `claude-agent-acp` | Agent harness for the parent run and every arm (`oracle` is rejected: a branch child connects an agent session, and `solve.sh` has none) | +| `--model` | agent registry | Model for the parent run and every arm | +| `--reasoning-effort` | agent default | Agent reasoning/thinking effort when the agent exposes one (e.g. `max`) — the same control as `bench eval run`, recorded in the parent's and every arm's `config.json` | +| `--sandbox` | `docker` | Sandbox backend; it must implement container snapshot/restore (docker, daytona direct) or the stage capture fails closed | +| `--environment-manifest` | task-declared | Environment-plane manifest for the parent and every arm — a manifest path or a `name@version` registry spec, same semantics as `bench eval run`. An explicit flag beats the task-declared `benchflow.environment.manifest` (the run path's precedence), and the bound environment is stamped into `ablation.json` | +| `--at-stage` | `env-ready` | Stage boundary to fork: `env-ready`, `post-research`, `pre-verify`, `post-verify`. `post-research` is a mid-`execute()` cut point only `Rollout.mark_stage()` can record, so it needs a research-end trigger: pair it with `--mark-research-end-on` (without one, the rejection names both that flag and the SDK path — `mark_stage` at the cut point, then `branch_at_stage`) | +| `--mark-research-end-on` | — | Workspace file whose first appearance marks `post-research` (e.g. `/app/PLAN.md`, the FrontierPhysics convention): the engine polls the sandbox (`test -e`, every ~2s, plus a final check when the agent quiesces) and snapshots the stage the first time the file exists — so the capture lands within one poll of the file appearing, and may include up to that much post-plan agent work. Required for `--at-stage post-research`; rejected for any other stage. If the file never appears, the ablation reports that the marker never appeared instead of forking | +| `--arms` | `with-skill,no-skill` | Comma-separated arms, one branch child each: `with-skill`, `no-skill`, `inject:`, `config:`, `env:`. At least two, no duplicates; commas inside a `config:` arm's JSON braces are content, not separators (quoted string values included); a comma inside a `config:@` path is not representable | +| `--out-dir`, `-o` | `jobs/ablate-` | Output directory: the parent rollout lands in `/ablation//`, the report in `/ablation.json` | +| `--keep-snapshots` | `false` | Export the branched stage's sandbox snapshot image (`docker save`) to `/snapshots/.tar` before cleanup destroys it; `ablation.json` records the tar's path, sha256 and image id (`stage_snapshot.exported`), mirrored into the parent run's `stage_snapshots.json`. Without the flag the committed image dies with the run, and the report records the handle as `ephemeral: true` with `exported: null` so a reader knows the ref no longer resolves. Load the exported image back later with [`bench eval import-snapshots`](#bench-eval-import-snapshots) | +| `--json` | `false` | Emit the report as JSON on stdout instead of the table | + +Arm kinds map onto the `BranchDelta` fields the branch engine executes: +`with-skill` / `no-skill` set `skill_mode`, which re-runs `install_agent()` +under the switched mode — so they are only valid at `env-ready`, the boundary +captured *before* installation. `inject:` sets `injected_prompt`, +delivered as that child's user-visible continuation prompt and recorded as a +content hash only. `config:` sets `config_override`: the arm runs as a +fresh rollout whose config is the parent's with the patch deep-merged on top +through the run-level overlay machinery (#790) — same value-or-`@file` form, +same fail-closed allowlist (never the scorer: a `verifier`/`reward` key is +rejected at parse time), same content addressing (the arm's provenance records +the patch's sha256 and its allowlisted keys, and the child's `config.json` +records the merged overlay). Like the skills arms it is only valid at +`env-ready`, because the child's own `setup()` is what applies the patch. + +`env:` sets `environment_ref`: the arm runs against a *different +environment-registry manifest* — the documented `env0@prod` vs `env0@outage` +tool-outage pattern. The executable slice is **service-level only**: the child +manifest must name the same image as the parent's and both must use +framework-started services (`owns_lifecycle = false`); the arm then restores +the parent's container and provisions the child manifest's service set over +it. A manifest that changes the image fails closed — the snapshot commits the +parent's image, so an image-changing environment delta would need a rebuild, +which contradicts branching from a snapshot — as does an entrypoint-owned +lifecycle on either side (the framework cannot subtract a service the +entrypoint restarts). These gates run before the parent run costs anything, +and the arm is only valid at `env-ready`, where the fresh child provisions its +own environment plane. + +*How* an arm runs depends on `--at-stage`, not on its kind. `env-ready` +precedes `install_agent()`, so **every** arm forked there runs as a fresh child +rollout over the restored snapshot and installs the agent for itself — +including an `inject:` arm, which would otherwise be scored in a world with no +agent, no skills and no path lockdown and reported as "the parent's world plus +one prompt". At `post-research` / `pre-verify` / `post-verify` the agent is +already installed and the arms continue in place: each child restores the +stage's workspace snapshot and connects a **fresh agent session** over it +(agent-session memory is not snapshotted — RFC §7), so a `post-research` fork +compares continuations of the recorded *workspace* boundary; to rebuild the +agent's exact memory up to that boundary instead, replay to it with +`bench eval continue --cut-stage post-research`. + +The table shows one row per arm — reward, pass/fail, wall clock, and a +one-line attribution such as *"fails (0.00) where with-skill passes (1.00) at +env-ready — this delta decides the outcome when applied at env-ready (1 run per +arm)"*. Wording is deliberately confined to the comparison that was run: the +skills arms are read against each other, any other arm against the parent's own +reward, and nothing is claimed about boundaries this invocation did not fork. +Localizing a failure to a stage takes a second ablation at a second boundary. + +**Sub-test attribution.** A binary reward is a lossy summary of what an arm +did: two arms can both score 0.00 while one passes `test_a` and fails `test_b` +and the other does the reverse — a large, reproducible behavioral difference +that nets to exactly zero on the scalar. So a second table follows the first, +listing **only the tests whose outcome differs across the arms**, with each +arm's status per test; tying tests are counted and omitted. Outcomes are read +from each branch child's own verifier CTRF report +(`/verifier/ctrf.json`), the same artifact the eval report mines its +failure lines from — statuses (`passed` / `failed` / `skipped`) pass through +verbatim, and a test one arm's report does not name at all shows as *not +reported* rather than as a failure. + +When the sub-tests disagree while the rewards tie, the arm's own attribution +line says so — *"matches no-skill at env-ready (both 0.00) — scalar rewards +tie, but 2 sub-test outcome(s) differ: test_dominant_factor, +test_trend_result"* — so "no difference in this comparison" can never be +printed over a difference that was measured. A verifier that emits no CTRF +report yields no rows and no invented ones: the section says *scalar-only +attribution* and names the arms it could not read. + +`ablation.json` carries the task id, the stage, the parent's own reward, the +**bound environment** (`environment`: manifest name, the ref exactly as +given — flag value or `task.md` declaration, never a machine-local path — its +`sha256` content address, and the image it names; `null` when no manifest was +bound), the **branched stage's snapshot refs** (`stage_snapshot`: the +committed sandbox image ref, the environment snapshot id, and the captured +layers — the handles to restore that world and re-branch it by hand later, +also on disk in the parent run's `stage_snapshots.json`, which carries +the same `ephemeral`/`exported` lifetime annotation the report does), and +per arm its content-addressed delta provenance, reward, status, wall clock, +the swapped-in environment stamp when an `env:` delta changed it +(`environment`, absent for arms that inherit the parent's), +per-test outcome map (`tests`, `null` when that arm's verifier reported none) +and the branch-child artifact directory +(`/branches//children//`, +holding that child's own `result.json` / `timing.json` / `provenance.json` — +a fresh-rollout child adds the `config.json` its own `setup()` wrote, an +in-place child the `mounted/` archive of what it wrote through the parent's +bind mounts. An in-place child's `result.json` is synthesized from the +child's own scoped state, so fields the child did not produce — token usage, +a run-level error — are `null` rather than inherited from the parent; either +way the child directory answers "what happened in this arm" without reading +`tree.json`). +The top-level `test_attribution` section repeats the differing tests with their +per-arm outcomes, the tying test names, which arms reported per-test data at +all, and whether the scalar rewards tied. +Exit code is 1 when any arm errors or is skipped — a child failure propagates +out of the branch, so the arms after it do not run and are never scored 0.0. A +parent run that failed *after* the boundary is reported but does not fail the +command: attributing a failed run is what the ablation is for. + +### bench eval import-snapshots + +Load a completed run's exported stage-snapshot images back into the local +Docker image store, so a recorded stage boundary can be restored and branched +after the run that captured it is gone. + +```bash +# Everything the run exported +bench eval import-snapshots jobs/2026-08-29__10-00-00/my-task__ab12cd34 + +# One stage only +bench eval import-snapshots jobs/.../my-task__ab12cd34 --stage pre-verify +``` + +Reads `/stage_snapshots.json`; the tars exist only when the run was +made with `--keep-snapshots` (on `bench eval run` or `bench eval ablate`). +Each import verifies the tar's recorded sha256, `docker load`s it, and +verifies the loaded image id matches the recorded one — a snapshot is only +reported restored when `docker image inspect` agrees. Entries recorded +`ephemeral: true` fail closed, naming the flag to re-run with. A run folder +copied from another machine works: when the recorded absolute tar path does +not exist, the tar is found at `/snapshots/`. + +After a successful import the recorded `bf-snap-*` ref resolves again and the +snapshotted world can be branched — restore it into a live sandbox +(`DockerSandbox.restore`) or inspect it directly (`docker run `). SDK +equivalent: `benchflow.snapshot_import.import_stage_snapshots(run_dir)`. + +| Flag | Default | Description | +|------|---------|-------------| +| `--stage` | every exported stage | Import only this stage's snapshot; repeatable | + + ## bench review Grade finished rollouts against a rubric with a reviewer agent. Reviews run @@ -1102,8 +1271,14 @@ The original top-level `bench continue` still works as a hidden, deprecated alia Key options: `--model` (override the live-continuation model; defaults to the original run's model), `--timeout`, `--output`, `--require-timeout`, `--strict-divergence`, `--replay-only` (rebuild via replay and stop at the -cut-point — no live model or API key needed), and `--proxy-mode` (replay -proxy placement: `auto`, `host`, or `sandbox`; default `auto` uses +cut-point — no live model or API key needed), `--max-exchanges` (replay only +the first K recorded LLM exchanges, then go live; default: all recorded — see +[Cut-points](../continue-runs.md#cut-points)), `--cut-stage` (cut at a +recorded stage boundary by name, e.g. `post-research`: resolves the exchange +index the original run's `stage_snapshots.json` recorded when that stage +closed; mutually exclusive with `--max-exchanges`, and an unrecorded stage +fails closed listing the stages the run did record), and `--proxy-mode` +(replay proxy placement: `auto`, `host`, or `sandbox`; default `auto` uses sandbox-local replay for Daytona/Modal and host replay for Docker). ### bench eval continue-batch diff --git a/docs/rollout-branching-rfc.md b/docs/rollout-branching-rfc.md new file mode 100644 index 000000000..39fefd7cf --- /dev/null +++ b/docs/rollout-branching-rfc.md @@ -0,0 +1,342 @@ +# RFC: Composed environment snapshotting + rollout branching + +- **Status:** Draft for review +- **Tracking:** [FrontierPhysics #73](https://github.com/benchflow-ai/FrontierPhysics/issues/73) (the `[12pt - rollout]` infra ticket from the 2026-08-13 sync) +- **Author:** Jicheng Wang (@JeremyJC67) + +## 1. Motivation + +A failed rollout today is an undifferentiated zero. The FrontierPhysics failure cascade +names four stages a research-style agent run can die in — **(1) env/tool init, (2) +research, (3) execution, (4) self-judgment** — and reviewers need to attribute a failure +to a stage, and to ask counterfactuals: *would this run have passed with the skill pack? +with the tool available? with the oracle's plan?* + +The mechanism for both is the same: **snapshot the rollout at stage boundaries, then +branch — resume from a snapshot with exactly one controlled change — and diff the +outcomes.** This RFC specifies that mechanism as a composition of subsystems benchflow +already ships, plus the missing glue. + +Concretely, after this RFC a task author or reviewer can run: + +- *skills ablation without a full re-run*: branch the env-init snapshot into a + `with-skill` child and a `no-skill` child; every other bit of the world is identical. +- *tool-outage perturbation*: branch with an environment-manifest delta + (the documented `env0@prod` vs `env0@outage` pattern) at any stage boundary. +- *plan injection*: branch the post-research boundary with the oracle's `PLAN.md` + substituted, separating "researched wrong" from "executed wrong". +- *stage attribution for a real failure*: replay the recorded trajectory to a cut-point, + then go live under a delta, localizing the first stage whose fix flips the outcome. + +## 2. What already exists (this RFC composes; it does not invent) + +| Substrate | Where | State | +|---|---|---| +| Branch engine: `Rollout.branch(n)` — quiesce → checkpoint → fork → restore → aggregate over a tree-native rollout | `src/benchflow/rollout_branch.py`, `src/benchflow/branch.py`, `docs/architecture.md` ("Branch lifecycle") | **Live**; composes the checkpoint from the layers the fork requests (`snapshot_layers` — Environment state by default, the container layer on request, §3.1) | +| Container snapshots: `Sandbox.snapshot()/restore()`, `SandboxImage`, `supports_snapshot` (fail-closed default) | `src/benchflow/sandbox/protocol.py` (from #384/#470); Docker = `docker commit`, Daytona-direct = provider snapshots | **Live**; gated on `supports_snapshot` and composed into the branch checkpoint whenever the `sandbox` layer is requested | +| Environment-state snapshots (declared sqlite files, `sqlite3 .backup`, fail-closed `EnvironmentSnapshotError`) | `src/benchflow/environment/manifest_env.py` (from #387/#486) | **Live**; the branch engine's default restore layer, composed with the container layer when the fork requests both | +| Record-replay of a finished run: replay `llm_trajectory.jsonl` responses by index through a proxy into a fresh sandbox | `src/benchflow/continue_run/` (`bench eval continue`) | **Live** (openhands-only); replays the full prefix, or cuts short at `--max-exchanges` / `--cut-stage` (§3.5) | +| Per-run variation axes: env-registry refs (S-axis), allowlisted `--config-override` patches (C-axis), `skill_mode` | `environment/_registry/`, `_utils/config_override.py`, `rollout/_config.py` | **Live**; bound at rollout setup, and executable per branch child as §3.3 deltas | +| Task-authoring surface: `branch_execution: forked-snapshot` | `docs/task-standard.md` | **Live** — compiles to a stage-capture request the launch policy honors (snapshot-capable sandboxes only; see task-standard) | + +The gap, as `docs/architecture.md` stated it when this RFC was drafted: *"container +and agent-session checkpoint composition remain future work."* Container + environment +composition has since shipped (§3.1's `snapshot_layers`); agent-session composition is +still future work (§7). Branch children, which then left **no artifacts** — no +per-child `result.json`, no serialized tree, no record of what differed — now leave +the §3.4 lineage. + +## 3. Design + +### 3.1 Composed checkpoints (three layers, fixed order) + +A **stage snapshot** is the ordered composition already sketched in `architecture.md`: + +``` +quiesce (agent disconnected / not yet connected) + → environment.snapshot() # declared state (sqlite), fail-closed integrity + → sandbox.snapshot() # container filesystem (docker commit / provider) + → record StageSnapshot{stage, env_ref, sandbox_ref, meta} +``` + +Restore is symmetric and reversed: restore container, then environment state, then +(for service-topology deltas) stop/start services around the state restore, reusing +`ManifestEnvironment.reset()` semantics. + +**Capability discipline (unchanged pattern):** each layer keeps its `supports_*` flag +and typed error (`SandboxSnapshotNotSupported`, `EnvironmentSnapshotError`). A branch +request declares which layers it requires; missing capability **fails closed with a +diagnostic**, never silently degrades. This generalizes today's behavior instead of +changing it: an env-state-only checkpoint (the historical engine, still the default) +remains expressible as `snapshot_layers={"environment"}`. It also *relaxes* the prior +hard constraint that any branch requires declared sqlite state: a stateless env + +snapshot-capable sandbox can branch with `snapshot_layers={"sandbox"}`. + +Agent-session state is **explicitly layer three and out of scope for v1** (§7). + +### 3.2 Stage boundaries map onto existing lifecycle phases + +No new phase system. The four cascade stages pin to existing transitions in +`rollout/__init__.py`: + +| Cascade stage | Lifecycle boundary | Snapshot point | +|---|---|---| +| env/tool init | end of `start()` (sandbox up, env plane provisioned, readiness gate passed) — **before `install_agent()`** | `env-ready` | +| research | cursor/step boundary inside `execute()` (e.g. the step that finalizes `PLAN.md`) | `post-research` | +| execution | agent finished / quiesced, **before `harden_before_verify`** | `pre-verify` | +| self-judgment | after verify, before review | `post-verify` | + +`env-ready` deliberately precedes `install_agent()` so skills-on/off branches re-run +skill deployment from a skill-free world (skills are baked in at install time). +That world is skill-free only when the *branched run itself* is `no-skill`: a +`with-skill` run's `setup()` injects the pack into the Dockerfile it builds, so its +`env-ready` image already carries `/skills` and a `no-skill` child would restore the +pack, deploy nothing on top of it, and still be labelled `no-skill` everywhere. A +`skill_mode` delta therefore requires a `no-skill` parent and fails closed +(`BranchParentSkillModeConflict`) otherwise. +Mid-`execute()` boundaries are cursor positions in the existing tree — the branch +engine already forks at the cursor; named cut-points are recorded as stage-tagged +exchange indices (§3.5): every stage capture stores the completed-LLM-exchange +index of its moment in `stage_snapshots.json`. `post-research` is markable from +the CLI too: `bench eval ablate --at-stage post-research --mark-research-end-on +/app/PLAN.md` marks the stage the first time the named workspace file exists. + +Snapshot policy is opt-in per run (`RolloutConfig.snapshot_stages`) or per task +(`benchflow.nudges.branch_execution: forked-snapshot`, optionally narrowed with +`branch_stages` — see `docs/task-standard.md`). + +### 3.3 Per-child deltas (reuse the three run-level axes) + +A branch child's delta is a recorded tuple; every member reuses an existing, +content-addressed mechanism: + +```python +BranchDelta( + environment_ref: str | None, # S-axis: registry ref (env0@prod → env0@outage) + config_override: dict | None, # C-axis: allowlisted patch, hashed like #790 + skill_mode: SkillMode | None, # no-skill | with-skill (re-runs install_agent) + injected_prompt: str | None, # e.g. oracle PLAN.md; recorded, never silent +) +``` + +Injection points dictate cost and validity: + +- `skill_mode` / tool-set deltas ⇒ branch from `env-ready` (install re-runs). +- `config_override` ⇒ applied at child setup, same allowlist and hash trail as #790 + (never the scorer). Since the child's own `setup()` is what applies the overlay, + the delta executes only from `env-ready` — the fresh-child boundary — and a + non-allowlisted key fails closed before any child runs. +- `environment_ref` with `[[services]]` changes ⇒ service stop/start bracketing around + state restore (§3.1). **The image-vs-services boundary:** the executable slice is + *service topology only* — same image, `owns_lifecycle = false` on both sides (the + `env0@prod` → `env0@outage` outage pattern), because the restored container's + entrypoint starts nothing and the fresh child's provisioning step decides the + running services. A manifest that changes the image breaks the + restore-the-parent-container premise — it would need a rebuild path, which + contradicts branching from a snapshot — and fails closed + (`BranchEnvironmentImageConflict`); an entrypoint-owned lifecycle on either side + fails closed too, since the framework cannot subtract a service the entrypoint + restarts. Image-changing environment comparisons are two independent runs, not a + branch. +- `injected_prompt` ⇒ delivered as an explicit user-visible message in the child's + session and recorded in provenance. Precedent: 0.6.5 removed silent prompt-level + skill injection (#908); we do not reintroduce it — injection is always a recorded, + first-class delta. + +**How a child executes is decided by the boundary, not by the delta.** `env-ready` is +captured before `install_agent()`, so a world restored to it has no agent binary, no +sandbox user, no seeded verifier workspace, no path lockdown and no skill pack. Every +child the engine runs from that boundary is therefore a *fresh rollout* over the +restored sandbox that installs for itself — including an `injected_prompt` child and a +child with no delta at all, which re-installs the parent's own recorded skill mode. An +in-place child there would either die connecting to an agent the restore deleted or +score a differently-provisioned world and have it reported as a one-delta comparison; +the fork requires the container layer in the stage snapshot and fails closed +(`BranchChildExecutionNotSupported`) without it. At every later boundary the agent is +already installed and children continue in place. + +### 3.4 Lineage: branched runs must be auditable and trainable + +`RolloutTree` used to live and die in memory. This RFC makes branching leave the same +quality of evidence as a linear run: + +- **`tree.json`** in the run folder: nodes, edges, stage tags, snapshot refs, + per-child delta hashes, per-child rewards, aggregate `V(parent)`. +- **Per-child artifact directories** (`children//`) each with standard + `result.json` / `config.json` / trajectory files — a child is a first-class rollout + (implementation seam: `use_prebuilt_env` + the existing child-runner). +- **The parent's own artifacts survive its children.** Children share the parent's + sandbox, and the sandbox bind-mounts the parent rollout's `agent/` / `artifacts/` / + `verifier/` directories into the container — which `restore()` deliberately replays. + Left alone, every child writes (and, via `clear_verifier_output_dir`, *deletes*) the + parent's own evidence, so a finished two-arm ablation shows the second arm's + `reward.txt` where the parent's belongs. The engine therefore takes custody of those + three directories for the duration of the fork: + + ``` + / + agent/ artifacts/ verifier/ # the parent's own, restored at the end + branches// + parent/{agent,artifacts,verifier}/ # transient: held only while children run + children// + mounted/{agent,artifacts,verifier}/ # what THIS child wrote to the shared mounts + ``` + + The mount roots are emptied entry by entry, never removed (deleting a live bind-mount + source detaches it from the running container), and `parent/` is transient — finding + it after a run means the branch did not complete and the parent's evidence is in + there rather than at its canonical path. +- **`source_provenance`** on every child, extending the existing seam (the same one + `benchflow-continue` uses): + +```json +{ + "kind": "benchflow-branch", + "parent_rollout": "", + "parent_task_digest": "sha256:…", + "branch_stage": "env-ready | post-research | pre-verify | post-verify | cursor:", + "snapshot_ref": {"sandbox": "", "environment": ""}, + "cut_point": {"n_replayed_exchanges": 41}, + "delta": {"skill_mode": "with-skill", "config_override_sha256": "…", "environment_ref": null, "injected_prompt_sha256": null} +} +``` + +- The `branched` phase joins the terminal-phase set so `Rollout.result` has clean + semantics for branch-first workflows. + +### 3.5 Replay cut-point API (the cheap bridge for mid-stage branching) + +`ReplayRouter` already serves recorded LLM responses by index; continue-runs prove the +proxy seam end-to-end. Two additions: + +1. **`max_exchanges: int`** — replay the first K exchanges, then switch the proxy to + live passthrough. ("Replay research verbatim, go live at execution.") +2. **Stage-tagged indices** — when a run records stage boundaries (§3.2), the exchange + index that closed each stage is stored (`exchanges_completed` per stage in + `stage_snapshots.json`), so cut-points can be named by stage instead of by + number: `bench eval continue --cut-stage post-research` resolves the recorded + index (an unrecorded stage, or one recorded without an index, fails closed + with a typed error naming what was recorded). + +Divergence accounting: at the cut-point, record content digests of the last replayed +request — both the request actually served and the recorded one, named separately — +and a deterministic digest of the continuation workspace (the `cut_point` provenance +block: `served_request_digest` / `recorded_request_digest` / `workspace_digest`, null +with a reason when no live sandbox is reachable at the cut), plus a per-exchange +content comparison whose divergence events land in the same block — so a +silently-diverged replay is detectable in artifacts. Fidelity caveats inherit from +continue-runs and are recorded, not hidden. + +Generalizing replay beyond openhands (all agents already *record* through the same +LiteLLM gateway) is desirable but independent; it is a named follow-on, not v1. + +### 3.6 Snapshot lifecycle + +`docker commit` images (`bf-snap-*`) and in-sandbox state dirs die with the +rollout (or linger unmanaged). v1 adds: snapshot refs recorded in `tree.json`; optional +`--keep-snapshots` (`bench eval ablate` **and** `bench eval run`) to export captured +container images (`docker save`) into `snapshots/.tar` — the ablation's +``, a plain run's own run directory — for cross-run branching, with each +tar's path, content sha256 and image id recorded in `ablation.json` / +`stage_snapshots.json`. Rollout cleanup stamps every recorded ref's lifetime into +`stage_snapshots.json` before the images are destroyed: without the flag the entry +reads `ephemeral: true` with `exported: null` — a recorded ref whose image cleanup +has destroyed must say so rather than read as restorable. The import half closes the +loop: `bench eval import-snapshots ` +(`benchflow.snapshot_import.import_stage_snapshots`) verifies the exported tar's +sha256, `docker load`s it, and checks the loaded image id against the recorded one, +so a completed run's stage boundary is genuinely branchable later. A remote snapshot +registry is out of scope. + +## 4. Capability matrix (v1) + +| Backend | container layer | env-state layer | notes | +|---|---|---|---| +| Docker | ✅ `docker commit` | ✅ sqlite | reference implementation | +| Daytona (direct) | ✅ provider snapshots | ✅ sqlite | immutable snapshots; restore = recreate | +| Daytona (DinD) | ❌ fail-closed | ✅ | | +| Apple Container | ❌ fail-closed | ✅ | | +| AgentCore / Modal | ❌ fail-closed | ✅ | | + +Known scope limits carried over from the substrate (documented, unchanged): container +snapshots exclude host-mounted volumes and sibling compose services; env-state covers +declared sqlite only. + +## 5. Validation plan + +Three tiers, cheapest first; all deterministic and credential-free: + +- **T1 — mechanical correctness (unit):** snapshot→restore reproduces workspace and + env-DB digests; unsupported backends raise typed errors; per-child artifacts and + `tree.json` validate against the schema; regression tests name this PR per house + convention. +- **T2 — oracle invariants (integration, docker):** on a small task, (a) **zero-delta + branch ⇒ child verifier reward equals parent's** at every stage boundary — an + executable end-to-end proof that restore is lossless; (b) a known-breaking delta + (removing a required tool) ⇒ the child fails at the expected stage with the expected + diagnostic. +- **T3 — attribution demo (evidence for the FrontierPhysics paper):** re-run the + documented no-skill failure of the `surface-ion-trap-shuttling` reference task and + produce the stage-attribution table (branch at `env-ready` with skills ⇒ pass; + branch at `post-research` without ⇒ still fails execution). + +### Stage-level ablation (the T3 surface) + +T3 is not a script anyone re-derives per experiment: it is a command, +`bench eval ablate` (flags in +[docs/reference/cli.md](reference/cli.md#bench-eval-ablate)). One invocation +runs the task once with `--at-stage` captured, forks that recorded world into +one child per `--arms` entry (`with-skill`, `no-skill`, `inject:` — the +§3.3 deltas the engine executes), and writes `ablation.json` beside the +per-child lineage artifacts of §3.4: + +```bash +bench eval ablate --tasks-dir tasks/surface-ion-trap-shuttling \ + --at-stage env-ready --arms with-skill,no-skill +``` + +Two properties keep the output publishable. **The arms are comparable:** they +restore the same snapshot rather than re-running the task, so the world they +differ in is exactly the recorded delta. **The verdict is an observation:** +each row states the two rewards it compares, the boundary they were forked +from, and that it rests on one run per arm — the cross-stage claim ("the +intervention matters at or before stage X") is only earned by a *second* +ablation at a second boundary, which is the T3 table, not one invocation of the +command. + +Out of the command's reach, by construction: `post-research` without a +research-end trigger (only an explicit `Rollout.mark_stage()` records a +mid-`execute()` cut point, §3.2 — `--mark-research-end-on ` +is what supplies that mark from the CLI), image-changing `env:` arms (the +image-vs-services boundary, §3.3 — an image +delta needs a rebuild, which contradicts branching from a snapshot), and +repeated arms for variance — one run per arm is one observation, not an +estimate. `config:` and service-level `env:` arms execute (§3.3) at +`env-ready`, riding the same fresh-child path as the skills arms. + +## 6. Compatibility + +- Targets the #470 `Sandbox` contract as-is — stable across the 0.7 line (#827). +- Any new task frontmatter uses the post-#966 `sandbox:` spelling. +- No prompt-content changes to existing modes (respects #908). +- Branch trees are designed to render in the trace-viewer work (benchflow#987). +- Makes `branch_execution: forked-snapshot` (task-standard) real instead of fail-closed. + +## 7. Out of scope (v1), named follow-ons + +1. **Agent-session snapshot** — documented as the unsolved hard part; v1 children get a + fresh session with replayed-or-injected context. Follow-on: ACP `session/load`. +2. Replay for ACP-native agents (record side already agent-agnostic). +3. Remote snapshot registry / cross-host branching. +4. Verifier-stage re-judgment under alternative judges (needs verifier-isolation + materializer, tracked in task-standard). + +## 8. Workstreams + +| WS | Content | Size | +|---|---|---| +| WS-1 | Composed checkpoint layer (§3.1) + capability matrix + T1 tests | S | +| WS-2 | Deltas (§3.3) + lineage artifacts (§3.4) + T1 tests | M | +| WS-3 | Replay cut-point (§3.5) + T2 oracle invariants + demo (T3) | S/M | + +Matching the sync's "2–3 people" sizing; WS-2/WS-3 are parallelizable after WS-1. diff --git a/docs/task-standard.md b/docs/task-standard.md index 3145f3adc..033a51776 100644 --- a/docs/task-standard.md +++ b/docs/task-standard.md @@ -726,14 +726,35 @@ Today the first model-linear slice accepts `claude-*`, `gpt-*`, and ACP permission handler unless the caller supplies an explicit `on_ask_user` handler, and the ACP `ask_user` bridge preserves both option IDs and option kinds so reject/allow choices are explicit branchable evidence. Authors may -spell the current executable branch slice as -`branch_execution: option-kinds-preserved`; `branch_execution: forked-snapshot` -fails closed until the user loop is integrated with the Environment snapshot -branch engine. The first sequential shared-workspace team handoff slice records -`scene`, `role`, `handoff_from`, and `handoff_to` metadata per user round. -`branchable` is still not automatic branch execution; interactive approval UI, -parallel teams, handoff artifacts, full trajectory sharing, and -branch/message-routing policy remain fail-closed target work. +spell that executable branch slice as +`branch_execution: option-kinds-preserved`. `branch_execution: +forked-snapshot` declares the Environment/sandbox snapshot slice (the +rollout-branching engine): it compiles to a stage-capture request the launch +policy honors — the rollout snapshots the auto-capturable stage boundaries +(`env-ready`/`pre-verify`/`post-verify`; narrow or extend the set with +`branch_stages: [...]`, validated against the branch-stage taxonomy — +declaring `post-research` says the harness will `mark_stage()` it) and +records each boundary's completed-LLM-exchange index in +`stage_snapshots.json`, so `bench eval ablate`, `Rollout.branch_at_stage()`, +or `bench eval continue --cut-stage` can fork the recorded boundaries later. +"Later" is scoped by each ref's recorded lifetime: within the run the +snapshot images are live, but a plain `bench eval run` destroys them at +cleanup and marks every ref `ephemeral: true` in `stage_snapshots.json` — +branching *after* the run completes therefore needs `bench eval run +--keep-snapshots` (which exports each captured image to +`snapshots/.tar` in the run directory) followed by `bench eval +import-snapshots ` to load and identity-check the image before +forking it. `--cut-stage` replay needs only the recorded exchange indices, +no image. Both values require `branchable: true`. Forked-snapshot stays fail-closed +where the engine genuinely cannot honor it: a backend whose sandboxes cannot +take container snapshots (modal, apple-container, agentcore — and Daytona's +DinD strategy at run time) is rejected by task validation / the capture +gate rather than run without the requested captures. The first sequential +shared-workspace team handoff slice records `scene`, `role`, `handoff_from`, +and `handoff_to` metadata per user round. `branchable` is still not automatic +branch execution; interactive approval UI, parallel teams, handoff artifacts, +full trajectory sharing, and branch/message-routing policy remain fail-closed +target work. ## Compatibility @@ -834,7 +855,7 @@ Current implementation status: | verifier `ors-episode` strategy | yes | partial | runtime helper writes ORS tool-output rewards to `trajectory/ors-rewards.jsonl`; declared reward responses/event streams normalize into `reward.json` and `reward-details.json`; fuller OpenReward environment import/export remains target work | | `agents.roles` | yes | partial | `TaskRuntimeView` carries parsed scenes; explicit sequential shared-workspace handoff can switch roles through the user loop | | `scenes` | yes | partial | prompt composition compiles; multi-role document-user scenes execute only with explicit turns and supported team handoff | -| `user` / `## user-persona` | yes | partial | `model: scripted` + string `private_facts` compiles to `DocumentNudgeUser`; bounded model-linear users compile to `ModelDocumentNudgeUser`; linear single- and multi-scene user loops execute when every scene is single-role, or when explicit multi-role turns opt into sequential shared-workspace team handoff; `confirmation_policy: human` gates ACP permissions fail-closed without an explicit handler; `branch_execution: option-kinds-preserved` preserves option IDs and kinds; forked branch execution, interactive approval UI, parallel teams, and rich handoff artifacts fail closed | +| `user` / `## user-persona` | yes | partial | `model: scripted` + string `private_facts` compiles to `DocumentNudgeUser`; bounded model-linear users compile to `ModelDocumentNudgeUser`; linear single- and multi-scene user loops execute when every scene is single-role, or when explicit multi-role turns opt into sequential shared-workspace team handoff; `confirmation_policy: human` gates ACP permissions fail-closed without an explicit handler; `branch_execution: option-kinds-preserved` preserves option IDs and kinds; `branch_execution: forked-snapshot` compiles to a stage-capture request (snapshot-capable sandboxes only; optional `branch_stages`); interactive approval UI, parallel teams, and rich handoff artifacts fail closed | | `benchflow.teams` | yes | partial | supports exactly one `handoff` with `mode: sequential`, `workspace_visibility: shared`, and `trajectory_visibility: none|metadata`; richer team keys fail closed | | `benchflow:` | raw | no | typed document schema after v0.3 stabilizes | | imported `steps` | yes | no/partial | fail closed per sandbox until implemented | diff --git a/src/benchflow/_utils/scoring.py b/src/benchflow/_utils/scoring.py index 0bb3f0d04..65b60c29d 100644 --- a/src/benchflow/_utils/scoring.py +++ b/src/benchflow/_utils/scoring.py @@ -38,6 +38,16 @@ # tokens) is classified by ``_maybe_classify_api_error`` as # ``api_error[rejected_request/permanent]`` and is already non-retryable. PROVIDER_REJECTED = "provider_rejected" +# The agent rejected a *request-global* setting — the model or the reasoning +# effort — that every retry and every branch child would re-request +# identically (PR #1046 second review, P2-A: Gemini + --reasoning-effort high +# provisioned, snapshotted, installed, and only then hit the ACP effort +# rejection, which ablation then retried in a child). Deterministic and +# non-task-attributable, so non-retryable, following the provider_auth (#917) +# pattern: a category plus membership in ``RetryConfig.exclude_categories``. +# ``bench eval ablate`` additionally skips the branch children outright when +# the parent failed with this category (see ``benchflow.ablate``). +REQUEST_GLOBAL = "request_global" TIMED_OUT = "timeout" # Provider API failures detected post-rollout (rate limit, quota, rejected # request, 5xx). "api_error" is proxy-proven (every captured provider request @@ -83,6 +93,13 @@ "provider rejected request", "http 400", ) +# Stamped by the ACP runtime (``ACPRequestGlobalError`` in +# ``benchflow.acp.runtime``) on every deterministic rejection of a global +# request setting. Checked before the "acp error" branch of +# ``classify_error``: these messages often embed the agent's raw +# ``ACP error -326xx`` text, which would otherwise classify as retryable +# ``acp_error``. +REQUEST_GLOBAL_MARKER = "request-global setting rejected" # Verifier error category constants VERIFIER_FAILED = "verifier_failure" @@ -121,6 +138,11 @@ def classify_error(error: str | None) -> str | None: if not error: return None lower = error.lower() + # First: the marker is stamped only by ACPRequestGlobalError, and the + # message usually embeds agent text ("ACP error -326xx…") that the later + # branches would misread as a retryable acp_error. + if REQUEST_GLOBAL_MARKER in lower: + return REQUEST_GLOBAL if "agent idle for" in lower: return IDLE_TIMEOUT if "install failed" in lower: diff --git a/src/benchflow/ablate.py b/src/benchflow/ablate.py new file mode 100644 index 000000000..403e8b3fb --- /dev/null +++ b/src/benchflow/ablate.py @@ -0,0 +1,674 @@ +"""Stage-level ablation — run a task once, branch a stage, compare the arms. + +The library half of ``bench eval ablate`` (rollout-branching RFC §5). The +branch machinery underneath is already complete: a rollout captures a stage +boundary (RFC §3.2), :meth:`~benchflow.rollout.Rollout.branch_at_stage` forks +that recorded world into one child per +:class:`~benchflow.branch_delta.BranchDelta` (RFC §3.3), and every child leaves +lineage artifacts (RFC §3.4). What was missing is the user-facing shape of the +experiment: *arms*. + +An **arm** is one delta plus the name a reader recognizes it by +(``with-skill``, ``no-skill``, ``inject:``). This module drives the +parent rollout to the requested boundary, forks it once into all arms, and +turns the per-arm rewards into an :class:`AblationReport` — a deterministic +``ablation.json`` plus a one-line, observation-only verdict per arm. + +Attribution runs at two granularities, because a binary reward is a lossy +summary of what an arm did: the scalar comparison, and the per-test outcomes +mined from each arm's own verifier report (:func:`differing_tests`, +:func:`sub_test_attribution`). A measured skills ablation that scored 0.00/0.00 +had *both* its sub-tests flip in opposite directions — attributing on the +scalar alone would have reported "no difference" about a large, reproducible +behavioral one. + +Everything decidable from the request alone is decided *before* the parent +rollout runs, in :mod:`benchflow.ablate_arms` (:func:`parse_arms`, +:func:`validate_arms_for_stage` — re-exported here as the one ablation +surface): an ablation costs a full task run before the branch, so a request +the branch engine would reject at fork time must not cost that run first. The +engine keeps its own gates — the pre-flight is a mirror, never a replacement. +""" + +from __future__ import annotations + +import asyncio +import logging +import shlex +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from benchflow.ablate_arms import ( # noqa: F401 (pre-flight API re-exported here) + ARM_KIND_CONFIG, + ARM_KIND_ENV, + ARM_KIND_INJECT, + ARM_KIND_SKILL_MODE, + CAPTURABLE_STAGES, + CONFIG_PREFIX, + ENV_PREFIX, + INJECT_PREFIX, + AblationArm, + AblationError, + AblationRunError, + AblationSpecError, + _split_arm_specs, + parse_arm, + parse_arms, + resolve_ablation_environment_binding, + resolve_ablation_task, + validate_arms_for_stage, +) +from benchflow.branch_report import ( # noqa: F401 (report API re-exported here) + PASS_REWARD, + REFERENCE_PARENT, + REPORT_FILENAME, + SCHEMA_VERSION, + STATUS_ERROR, + STATUS_FAIL, + STATUS_PASS, + STATUS_SKIPPED, + AblationReport, + ArmOutcome, + attribute, + branch_children_of, + differing_tests, + environment_stamp, + outcomes_for_arms, + sub_test_attribution, + write_ablation_report, +) +from benchflow.branch_skill import SKILL_DELTA_LAYER +from benchflow.branch_stage import STAGE_ENV_READY, STAGE_POST_RESEARCH +from benchflow.skill_policy import SKILL_MODE_NO_SKILL + +logger = logging.getLogger(__name__) + +#: The parent rollout's job name under ``--out-dir`` — fixed, not stamped, so +#: the run directory of an ablation is derivable from its output directory. +PARENT_JOB_NAME = "ablation" + + +# Request + + +@dataclass(frozen=True) +class AblationRequest: + """One ablation: a task, a stage boundary, and the arms to fork into.""" + + task_path: Path + arms: Sequence[AblationArm] + agent: str + stage: str = STAGE_ENV_READY + model: str | None = None + # Agent reasoning/thinking effort (``--reasoning-effort``) — the same + # normalized control ``bench eval run`` resolves, threaded through the + # canonical plan so parent and child configs record the effort the arms + # actually ran under. ``None`` = the agent's own default. + reasoning_effort: str | None = None + sandbox: str = "docker" + out_dir: Path = Path("jobs") + # Explicit environment binding (``--environment-manifest``): a manifest + # path or ``name@version`` registry spec. Beats the task-declared + # manifest — the same precedence as ``bench eval run``. ``None`` = bind + # whatever the task declares (or nothing). + environment_manifest: Path | str | None = None + # The research-end trigger (``--mark-research-end-on``): a workspace path + # (absolute, or relative to the agent's cwd) whose first appearance marks + # ``post-research`` during the parent's run — the FrontierPhysics + # convention is the agent materializing its plan as ``/app/PLAN.md``. + # Required for ``stage='post-research'``; meaningless (fail closed) for + # any other stage. ``None`` = no trigger. + mark_research_end_on: str | None = None + # The layers the stage snapshot composes (RFC §3.1). The container layer is + # mandatory for a skills arm and sufficient on its own; the environment + # layer needs an Environment plane, which this command does not bind, so + # requesting it by default would fail closed at capture time. + snapshot_layers: frozenset[str] = frozenset({SKILL_DELTA_LAYER}) + # Durable snapshot retention (RFC §3.6): export the branched stage's + # committed sandbox image (``docker save``) to + # ``/snapshots/.tar`` before cleanup destroys it, and record + # the tar's path + sha256 in the report. Without it the snapshot dies with + # the rollout and the report marks its handle ephemeral instead. + keep_snapshots: bool = False + + +# Execution + + +def resolve_canonical_parent_config( + request: AblationRequest, + *, + stage: str, + environment_manifest: Any = None, +) -> Any: + """The parent's RolloutConfig: the canonical eval plan + the ablation axis. + + An ablation's parent is a *normal evaluation run* of the task — the arms + fork its world, so any control the parent dropped is dropped for every arm + too. The request is therefore resolved through the same two stages ``bench + eval run`` uses: :func:`~benchflow.eval_plan.build_eval_plan` (normalized + agent/model/effort/sandbox/usage settings, fail-closed validation) and + :func:`~benchflow.evaluation.task_rollout_config` (task digest, dataset and + source identity, prompts, the task-declared environment fallback). A + hand-rolled reduced config here is the PR #1046 review finding: the real + E2E parent and child configs published ``task_digest: null`` and + ``reasoning_effort: null``. + + Overlaid on top — the only fields the ablation owns: + + * the stage-capture request (``snapshot_stages={stage}`` plus the request's + layers), which is *why* this rollout exists; + * ``skill_mode='no-skill'`` — stated, not defaulted: the arms fork the + parent's own ``env-ready`` image, and a with-skill parent bakes its pack + into that image, so a ``no-skill`` arm would restore the pack and still + be labelled no-skill. The branch engine refuses that fork; pinning the + parent here keeps every ablation on the side of the gate that runs; + * out-dir / job naming (``/ablation/``, so the run + directory is derivable from the output directory); + * the resolved environment binding (explicit flag beats the task's own + declaration — resolved fail-closed by + :func:`resolve_ablation_environment_binding` before this is called). + + Plan-validation failures re-raise as :class:`AblationSpecError`: they are + request defects, decidable before the parent run costs anything. + """ + from benchflow.eval_plan import EvalCreateRequest, EvalPlanError, build_eval_plan + from benchflow.evaluation import task_rollout_config + + task_path = Path(request.task_path) + try: + plan = build_eval_plan( + EvalCreateRequest( + tasks_dir=task_path, + agent=request.agent, + model=request.model, + reasoning_effort=request.reasoning_effort, + environment=request.sandbox, + jobs_dir=str(request.out_dir), + # One parent rollout at a time — the truthful value for a + # single-task experiment, not the batch default. + concurrency=1, + ) + ) + except EvalPlanError as exc: + raise AblationSpecError(str(exc)) from exc + return task_rollout_config( + plan.make_eval_config(), + task_path, + job_name=PARENT_JOB_NAME, + jobs_dir=plan.output_jobs_dir, + rollout_name=task_path.name, + environment_manifest=environment_manifest, + skill_mode=SKILL_MODE_NO_SKILL, + snapshot_stages={stage}, + snapshot_layers=request.snapshot_layers, + ) + + +#: How often the research-end watcher polls the workspace for the marker file. +RESEARCH_END_POLL_SEC = 2.0 + + +def _resolve_marker_path(rollout: Any, marker: str) -> str: + """An absolute in-sandbox path for the research-end marker file.""" + if marker.startswith("/"): + return marker + cwd = getattr(rollout, "_agent_cwd", None) or "/app" + return f"{str(cwd).rstrip('/')}/{marker}" + + +async def _research_marker_exists(rollout: Any, path: str) -> bool: + """One cheap ``test -e`` for the marker file in the parent's sandbox.""" + sandbox = getattr(rollout, "env", None) + if sandbox is None: + return False + result = await sandbox.exec(f"test -e {shlex.quote(path)}", timeout_sec=10) + return getattr(result, "return_code", 1) == 0 + + +async def watch_research_end( + rollout: Any, marker: str, *, poll_interval: float = RESEARCH_END_POLL_SEC +) -> bool: + """Mark ``post-research`` the first time ``marker`` exists in the workspace. + + The concrete trigger behind ``--mark-research-end-on`` (RFC §3.2): a + research-style agent materializes its plan as a workspace file (the + FrontierPhysics convention is ``/app/PLAN.md``), so the file's first + appearance *is* the research→execution boundary. The engine has no + per-LLM-exchange hook from outside the agent process, so the check runs + on a cheap wall-clock poll (one sandbox ``test -e`` per + ``poll_interval``) concurrent with ``execute()``, plus a final check when + the agent quiesces — **the capture therefore lands within one poll of the + file appearing, and the snapshot may include up to that much post-plan + agent work.** That cadence bound is the documented tradeoff of marking + from outside the agent; the exchange index recorded with the mark + (``capture_stage``) is exact for the moment the capture actually ran. + + Transient poll failures (a raced exec, teardown) keep polling; a + ``mark_stage()`` failure propagates — a capture the run was told to take + and could not is the same fail-closed rule the lifecycle's own boundaries + apply. Returns True once the stage is marked; cancellation is the normal + end when the agent finishes before the marker ever appears. + """ + path = _resolve_marker_path(rollout, marker) + while True: + try: + found = await _research_marker_exists(rollout, path) + except asyncio.CancelledError: + raise + except Exception: + logger.debug("research-end marker poll failed", exc_info=True) + found = False + if found: + await rollout.mark_stage(STAGE_POST_RESEARCH) + return True + await asyncio.sleep(poll_interval) + + +async def _execute_parent(rollout: Any, research_end_marker: str | None) -> None: + """``execute()``, with the research-end watcher beside it when requested. + + The watcher runs as a sibling task for the duration of the agent's run and + is cancelled when the agent quiesces; a marker that appeared between the + watcher's last poll and quiescence is caught by one final check, so the + trigger is decided by the file's presence, not by poll timing. A watcher + failure (its ``mark_stage`` raising) is logged and leaves the stage + uncaptured — the branch then fails closed on the missing stage rather + than this masking the agent's own outcome. + """ + if research_end_marker is None: + await rollout.execute() + return + watcher = asyncio.create_task(watch_research_end(rollout, research_end_marker)) + marked = False + watcher_failed = False + try: + await rollout.execute() + finally: + if not watcher.done(): + watcher.cancel() + try: + marked = await watcher + except asyncio.CancelledError: + marked = False + except Exception: + watcher_failed = True + logger.warning( + "research-end watcher failed — 'post-research' was not " + "captured mid-run", + exc_info=True, + ) + if marked or watcher_failed: + return + if await _research_marker_exists( + rollout, _resolve_marker_path(rollout, research_end_marker) + ): + await rollout.mark_stage(STAGE_POST_RESEARCH) + + +async def _run_parent( + rollout: Any, stage: str, *, research_end_marker: str | None = None +) -> tuple[float | None, str | None]: + """Drive the parent past ``stage`` and score it; return ``(reward, error)``. + + Not ``Rollout.run()``: that cleans the sandbox up on the way out, and a + torn-down sandbox has nothing left to branch. The phases below are the + linear lifecycle in order, with the boundary captured by the rollout's own + ``snapshot_stages`` policy as it passes — except ``post-research``, which + is marked by the research-end watcher (:func:`watch_research_end`) when + ``research_end_marker`` names the trigger file. + + A failure *after* the boundary is recorded, not raised: the snapshot the + arms fork from was already taken, and attributing a failed run is the + reason this command exists (RFC §1). + """ + from benchflow._utils.text import describe_exception + + try: + await rollout.setup() + await rollout.start() + except Exception as exc: + raise AblationRunError( + f"the parent rollout failed before the {stage!r} boundary, so there " + f"is nothing to branch: {describe_exception(exc)}" + ) from exc + try: + await rollout.install_agent() + await rollout.connect() + await _execute_parent(rollout, research_end_marker) + rewards = await rollout.verify() + except Exception as exc: + error = describe_exception(exc) + logger.warning( + "ablation parent run failed after the %r boundary (%s) — the arms " + "still fork from the recorded snapshot", + stage, + error, + ) + return None, error + if not rewards: + return None, None + reward = rewards.get("reward") + return (None if reward is None else float(reward)), None + + +async def export_stage_snapshot( + sandbox: Any, *, sandbox_ref: str, out_dir: Path +) -> dict[str, Any]: + """``docker save`` the branched stage's sandbox image into ``out_dir``. + + The durable half of ``--keep-snapshots`` (RFC §3.6): the tar lands at + ``/snapshots/.tar`` and the returned record carries its + path, content sha256 and image id — enough for a reader to verify, + ``docker load`` and identity-check the world later (``bench eval + import-snapshots``). Raises :class:`AblationError` when the sandbox + backend cannot export (the caller records the failure; the report stays + truthful). Thin ablation-facing wrapper over the shared retention + machinery in :mod:`benchflow.branch_policy` — the same code path ``bench + eval run --keep-snapshots`` exports through, so the two artifacts cannot + drift. + """ + from benchflow.branch_policy import SnapshotExportUnsupported + from benchflow.branch_policy import export_stage_snapshot as _export_engine + + try: + return await _export_engine(sandbox, sandbox_ref=sandbox_ref, out_dir=out_dir) + except SnapshotExportUnsupported as exc: + raise AblationError(str(exc)) from exc + + +async def retain_stage_snapshot( + report: AblationReport, + *, + sandbox: Any, + keep_snapshots: bool, + out_dir: Path, + run_dir: Path | str | None = None, +) -> None: + """Make the report truthful about the stage snapshot's lifetime (RFC §3.6). + + Must run **before** ``rollout.cleanup()`` — cleanup's + ``compose down --rmi all`` is what destroys the committed ``bf-snap-*`` + image, and a report serialized afterwards once published a handle + ``docker image inspect`` could no longer resolve. With + ``keep_snapshots`` the image is exported to a tar and the entry records + ``ephemeral: false`` plus the tar's path, sha256 and image id; without it + (or when the export fails — recorded as ``export_error``) the entry + records ``ephemeral: true, exported: null`` so a reader knows the ref no + longer resolves. Never raises: the arms' rewards must survive a failed + export. + + ``run_dir`` (the parent rollout's run directory, when known) receives the + same annotation into its ``stage_snapshots.json`` entry — the on-disk + twin of ``report.stage_snapshot`` — so the parent artifact and the + ablation report tell one story; ``Rollout.cleanup()``'s own lifetime + finalization then preserves this entry instead of re-marking it. + """ + from benchflow.branch_policy import annotate_stage_snapshot_lifetime + + snap = report.stage_snapshot + if snap is None: + return + await annotate_stage_snapshot_lifetime( + snap, sandbox=sandbox, keep=keep_snapshots, out_dir=Path(out_dir) + ) + if run_dir is None: + return + try: + from benchflow.branch_lineage import annotate_stage_snapshots_file + + annotate_stage_snapshots_file( + run_dir=Path(run_dir), annotations={report.stage: snap} + ) + except Exception: + logger.warning( + "stage-snapshot lifetime annotation under %s failed — the report " + "is unaffected", + run_dir, + exc_info=True, + ) + + +# The run_ablation phases — resolve/prepare, parent leg, branch leg, +# retention, report — each with the failure-isolation rule it owns. + + +def _preflight_environment_arms( + arms: Sequence[AblationArm], environment_manifest: Any +) -> dict[str, dict[str, Any]]: + """Pre-flight the env arms' content gates; stamp the arms that pass. + + Runs once the parent's manifest is known, still *before* the parent run: + an unresolvable ref, an image-changing manifest, or an entrypoint-owned + lifecycle is decidable from the request alone, and the branch engine + would reject it only after a full parent run. An arm that passes the gate + is stamped with the environment it swaps in, so its report row names the + world it ran against. + """ + from benchflow.branch_skill import resolve_environment_ref_delta + from benchflow.environment.manifest import load_manifest_binding + + environment_stamps: dict[str, dict[str, Any]] = {} + for arm in arms: + if arm.kind != ARM_KIND_ENV or arm.delta.environment_ref is None: + continue + try: + resolve_environment_ref_delta( + environment_manifest, + arm.delta.environment_ref, + subject=f"arm {arm.name!r}", + ) + except NotImplementedError as exc: + raise AblationSpecError(str(exc)) from exc + stamp = environment_stamp(load_manifest_binding(arm.delta.environment_ref)) + if stamp is not None: + environment_stamps[arm.name] = stamp + return environment_stamps + + +async def _branch_into_arms( + rollout: Any, request: AblationRequest, report: AblationReport, stage: str +) -> str | None: + """The branch leg: fork the recorded boundary once into all the arms. + + Returns the branch error instead of raising — a branch failure becomes + reported arm errors, and the arms that did run keep their rewards. Two + gates skip the fork outright, each naming what actually happened rather + than what the engine would have raised: + + * a parent error classifying **request-global** (an unsupported reasoning + effort / model — PR #1046 second review, P2-A): every child would + restore the snapshot, reinstall the agent, and be rejected identically, + so the error is a property of the request, never of a task or an arm; + * a research-end trigger that never fired: branching would only raise + ``BranchStageNotCaptured``, and the marker file — not the stage + machinery — is what was missing. + """ + from benchflow._utils.scoring import REQUEST_GLOBAL, classify_error + from benchflow._utils.text import describe_exception + + if classify_error(report.parent_error) == REQUEST_GLOBAL: + branch_error = ( + "the parent run failed with a request-global configuration " + f"error, so the arms were not attempted: " + f"{report.parent_error} — every branch child would re-run " + "the same rejected agent/model configuration" + ) + logger.error("ablation branch at %r skipped: %s", stage, branch_error) + return branch_error + if request.mark_research_end_on is not None and stage not in getattr( + rollout, "_stage_snapshots", {} + ): + branch_error = ( + f"the research-end marker {request.mark_research_end_on!r} " + "never appeared in the workspace during the parent run, so " + f"{stage!r} was never captured and there is nothing to branch" + ) + logger.error("ablation branch at %r failed: %s", stage, branch_error) + return branch_error + try: + report.value = await rollout.branch_at_stage( + stage, + len(request.arms), + deltas=[arm.delta for arm in request.arms], + ) + except Exception as exc: + branch_error = describe_exception(exc) + logger.error("ablation branch at %r failed: %s", stage, branch_error) + return branch_error + return None + + +async def _retain_and_cleanup( + rollout: Any, request: AblationRequest, report: AblationReport, stage: str +) -> None: + """The retention leg, then — always — the parent sandbox's cleanup. + + The branched stage's recorded snapshot refs — the committed sandbox image + and environment snapshot id a reader needs to restore this world and + re-branch it by hand later (RFC §3.2; also on disk in the parent run's + ``stage_snapshots.json``) — are read, retained and annotated **before** + cleanup: cleanup destroys the committed image, and a report serialized + afterwards once published a snapshot ref that ``docker image inspect`` + could no longer resolve. Cleanup always runs, and never over a masked + retention error (:func:`retain_stage_snapshot` records failures instead + of raising). + """ + try: + stage_registry = getattr(rollout, "_stage_snapshots", None) + if stage_registry: + from benchflow.branch_lineage import stage_snapshots_payload + + report.stage_snapshot = stage_snapshots_payload(stage_registry).get(stage) + await retain_stage_snapshot( + report, + sandbox=getattr(rollout, "env", None), + keep_snapshots=request.keep_snapshots, + out_dir=Path(request.out_dir), + run_dir=getattr(rollout, "_rollout_dir", None), + ) + finally: + await rollout.cleanup() + + +def _finalize_report( + rollout: Any, + request: AblationRequest, + report: AblationReport, + *, + stage: str, + branch_error: str | None, + environment_stamps: dict[str, dict[str, Any]], +) -> AblationReport: + """The report leg: per-arm outcomes off the tree, attribution, artifacts. + + A branch that failed before it forked anything (an uncaptured stage, a + capability gap) has no arm to carry the error, so the report carries it. + The parent's own result materialization is failure-isolated — the arms' + rewards are already in the report and must survive a result-build error. + """ + run_dir = getattr(rollout, "_rollout_dir", None) + report.parent_run_dir = None if run_dir is None else str(run_dir) + report.arms = outcomes_for_arms( + request.arms, + branch_children_of(rollout.tree), + run_dir=run_dir, + branch_error=branch_error, + environment_stamps=environment_stamps, + ) + if branch_error is not None and all(arm.error is None for arm in report.arms): + report.error = branch_error + attribute(report.arms, parent_reward=report.parent_reward, stage=stage) + try: + materialized = rollout.result + except Exception: + logger.warning( + "ablation parent result artifacts failed to build — the arms' " + "rewards are unaffected", + exc_info=True, + ) + else: + if materialized is None: + logger.info("ablation parent reached no terminal result to materialize") + return report + + +async def run_ablation(request: AblationRequest) -> AblationReport: + """Run the task once, fork the requested stage into the arms, score them. + + The whole command in one call: validate (again — the library is the + contract, the CLI one caller), run the parent to the stage boundary, fork + it once into ``len(arms)`` children carrying the arms' deltas, and read the + per-arm rewards back off the tree the engine grew. The parent's sandbox is + always cleaned up, and a branch failure becomes reported arm errors rather + than an exception — the arms that did run keep their rewards. + + The phases, in order — each a named function carrying its own + failure-isolation rule: pre-flight (:func:`validate_arms_for_stage`, + :func:`resolve_ablation_environment_binding`, + :func:`_preflight_environment_arms`, :func:`resolve_canonical_parent_config` + — all before the parent run costs anything), the parent leg + (:func:`_run_parent`), the branch leg (:func:`_branch_into_arms`), + retention + cleanup (:func:`_retain_and_cleanup` — always runs), and the + report (:func:`_finalize_report`). + """ + from benchflow.rollout import Rollout + + stage = validate_arms_for_stage( + request.arms, + request.stage, + snapshot_layers=request.snapshot_layers, + research_end_marker=request.mark_research_end_on, + ) + if request.agent == "oracle": + raise AblationSpecError( + "bench eval ablate needs an ACP agent: every branch child connects " + "an agent session over the restored snapshot, and the oracle path " + "(solve.sh) has no session to fork" + ) + task_path = Path(request.task_path) + environment_binding = resolve_ablation_environment_binding( + task_path, explicit=request.environment_manifest + ) + environment_manifest = ( + None if environment_binding is None else environment_binding.manifest + ) + environment_stamps = _preflight_environment_arms(request.arms, environment_manifest) + # The canonical evaluation configuration with the ablation axis overlaid — + # the bound world (an explicit ``--environment-manifest`` when given, else + # the task's own declaration) rides along: every arm forks the parent's + # snapshot and (at ``env-ready``) re-runs from the parent's config, so + # binding it here binds it for the whole experiment. + config = resolve_canonical_parent_config( + request, stage=stage, environment_manifest=environment_manifest + ) + rollout = Rollout(config) + report = AblationReport( + task_id=task_path.name, + task_path=str(task_path), + stage=stage, + snapshot_layers=sorted(request.snapshot_layers), + agent=config.agent, + model=config.model, + sandbox=request.sandbox, + arms=[], + environment=environment_stamp(environment_binding), + ) + branch_error: str | None = None + try: + report.parent_reward, report.parent_error = await _run_parent( + rollout, stage, research_end_marker=request.mark_research_end_on + ) + branch_error = await _branch_into_arms(rollout, request, report, stage) + finally: + await _retain_and_cleanup(rollout, request, report, stage) + return _finalize_report( + rollout, + request, + report, + stage=stage, + branch_error=branch_error, + environment_stamps=environment_stamps, + ) diff --git a/src/benchflow/ablate_arms.py b/src/benchflow/ablate_arms.py new file mode 100644 index 000000000..c596af4a2 --- /dev/null +++ b/src/benchflow/ablate_arms.py @@ -0,0 +1,481 @@ +"""Ablation pre-flight — everything decidable from the request alone. + +The parse/validate half of ``bench eval ablate`` (rollout-branching RFC §5), +split out of :mod:`benchflow.ablate` so that module stays the *orchestration* +of an experiment and this one stays its *admission gate*. An ablation costs a +full task run before the branch, so a request the branch engine would reject +at fork time must not cost that run first: arm specs parse into executable +:class:`~benchflow.branch_delta.BranchDelta` values here (:func:`parse_arm`, +:func:`parse_arms`), stage/arm combinations are rejected before anything runs +(:func:`validate_arms_for_stage`), and the one task and its environment +binding resolve fail-closed (:func:`resolve_ablation_task`, +:func:`resolve_ablation_environment_binding`). The engine keeps its own gates +— everything here is a pre-flight mirror, never a replacement. + +The ablation error hierarchy lives here too, beside its earliest raisers; +:mod:`benchflow.ablate` re-exports every public name, so callers keep +importing the one ablation surface. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +from benchflow.branch_delta import BranchDelta +from benchflow.branch_skill import ( + FRESH_CHILD_LAYER, + FRESH_CHILD_STAGE, + SKILL_DELTA_STAGE, +) +from benchflow.branch_stage import ( + BRANCH_STAGES, + MARKED_STAGES, + STAGE_POST_RESEARCH, + validate_stage, +) +from benchflow.skill_policy import SKILL_MODE_NO_SKILL, SKILL_MODE_WITH_SKILL + +if TYPE_CHECKING: + from benchflow.environment.manifest import ManifestBinding + +#: Arm spec prefix for a plan-injection arm: ``inject:``. +INJECT_PREFIX = "inject:" + +#: Arm spec prefix for a config-override arm: ``config:`` — +#: the same dual "value or ref" form as the run-level ``--config-override``. +CONFIG_PREFIX = "config:" + +#: Arm spec prefix for an environment arm: ``env:`` +#: (the ``env0@prod`` vs ``env0@outage`` tool-outage pattern). +ENV_PREFIX = "env:" + +ARM_KIND_SKILL_MODE = "skill-mode" +ARM_KIND_INJECT = "inject" +ARM_KIND_CONFIG = "config-override" +ARM_KIND_ENV = "environment-ref" + +#: The stages ``bench eval ablate`` can capture on its own — everything the +#: lifecycle reaches without an explicit ``Rollout.mark_stage()`` call. +CAPTURABLE_STAGES: tuple[str, ...] = tuple( + stage for stage in BRANCH_STAGES if stage not in MARKED_STAGES +) + +_SKILL_ARM_NAMES = (SKILL_MODE_WITH_SKILL, SKILL_MODE_NO_SKILL) +_ARM_SPEC_HELP = ( + f"supported arms are {SKILL_MODE_WITH_SKILL!r}, {SKILL_MODE_NO_SKILL!r}, " + f"'{INJECT_PREFIX}', " + f"'{CONFIG_PREFIX}', and " + f"'{ENV_PREFIX}'" +) + + +class AblationError(Exception): + """Base class for ablation failures the CLI reports without a traceback.""" + + +class AblationSpecError(AblationError, ValueError): + """The requested ablation cannot run as asked (fail closed, before any run). + + Raised for a malformed arm spec, an arm the branch engine could not execute + at the requested stage, a stage this command cannot capture, or a + ``--tasks-dir`` that does not name exactly one task. + """ + + +class AblationRunError(AblationError, RuntimeError): + """The parent rollout never reached the stage boundary to branch from.""" + + +# Arms + + +@dataclass(frozen=True) +class AblationArm: + """One arm of the ablation: a recognizable name and the delta it runs under. + + ``name`` is the spec exactly as the caller wrote it — it is the row label + in the table, the key in ``ablation.json``, and the reference other arms + are compared against, so it survives the round trip verbatim. ``source`` + records where an injection arm read its text from; the text itself is + never recorded (provenance hashes it, per #908). + """ + + name: str + kind: str + delta: BranchDelta + source: str | None = None + + +def parse_arm(spec: str) -> AblationArm: + """Parse one arm spec into an :class:`AblationArm` (fail closed). + + Five kinds, each mapping onto exactly one executable + :class:`BranchDelta` field: ``with-skill`` / ``no-skill`` become + ``skill_mode`` (the child re-runs installation as a fresh rollout from the + ``env-ready`` snapshot), ``inject:`` reads the file and becomes + ``injected_prompt`` (the child's user-visible continuation prompt) — which + at ``--at-stage env-ready`` is *also* delivered by a fresh rollout, since + every child of that boundary installs the agent for itself — + ``config:`` parses through the run-level overlay loader + and becomes ``config_override`` (the child runs fresh under the parent's + config with the allowlisted patch deep-merged on top, #790), and + ``env:`` becomes ``environment_ref`` (the service-level + environment swap; resolution against the parent's manifest happens in + :func:`~benchflow.ablate.run_ablation`, where the parent's manifest is + known). An unknown kind, an empty spec, an injection file that is missing + or blank, or a config patch that is unparsable or touches a + non-allowlisted section raises :class:`AblationSpecError` — a silently + dropped arm would publish an ablation table with a missing comparison, and + a scorer-touching patch must die here, before the parent run costs + anything. + """ + name = spec.strip() + if not name: + raise AblationSpecError(f"empty arm in --arms — {_ARM_SPEC_HELP}") + if name in _SKILL_ARM_NAMES: + return AblationArm( + name=name, kind=ARM_KIND_SKILL_MODE, delta=BranchDelta(skill_mode=name) + ) + if name.startswith(CONFIG_PREFIX): + return _parse_config_arm(name) + if name.startswith(ENV_PREFIX): + ref = name[len(ENV_PREFIX) :].strip() + if not ref: + raise AblationSpecError( + f"arm {name!r} names no environment — the environment arm is " + f"'{ENV_PREFIX}' (e.g. '{ENV_PREFIX}env0@outage') " + "or a manifest file path" + ) + return AblationArm( + name=name, kind=ARM_KIND_ENV, delta=BranchDelta(environment_ref=ref) + ) + if name.startswith(INJECT_PREFIX): + raw = name[len(INJECT_PREFIX) :].strip() + if not raw: + raise AblationSpecError( + f"arm {name!r} names no file — the injection arm is " + f"'{INJECT_PREFIX}' (e.g. " + f"'{INJECT_PREFIX}oracle-plan.md')" + ) + path = Path(raw) + try: + text = path.read_text(encoding="utf-8") + except OSError as exc: + raise AblationSpecError( + f"arm {name!r} cannot read its injection file {raw!r}: {exc}" + ) from None + if not text.strip(): + raise AblationSpecError( + f"arm {name!r} reads an empty injection file {raw!r} — an " + "injection arm must carry the text it injects" + ) + return AblationArm( + name=name, + kind=ARM_KIND_INJECT, + delta=BranchDelta(injected_prompt=text), + source=str(path), + ) + raise AblationSpecError(f"unknown arm {name!r} — {_ARM_SPEC_HELP}") + + +def _parse_config_arm(name: str) -> AblationArm: + """Parse a ``config:`` arm through the #790 loader. + + The value after the prefix is exactly what ``--config-override`` accepts — + inline JSON/YAML/TOML or an ``@file`` ref — parsed by the same + :func:`~benchflow._utils.config_override.load_config_override` so the two + surfaces cannot drift. The allowlist runs *here*, at parse time: an arm + that patches the scorer must fail before the parent run costs anything, + exactly as the branch engine would fail it at fork time. + """ + from benchflow._utils.config_override import load_config_override, validate_overlay + + raw = name[len(CONFIG_PREFIX) :].strip() + if not raw: + raise AblationSpecError( + f"arm {name!r} carries no patch — the config arm is " + f"'{CONFIG_PREFIX}' (e.g. " + f"""'{CONFIG_PREFIX}{{"agent": {{"timeout_sec": 60}}}}' or """ + f"'{CONFIG_PREFIX}@overlay.yaml')" + ) + try: + overlay = load_config_override(raw) + except (ValueError, OSError) as exc: + raise AblationSpecError( + f"arm {name!r} cannot load its config patch: {exc}" + ) from None + if not overlay: + raise AblationSpecError( + f"arm {name!r} parses to an empty patch — a config arm must " + "change at least one allowlisted section" + ) + try: + validate_overlay(overlay) + except ValueError as exc: + raise AblationSpecError(f"arm {name!r} is not executable: {exc}") from None + return AblationArm( + name=name, + kind=ARM_KIND_CONFIG, + delta=BranchDelta(config_override=overlay), + source=raw[1:] if raw.startswith("@") else None, + ) + + +def _split_arm_specs(spec: str) -> list[str]: + """Split ``--arms`` on commas, ignoring commas nested in JSON braces. + + A ``config:`` arm may carry inline JSON (``config:{"agent": {"a": 1, + "b": 2}}``) whose commas are content, not separators. Depth counting over + ``{}``/``[]`` keeps every historical spec splitting exactly as before — + no other arm kind can contain a brace. Inside a JSON string literal, + braces, brackets and commas are content too, so the walk tracks quote + state (with ``\\``-escape handling) and ignores structure until the + string closes. Commas in a ``config:@`` file path remain + unrepresentable — that grammar limit is documented on ``--arms``. + """ + parts: list[str] = [] + current: list[str] = [] + depth = 0 + in_string = False + escaped = False + for char in spec: + if in_string: + if escaped: + escaped = False + elif char == "\\": + escaped = True + elif char == '"': + in_string = False + current.append(char) + continue + if char == '"' and depth > 0: + in_string = True + elif char in "{[": + depth += 1 + elif char in "}]": + depth = max(0, depth - 1) + if char == "," and depth == 0: + parts.append("".join(current)) + current = [] + else: + current.append(char) + parts.append("".join(current)) + return parts + + +def parse_arms(spec: str) -> list[AblationArm]: + """Parse a comma-separated ``--arms`` value into ordered arms. + + Two whole-request rules beyond the per-arm ones. A fork needs at least two + children (the branch engine's own ``n >= 2``), so a single arm is rejected + here rather than after a full parent run; and duplicate arm names are + rejected because the report keys arms by name, which would make the + attribution ambiguous. + """ + if not spec.strip(): + raise AblationSpecError( + "--arms is empty — pass at least two arms, e.g. " + f"'{SKILL_MODE_WITH_SKILL},{SKILL_MODE_NO_SKILL}'" + ) + arms = [parse_arm(entry) for entry in _split_arm_specs(spec)] + seen: set[str] = set() + for arm in arms: + if arm.name in seen: + raise AblationSpecError( + f"duplicate arm {arm.name!r} in --arms — each arm is one " + "controlled change, and the report keys arms by name" + ) + seen.add(arm.name) + if len(arms) < 2: + raise AblationSpecError( + f"--arms needs at least two arms to compare, got one " + f"({arms[0].name!r}) — a branch forks into >= 2 children. Pair it " + f"with a counterpart arm, e.g. " + f"'{SKILL_MODE_WITH_SKILL},{SKILL_MODE_NO_SKILL}'" + ) + return arms + + +def validate_arms_for_stage( + arms: Sequence[AblationArm], + stage: str, + *, + snapshot_layers: frozenset[str] | set[str] | None = None, + research_end_marker: str | None = None, +) -> str: + """Reject a stage/arm combination before the parent rollout runs. + + Pre-flight gates, each mirroring a rule that lives elsewhere: + + * ``post-research`` is a mid-``execute()`` cut point only an explicit + ``Rollout.mark_stage()`` can record (RFC §3.2). This command drives the + lifecycle, not the agent's planning, so on its own it cannot mark it — + ``--mark-research-end-on `` supplies the concrete + trigger (the engine marks the stage the first time that file exists in + the workspace; see :func:`~benchflow.ablate.watch_research_end`), and + without it the rejection says so up front, which beats running the whole + task and dying at the fork. The marker is meaningless for any other + stage and fails closed there. + * a ``skill_mode`` arm executes only at ``env-ready``, because skills are + deployed by ``install_agent()`` (RFC §3.3). The branch engine fails + closed on this too; here it costs nothing instead of a full run. + * a ``config:`` arm executes only at ``env-ready`` for the same shape of + reason: the overlay is applied by the child's own ``setup()``, which + only a fresh child of that boundary re-runs. + * an ``env:`` arm executes only at ``env-ready`` likewise: the child + manifest's services are provisioned over the restored snapshot by the + fresh child. (Its *content* gates — resolvable ref, same image, + framework-started services — need the parent's manifest and run in + :func:`~benchflow.ablate.run_ablation`, still before the parent run.) + * an ``env-ready`` ablation runs *every* arm as a fresh rollout — that + boundary precedes ``install_agent()``, so each child installs the agent + for itself — which needs the container layer in the stage snapshot. The + engine raises :class:`~benchflow.rollout_branch.BranchChildExecutionNotSupported` + for this; the mirror here is checked only when the caller passes the + layers it will request. + + Returns the validated stage, so callers can use this as their one gate. + """ + validate_stage(stage, field="--at-stage") + if research_end_marker is not None and stage != STAGE_POST_RESEARCH: + raise AblationSpecError( + f"--mark-research-end-on only applies to --at-stage " + f"{STAGE_POST_RESEARCH!r}, not {stage!r}: the marker file's first " + "appearance is what defines the research-end boundary, and no " + "other stage is captured from it" + ) + if stage in MARKED_STAGES and research_end_marker is None: + raise AblationSpecError( + f"--at-stage {stage!r} cannot be captured by this command without " + "a research-end trigger: it is a mid-execute() cut point only an " + "explicit Rollout.mark_stage() call can record. Pass " + "--mark-research-end-on (e.g. /app/PLAN.md) so " + "the engine marks the stage the first time that file appears " + "during the agent's run; or drive the run through the SDK — " + f"await rollout.mark_stage({stage!r}) at the cut point during the " + f"agent's run, then await rollout.branch_at_stage({stage!r}, n, " + "deltas=[...]) — or ablate one of " + f"{list(CAPTURABLE_STAGES)!r}" + ) + for arm in arms: + if arm.kind == ARM_KIND_SKILL_MODE and stage != SKILL_DELTA_STAGE: + raise AblationSpecError( + f"arm {arm.name!r} needs --at-stage {SKILL_DELTA_STAGE!r}, not " + f"{stage!r}: skills are deployed by install_agent(), which has " + f"already run by {stage!r}, so the arm would measure nothing" + ) + if arm.kind == ARM_KIND_CONFIG and stage != FRESH_CHILD_STAGE: + raise AblationSpecError( + f"arm {arm.name!r} needs --at-stage {FRESH_CHILD_STAGE!r}, not " + f"{stage!r}: the config patch is applied by the child's own " + f"setup(), and by {stage!r} the parent's config has already " + "been consumed — the arm would record the override and run " + "without it" + ) + if arm.kind == ARM_KIND_ENV and stage != FRESH_CHILD_STAGE: + raise AblationSpecError( + f"arm {arm.name!r} needs --at-stage {FRESH_CHILD_STAGE!r}, not " + f"{stage!r}: the child manifest's services are provisioned " + "over the restored snapshot by the fresh child rollout, and " + f"by {stage!r} the parent's provisioned services survive the " + "fork — the arm would record the swap and run without it" + ) + if ( + stage == FRESH_CHILD_STAGE + and snapshot_layers is not None + and FRESH_CHILD_LAYER not in snapshot_layers + ): + raise AblationSpecError( + f"--at-stage {FRESH_CHILD_STAGE!r} needs the {FRESH_CHILD_LAYER!r} " + f"snapshot layer, got {sorted(frozenset(snapshot_layers))!r}: that " + "boundary precedes install_agent(), so every arm re-installs the " + "agent for itself and the container layer is what rolls one arm's " + "installation back before the next one runs" + ) + return stage + + +def resolve_ablation_task(tasks_dir: Path) -> Path: + """The one task an ablation runs, resolved from ``--tasks-dir``. + + An ablation compares arms *within* one task — the arms are the axis, the + task is fixed — so a directory holding several tasks is a request this + command cannot answer, not a batch to expand. It fails closed naming the + tasks it found. + """ + from benchflow.task.discovery import is_task_dir, resolve_task_collection_root + + path = Path(tasks_dir) + if not path.is_dir(): + raise AblationSpecError(f"--tasks-dir {str(path)!r} is not a directory") + root = resolve_task_collection_root(path) + if is_task_dir(root): + return root + tasks = sorted( + child for child in root.iterdir() if child.is_dir() and is_task_dir(child) + ) + if not tasks: + raise AblationSpecError( + f"no task found under --tasks-dir {str(path)!r} — a task directory " + "carries a task.md or task.toml" + ) + if len(tasks) > 1: + names = [task.name for task in tasks] + raise AblationSpecError( + f"--tasks-dir {str(path)!r} holds {len(names)} tasks ({names!r}) — " + "an ablation compares arms within one task; point --tasks-dir at " + "the task directory itself" + ) + return tasks[0] + + +def resolve_ablation_environment_binding( + task_path: Path, *, explicit: Path | str | None = None +) -> ManifestBinding | None: + """The environment the ablation binds — flag first, then the task's own. + + ``explicit`` is the ``--environment-manifest`` value (a manifest path or a + ``name@version`` registry spec): when given it wins outright and the + task-declared manifest is not even resolved — the same precedence ``bench + eval run`` applies (an explicit run-level manifest suppresses + ``manifest_from_task_document``). Otherwise the task's own ``task.md`` + declaration is resolved + (:func:`~benchflow.environment.manifest.manifest_binding_from_task_document`), + exactly as a normal evaluation resolves it; a stateful task ablated + without it would run the parent — and therefore every arm forked from it — + in a *different* environment than the run it is meant to explain. + + The returned :class:`~benchflow.environment.manifest.ManifestBinding` + keeps the ref verbatim and the manifest's content address, which is what + :func:`~benchflow.branch_report.environment_stamp` writes into + ``ablation.json``. Resolution failures are fatal rather than degrading to + ``None``: an ablation whose declared environment could not be built is not + an ablation that ran without services, and it fails before the parent run + costs anything. + """ + from benchflow._utils.text import describe_exception + from benchflow.environment.manifest import ( + load_manifest_binding, + manifest_binding_from_task_document, + ) + + if explicit is not None: + try: + return load_manifest_binding(explicit) + except Exception as exc: + raise AblationSpecError( + f"--environment-manifest {str(explicit)!r} does not resolve to " + f"an environment manifest: {describe_exception(exc)}" + ) from exc + try: + return manifest_binding_from_task_document(task_path) + except Exception as exc: + raise AblationSpecError( + f"{task_path.name} declares an environment manifest in its task.md " + f"that could not be resolved: {describe_exception(exc)}. Every arm " + "forks the parent's environment, so this ablation would compare " + "arms in a world the task says is the wrong one" + ) from exc diff --git a/src/benchflow/acp/runtime.py b/src/benchflow/acp/runtime.py index a4cd802a8..f145a3390 100644 --- a/src/benchflow/acp/runtime.py +++ b/src/benchflow/acp/runtime.py @@ -23,11 +23,13 @@ import logging from pathlib import Path +from benchflow._utils.scoring import REQUEST_GLOBAL_MARKER from benchflow.acp.client import ACPClient from benchflow.acp.container_transport import ContainerTransport from benchflow.acp.selection import selected_acp_transport from benchflow.acp.types import McpServerSpec from benchflow.acp.watchdog import IdleWatchdog +from benchflow.agents.errors import AgentProtocolError from benchflow.agents.protocol import ACPSessionAdapter from benchflow.agents.providers import ( find_provider, @@ -53,16 +55,77 @@ # import ``IdleTimeoutError`` from this module. The canonical definition # lives in :mod:`benchflow.diagnostics` (issue #503). __all__ = [ + "ACPRequestGlobalError", "AgentPromptTimeoutError", "IdleTimeoutError", "connect_acp", "execute_prompts", + "reasoning_effort_preflight_error", "selected_acp_transport", ] logger = logging.getLogger(__name__) +class ACPRequestGlobalError(RuntimeError): + """The agent rejected a *request-global* setting (model or effort). + + PR #1046 second review, P2-A: with Gemini and ``--reasoning-effort high`` + the run provisioned the environment, snapshotted, installed and connected + the agent, and only then hit this rejection — which retries and branch + children then re-provoked identically. The rejection is a property of the + request/agent pairing, not of the task, so it is non-task-attributable + and non-retryable. + + Every message carries :data:`~benchflow._utils.scoring.REQUEST_GLOBAL_MARKER` + so :func:`~benchflow._utils.scoring.classify_error` maps it to the + ``request_global`` category however the exception is re-stringified — the + provider_auth (#917) marker pattern. ``RuntimeError`` base keeps the + existing fail-closed callers and tests intact. + """ + + +def reasoning_effort_preflight_error( + agent: str, reasoning_effort: str | None +) -> str | None: + """Why ``reasoning_effort`` can never be applied for ``agent`` — or None. + + The static half of the P2-A fix: :func:`_configure_acp_session` decides + effort support from facts knowable before any provisioning — the + registry's ``acp_effort_config_id`` and the codex-acp effort-via-model-id + convention — so planning can fail the same way *before* a sandbox is + built and an agent installed. Kept next to the dispatch it mirrors: if + the runtime learns a new way to deliver effort, this predicate must learn + it in the same change. + + Returns ``None`` for agents the registry cannot vouch for (existence is + validated elsewhere; manifest-registered agents land in ``AGENTS`` with + their declared option ids) and for non-ACP agents, whose runs never reach + the ACP effort dispatch. + """ + if not reasoning_effort: + return None + agent_cfg = AGENTS.get(agent) + if agent_cfg is None: + return None + if agent_cfg.protocol != "acp": + return None + if agent == "codex-acp": + # Effort rides the ``model[effort]`` id, resolved against the live + # session's advertised catalog — not statically decidable. + return None + if agent_cfg.acp_effort_config_id: + return None + return ( + f"--reasoning-effort {reasoning_effort!r} is not supported by agent " + f"{agent!r}: it declares no ACP effort config option and does not " + "encode effort in its model ids, so the run would only be rejected " + "after the sandbox was built and the agent installed. Drop " + "--reasoning-effort or choose an agent that supports it (e.g. " + "claude-agent-acp, codex-acp)." + ) + + _ACP_CONNECT_MAX_RETRIES = 3 _ACP_CONNECT_BASE_DELAY = 2.0 _PROMPT_CANCEL_DRAIN_TIMEOUT_SEC = 0.25 @@ -390,6 +453,67 @@ def _resolve_acp_model_input(agent: str, model: str, agent_env: dict[str, str]) return agent_env.get(mapped_model_env, model) +def _codex_launch_config_model(agent_env: dict[str, str]) -> str | None: + """The model BenchFlow's own launch config injected into codex, if any. + + ``apply_codex_provider_config`` writes the gateway route into the + ``CODEX_CONFIG`` env var (``{"model": , "model_providers": ...}``), + and ``@agentclientprotocol/codex-acp`` applies it at startup — the session + then advertises that model as ``currentModelId``. + """ + import json + + from benchflow.agents.codex_config import CODEX_CONFIG_ENV + + raw = agent_env.get(CODEX_CONFIG_ENV) + if not raw: + return None + try: + config = json.loads(raw) + except json.JSONDecodeError: + return None + if not isinstance(config, dict): + return None + model = config.get("model") + return model if isinstance(model, str) and model else None + + +def _codex_current_model_id(session: object | None) -> str | None: + state = getattr(session, "model_state", None) + if not isinstance(state, dict): + return None + current = state.get("currentModelId") + return current if isinstance(current, str) and current else None + + +def _codex_model_owned_by_launch_config( + agent: str, + acp_model_id: str, + session: object | None, + agent_env: dict[str, str], +) -> bool: + """True when set_model is doomed AND the launch config already routed it. + + codex-acp@1.6.0 validates ``session/set_model`` strictly against its + built-in catalog: a model with no advertised ``model[effort]`` variant is + rejected outright — bare ids with "Unsupported format of modelId", + suffixed ids with "Unknown model", even when that exact id is already the + session's *current* model (verified live 2026-08-21 against + ``gpt-5.4-mini``, which 1.6.0 dropped from the catalog). When + ``_codex_session_model_id`` found no variant (the id stayed bare) the call + can only fail; when BenchFlow's own ``CODEX_CONFIG`` injection is what the + session already runs as current, the call is also *unnecessary* — the + gateway route is in place, so skip it instead of failing the rollout. + """ + if agent != "codex-acp" or "[" in acp_model_id: + return False + launch_model = _codex_launch_config_model(agent_env) + if launch_model is None: + return False + current = _codex_current_model_id(session) + return current is not None and _codex_model_name(current) == launch_model + + def _session_config_option_ids(session: object | None) -> set[str]: options = getattr(session, "config_options", None) if not isinstance(options, (list, tuple)): @@ -448,6 +572,13 @@ async def _set_acp_model( model_id, e, ) + if isinstance(e, AgentProtocolError): + # The agent answered and refused the model — deterministic, + # request-global (same split as _set_acp_config_option). + raise ACPRequestGlobalError( + f"{REQUEST_GLOBAL_MARKER}: model {model_id!r} via ACP for " + f"agent {agent!r}: {e}" + ) from e raise RuntimeError( f"Failed to set model {model_id!r} via ACP for agent {agent!r}: {e}" ) from e @@ -464,9 +595,12 @@ async def _set_acp_config_option( ) -> None: option_ids = _session_config_option_ids(session) if option_ids and config_id not in option_ids: - raise RuntimeError( - f"ACP agent {agent!r} does not expose {label} config option " - f"{config_id!r}; available options: {sorted(option_ids)!r}" + # The live session enumerates its config options and the wanted one + # is not among them — deterministic, request-global. + raise ACPRequestGlobalError( + f"{REQUEST_GLOBAL_MARKER}: ACP agent {agent!r} does not expose " + f"{label} config option {config_id!r}; available options: " + f"{sorted(option_ids)!r}" ) try: await asyncio.wait_for( @@ -481,6 +615,15 @@ async def _set_acp_config_option( value, e, ) + if isinstance(e, AgentProtocolError): + # The agent *answered* and refused the value — a deterministic + # rejection of a global setting. A timeout or transport failure + # proves nothing about compatibility and keeps the plain + # (retryable) RuntimeError below. + raise ACPRequestGlobalError( + f"{REQUEST_GLOBAL_MARKER}: ACP {label} config option " + f"{config_id!r}={value!r} for agent {agent!r}: {e}" + ) from e raise RuntimeError( f"Failed to set ACP {label} config option {config_id!r}=" f"{value!r} for agent {agent!r}: {e}" @@ -527,6 +670,26 @@ async def _configure_acp_session( value=acp_model_id, label="model", ) + elif _codex_model_owned_by_launch_config( + agent, acp_model_id, session, agent_env + ): + # codex-acp@1.6.0 rejects session/set_model for any model outside + # its built-in catalog, and BenchFlow's CODEX_CONFIG injection has + # already made the gateway route the session's current model — the + # call would fail the rollout to request a state that is already + # in place. A requested effort is satisfied only if the injected + # current id carries it; otherwise the effort step below fails + # closed exactly as before. + current_model_id = _codex_current_model_id(session) + logger.info( + f"Skipping ACP session/set_model for {agent} — launch config " + f"already owns the session model ({current_model_id!r})" + ) + effort_in_model_id = bool( + reasoning_effort + and current_model_id is not None + and _codex_reasoning_effort(current_model_id) == reasoning_effort + ) elif not agent_cfg or agent_cfg.supports_acp_set_model: # No model config option advertised/declared — use the legacy # session/set_model path. Fails closed if the agent rejects it. @@ -547,9 +710,10 @@ async def _configure_acp_session( ) return if not agent_cfg or not agent_cfg.acp_effort_config_id: - raise RuntimeError( - f"reasoning_effort={reasoning_effort!r} was requested for agent " - f"{agent!r}, but that agent does not declare an ACP effort config option" + raise ACPRequestGlobalError( + f"{REQUEST_GLOBAL_MARKER}: reasoning_effort={reasoning_effort!r} " + f"was requested for agent {agent!r}, but that agent does not " + "declare an ACP effort config option" ) await _set_acp_config_option( acp_client, diff --git a/src/benchflow/branch.py b/src/benchflow/branch.py index be9326d30..1a443c770 100644 --- a/src/benchflow/branch.py +++ b/src/benchflow/branch.py @@ -9,6 +9,12 @@ * ``restore`` — roll the environment back to a node's checkpoint. * ``aggregate`` — average the children's returns into V(node). +``checkpoint_composed`` / ``restore_composed`` are the layered variants from +the rollout-branching RFC (§3.1): they compose the container layer +(``Sandbox.snapshot``) with the environment-state layer into one +:class:`StageSnapshot`, in a fixed order — environment first on checkpoint, +sandbox first on restore. + These are the Branch *operations*. Wiring them into the rollout engine — running the forked children as sub-rollouts, quiescing the agent first — is the engine's job; these primitives stay pure and independently testable. @@ -16,14 +22,59 @@ from __future__ import annotations +from collections.abc import Sequence +from dataclasses import dataclass, field from typing import Any from benchflow.environment.protocol import StateSnapshot +from benchflow.sandbox.protocol import SandboxImage from benchflow.trajectories.tree import RolloutNode, RolloutTree, Step _SNAPSHOT_KEY = "snapshot" _REWARD_KEY = "reward" +#: Node-state key recording *why* a branch child produced no score. A node +#: carries either ``reward`` (it was scored) or this key (it was not) — never +#: both, and never a stand-in number for a score that does not exist. +UNSCORED_KEY = "unscored" + + +class UnscoredChildError(RuntimeError): + """A branch child ran to completion but produced no verifier reward. + + The one failure mode a scalar return cannot express: the child's rollout + did not raise, so there is nothing to propagate, but ``verify()`` yielded + no rewards — the verifier crashed, its output never reached the host, or + the reward file was never written. Returning ``0.0`` there is worse than + returning nothing: a lost reward becomes indistinguishable from a real + failing score, and downstream (``bench eval ablate``) reports it as an + observation. Child runners raise this instead; the engine records the + reason on the child node (:data:`UNSCORED_KEY`), leaves the node's + ``reward`` unset, and refuses to aggregate a value out of it. + """ + + def __init__(self, reason: str) -> None: + super().__init__(reason) + self.reason = reason + + +@dataclass(frozen=True) +class StageSnapshot: + """A composed checkpoint — one ref per snapshot layer (RFC §3.1). + + ``environment_ref`` and ``sandbox_ref`` each hold that layer's roll-back + handle iff the layer was requested at checkpoint time; ``None`` means the + layer was *not requested*, and :func:`restore_composed` rejects a live + object for it. ``stage`` optionally names the lifecycle boundary the + snapshot was taken at (``env-ready``, ``pre-verify``, ...); ``meta`` + carries diagnostics. + """ + + environment_ref: StateSnapshot | None + sandbox_ref: SandboxImage | None + stage: str | None = None + meta: dict[str, Any] = field(default_factory=dict) + async def checkpoint(node: RolloutNode, environment: Any) -> StateSnapshot: """Snapshot the environment at ``node`` and record the snapshot on it. @@ -36,6 +87,47 @@ async def checkpoint(node: RolloutNode, environment: Any) -> StateSnapshot: return snap +async def checkpoint_composed( + node: RolloutNode, + *, + environment: Any = None, + sandbox: Any = None, + stage: str | None = None, +) -> StageSnapshot: + """Snapshot the requested layers at ``node`` — the composed checkpoint. + + Layer order is fixed (RFC §3.1): ``environment.snapshot()`` first, then + ``sandbox.snapshot()``. A layer is included iff its live object is passed + (``None`` = layer not requested); quiescing the agent first is the + engine's job, not this op's. If either layer's snapshot raises, the error + propagates and *nothing* is recorded on ``node`` — a partial + :class:`StageSnapshot` is never a roll-back point. + """ + if environment is None and sandbox is None: + raise ValueError( + "checkpoint_composed() needs at least one layer — pass " + "environment=, sandbox=, or both" + ) + env_ref = await environment.snapshot() if environment is not None else None + sandbox_ref = await sandbox.snapshot() if sandbox is not None else None + snap = StageSnapshot(environment_ref=env_ref, sandbox_ref=sandbox_ref, stage=stage) + node.state[_SNAPSHOT_KEY] = snap + return snap + + +def adopt_checkpoint(node: RolloutNode, snapshot: StageSnapshot) -> RolloutNode: + """Record an already-taken :class:`StageSnapshot` as ``node``'s roll-back point. + + The stage-boundary policy (RFC §3.2) snapshots at a lifecycle boundary and + branches from it later; when it does, the recorded snapshot *is* the + node's checkpoint and must not be re-taken (the world has moved on since). + This is the one seam that writes a checkpoint the node did not take + itself, so the snapshot key stays private to this module. + """ + node.state[_SNAPSHOT_KEY] = snapshot + return node + + def fork(tree: RolloutTree, node: RolloutNode, n: int) -> list[RolloutNode]: """Fork ``node`` into ``n`` child continuations, making it a branch point.""" if n < 2: @@ -44,15 +136,78 @@ def fork(tree: RolloutTree, node: RolloutNode, n: int) -> list[RolloutNode]: async def restore(node: RolloutNode, environment: Any) -> None: - """Roll the environment back to ``node``'s recorded checkpoint.""" + """Roll the environment back to ``node``'s recorded checkpoint. + + Legacy environment-only restore: it accepts the bare ``StateSnapshot`` + recorded by :func:`checkpoint`. A :class:`StageSnapshot` recorded by + :func:`checkpoint_composed` is rejected *before* the environment is + touched — its layer refs are not an environment snapshot, and passing + one through would fail deep inside the environment implementation. + """ snap = node.state.get(_SNAPSHOT_KEY) if snap is None: raise ValueError( f"node {node.id!r} has no checkpoint — call checkpoint() before restore()" ) + if isinstance(snap, StageSnapshot): + raise ValueError( + f"node {node.id!r} was checkpointed with checkpoint_composed(); " + "use restore_composed() to roll back a composed StageSnapshot" + ) await environment.restore(snap) +async def restore_composed( + node: RolloutNode, + *, + environment: Any = None, + sandbox: Any = None, +) -> None: + """Roll the requested layers back to ``node``'s composed checkpoint. + + Restore order is the reverse of checkpoint (RFC §3.1): + ``sandbox.restore()`` first, then ``environment.restore()``. Accepts both + checkpoint shapes on ``node.state["snapshot"]`` — a legacy bare + :class:`StateSnapshot` recorded by :func:`checkpoint` (environment-only) + or a :class:`StageSnapshot`. Every layer present in the checkpoint must be + matched by its live object and vice versa; a mismatch is a caller bug and + raises ``ValueError`` before either layer is touched. + """ + snap = node.state.get(_SNAPSHOT_KEY) + if snap is None: + raise ValueError( + f"node {node.id!r} has no checkpoint — call checkpoint_composed() " + "before restore_composed()" + ) + if isinstance(snap, StateSnapshot): + # Legacy shape: checkpoint() recorded a bare env-state snapshot. + snap = StageSnapshot(environment_ref=snap, sandbox_ref=None) + elif not isinstance(snap, StageSnapshot): + raise ValueError( + f"node {node.id!r} holds an unrecognized checkpoint of type " + f"{type(snap).__name__!r} — expected StateSnapshot or StageSnapshot" + ) + # Validate both layers before restoring either — never a partial restore. + for layer, ref, live in ( + ("sandbox", snap.sandbox_ref, sandbox), + ("environment", snap.environment_ref, environment), + ): + if ref is not None and live is None: + raise ValueError( + f"node {node.id!r}'s checkpoint has a {layer} layer but no " + f"live {layer} was passed to restore_composed()" + ) + if ref is None and live is not None: + raise ValueError( + f"a live {layer} was passed to restore_composed() but node " + f"{node.id!r}'s checkpoint has no {layer} layer" + ) + if snap.sandbox_ref is not None: + await sandbox.restore(snap.sandbox_ref) + if snap.environment_ref is not None: + await environment.restore(snap.environment_ref) + + async def branch( tree: RolloutTree, node: RolloutNode, environment: Any, n: int ) -> list[RolloutNode]: @@ -61,14 +216,26 @@ async def branch( return fork(tree, node, n) -def aggregate(node: RolloutNode) -> float: +def aggregate(node: RolloutNode, *, over: Sequence[RolloutNode] | None = None) -> float: """V(node) — the mean of the children's returns. Each child carries its return in ``state["reward"]`` (written by the Reward plane after a child rollout is scored). Averaging them estimates the value of ``node``'s state — a reward function become a value function. + + ``over`` narrows the average to one fork's children. It matters once a + branch point can be a node that already has children: a stage-boundary + branch (RFC §3.2) forks the node the stage was captured on, which may + already carry the linear continuation — and that linear child has no + ``reward``, so averaging over ``node.children`` would silently drag V + toward zero. ``None`` keeps the every-child default. + + The engine never passes an *unscored* child (:data:`UNSCORED_KEY`) here: + a fork with one of those has no defined mean, so it records no value at + all rather than averaging a score nobody observed. """ - if not node.children: + children = list(node.children if over is None else over) + if not children: raise ValueError(f"node {node.id!r} has no children to aggregate") - returns = [float(child.state.get(_REWARD_KEY, 0.0)) for child in node.children] + returns = [float(child.state.get(_REWARD_KEY, 0.0)) for child in children] return sum(returns) / len(returns) diff --git a/src/benchflow/branch_artifacts.py b/src/benchflow/branch_artifacts.py new file mode 100644 index 000000000..593d65f61 --- /dev/null +++ b/src/benchflow/branch_artifacts.py @@ -0,0 +1,312 @@ +"""Keeping the parent's on-disk rollout artifacts intact across a branch. + +A branch child shares the parent's sandbox — that is the whole point: it +restores the parent's container and continues *that* world. But the sandbox +bind-mounts three of the parent rollout's own directories into the container +(``DockerSandbox.__init__``: ``host_agent_logs_path`` / ``host_artifacts_path`` +/ ``host_verifier_logs_path`` -> ``/logs/agent``, ``/logs/artifacts``, +``/logs/verifier``), and :meth:`DockerSandbox.restore` deliberately replays +those mounts so a restored container still writes where the host reads. So +every child writes its verifier output, its published trajectory and its +collected artifacts straight into **the parent's** directories. + +The result is not a wrong number — rewards are parsed in memory and persisted +per child — but it is wrong *evidence*: after a two-arm ablation the parent's +``verifier/`` holds the second arm's ``reward.txt`` / ``test-stdout.txt`` / +``judge_result.json``, and the parent's own scoring run has left no trace. It +is destructive rather than merely confusing, because a child whose rollout +paths are its own (a fresh-rollout child) sees ``has_host_mount() == False`` +and therefore runs ``clear_verifier_output_dir``, whose +``find /logs/verifier -mindepth 1 -delete`` empties the parent's directory +before the child writes anything of its own. + +This module is the fix, and it is deliberately mechanical: hold the parent's +entries aside for the duration of the fork, hand each child the entries it +produced, then put the parent's back. Layout under the run directory:: + + / + agent/ artifacts/ verifier/ # the parent's own, restored at the end + branches// + parent/{agent,artifacts,verifier}/ # transient: the parent's entries, + # held only while children run + children// + mounted/{agent,artifacts,verifier}/ # what THIS child wrote to the + # shared mounts + +Two properties worth stating. The mount roots are never removed or recreated, +only emptied entry by entry — deleting a live bind-mount source detaches it +from the running container. And ``parent/`` is transient: it exists on disk +only between the first child starting and the last child finishing, so finding +it after a run means the branch did not complete and the parent's evidence is +in there — preserved, never deleted — rather than at its canonical path. + +Custody fails **closed**: this is audit-critical evidence, so a failure at any +step of the chain preserves the hold directory (nothing is ever removed with +``rmtree`` unless every held entry was confirmed moved back) and surfaces as a typed +:class:`ArtifactCustodyError` whose message names the preserved path. The one +concession to the engine's "an artifact error must never replace the real +failure" invariant is *when* the error is raised: a custody failure observed +while a child's own exception is propagating is recorded on the guard and +logged — with the preserved path — instead of raised, so the child's failure +stays the caller's diagnosis (:meth:`MountedArtifacts.hand_off` records; +:meth:`MountedArtifacts.raise_pending` and :meth:`MountedArtifacts.release` +raise it as soon as no more important exception is in flight). +""" + +from __future__ import annotations + +import logging +import shutil +from dataclasses import dataclass, field +from pathlib import Path + +from benchflow.task.paths import RolloutPaths + +logger = logging.getLogger(__name__) + + +class ArtifactCustodyError(RuntimeError): + """A step of the mounted-artifact custody chain failed. + + Fail-closed marker for :class:`MountedArtifacts`: the hold directory named + by :attr:`hold_dir` has been **preserved on disk** — whatever evidence the + failed move left behind is in there, recoverable by hand — and the branch + must not pretend the custody chain completed. + """ + + def __init__(self, message: str, *, hold_dir: Path) -> None: + super().__init__(message) + self.hold_dir = Path(hold_dir) + + +#: Directory name under ``branches//`` holding the parent's own +#: mounted-artifact entries while its children run. Transient — see the module +#: docstring. +PARENT_HOLD_DIRNAME = "parent" + +#: Directory name under a child's artifact directory holding what that child +#: wrote to the shared mounts. A dedicated subdirectory, so it can never +#: collide with the ``verifier/`` / ``agent/`` a fresh-rollout child downloads +#: into its own run directory. +CHILD_MOUNT_DIRNAME = "mounted" + + +def mounted_artifact_dirs(rollout_dir: Path) -> tuple[Path, ...]: + """The rollout directories a sandbox bind-mounts into the container. + + Derived from :class:`~benchflow.task.paths.RolloutPaths` rather than + hard-coded, because these are exactly the three paths ``DockerSandbox`` + passes as ``host_*_path`` when it builds the compose environment: a fourth + mount added there shows up here by construction. + """ + paths = RolloutPaths(rollout_dir=Path(rollout_dir)) + return (paths.agent_dir, paths.artifacts_dir, paths.verifier_dir) + + +def parent_hold_dir(run_dir: Path, parent_id: str) -> Path: + """Where the parent's mounted entries wait out the fork.""" + return Path(run_dir) / "branches" / parent_id / PARENT_HOLD_DIRNAME + + +def child_mount_dir(run_dir: Path, parent_id: str, child_id: str) -> Path: + """Where one child's mounted output is kept, under its own artifact dir.""" + from benchflow.branch_lineage import branch_child_dir + + return branch_child_dir(Path(run_dir), parent_id, child_id) / CHILD_MOUNT_DIRNAME + + +def child_artifact_roots(child_dir: Path) -> tuple[Path, ...]: + """Rollout-shaped roots holding one child's own output, best first. + + A branch child's verifier output lands in one of two places, and *which* + one depends on how the engine ran the child — which a report-time reader + should not have to know: + + * a **fresh-rollout** child (every child of ``env-ready``) has a run + directory of its own, and that run directory *is* its artifact + directory, so its verifier output is downloaded to ``/verifier``; + * an **in-place** child (``pre-verify`` / ``post-verify``) never had one. + It wrote through the parent's bind mounts, and :meth:`MountedArtifacts.hand_off` + archived what it wrote under ``/mounted/`` — same + ``agent``/``artifacts``/``verifier`` shape, one level down. + + Both roots are :class:`~benchflow.task.paths.RolloutPaths`-shaped, so a + reader tries them in order and takes the first that carries what it wants. + Order is the child's *own* directory first: a fresh-rollout child can have + both (it writes through the restored parent mounts too, and the hand-off + archives that), and the copy it downloaded into its run directory is the + authoritative one. + + This is the one place that knows the layout — ``bench eval ablate`` reads + per-test outcomes through it rather than re-deriving ``mounted/`` for + itself, so a change to :data:`CHILD_MOUNT_DIRNAME` cannot leave a second + hard-coded path behind. + """ + child_dir = Path(child_dir) + return (child_dir, child_dir / CHILD_MOUNT_DIRNAME) + + +def _move_entries(source: Path, target: Path) -> int: + """Move every entry of ``source`` into ``target``; return how many moved. + + ``source`` itself is left in place and empty — it is a live bind-mount + source, and removing it would detach the mount from the running container. + An entry already present at the destination is replaced, so a retry is + idempotent rather than a name collision. + """ + if not source.is_dir(): + return 0 + entries = sorted(source.iterdir()) + if not entries: + return 0 + target.mkdir(parents=True, exist_ok=True) + for entry in entries: + destination = target / entry.name + if destination.is_symlink() or destination.is_file(): + destination.unlink() + elif destination.is_dir(): + shutil.rmtree(destination) + shutil.move(str(entry), str(destination)) + return len(entries) + + +@dataclass +class MountedArtifacts: + """Custody of the parent's mounted artifact entries during a branch. + + Created by :meth:`hold` before the first child runs. :meth:`hand_off` is + called after each child, moving whatever is now in the mounts into that + child's own directory — which both preserves the child's evidence and + leaves the mounts empty, so the next child cannot inherit the previous + one's files. :meth:`release` puts the parent's entries back at the end. + + Every step fails closed (:class:`ArtifactCustodyError`, module docstring): + the hold directory is preserved on any failure, and the only step allowed + to defer its raise is :meth:`hand_off`, which runs inside the child loop's + ``finally`` where a raise would replace the child's own exception — it + records into :attr:`custody_failures` instead, and :meth:`raise_pending` + surfaces the record the moment no child failure is in flight. + """ + + rollout_dir: Path + hold_dir: Path + held: list[str] = field(default_factory=list) + #: Custody failures observed where raising would have masked a more + #: important exception (``hand_off`` inside the child loop's ``finally``). + custody_failures: list[str] = field(default_factory=list) + + @classmethod + def hold(cls, *, run_dir: Path, parent_id: str) -> MountedArtifacts: + """Move the parent's mounted entries aside, before any child runs. + + Fails closed: a failure here means a child would overwrite the + parent's evidence, so the branch must not proceed. Whatever was moved + before the failure stays preserved in the hold directory — named in + the raised :class:`ArtifactCustodyError` — and is never deleted. + """ + guard = cls( + rollout_dir=Path(run_dir), hold_dir=parent_hold_dir(run_dir, parent_id) + ) + for source in mounted_artifact_dirs(guard.rollout_dir): + try: + if _move_entries(source, guard.hold_dir / source.name): + guard.held.append(source.name) + except Exception as exc: + message = ( + f"branch could not hold the parent's {source} directory " + f"aside ({exc}); refusing to run children over the " + "parent's evidence. Entries already held are preserved at " + f"{guard.hold_dir}" + ) + logger.error(message, exc_info=True) + raise ArtifactCustodyError(message, hold_dir=guard.hold_dir) from exc + return guard + + def hand_off(self, target: Path) -> None: + """Move what the child just wrote to the mounts into ``target``. + + Runs inside the child loop's ``finally`` — a raise here would replace + the child's own exception, so a failure is recorded on + :attr:`custody_failures` (and logged, naming the preserved hold + directory) for :meth:`raise_pending` / the caller to surface once the + child's exception, if any, has been dealt with. + """ + for source in mounted_artifact_dirs(self.rollout_dir): + try: + _move_entries(source, Path(target) / source.name) + except Exception as exc: + message = ( + f"branch could not capture a child's {source.name} output " + f"under {target} ({exc}); files left in the mounts would " + "leak into the next child. The parent's own evidence is " + f"preserved at {self.hold_dir}" + ) + self.custody_failures.append(message) + logger.error(message, exc_info=True) + + def raise_pending(self) -> None: + """Raise the custody failure :meth:`hand_off` had to swallow, if any. + + Called by the child loop after a child completes *without* raising — + the point where no more important exception exists to preserve. + """ + if self.custody_failures: + raise ArtifactCustodyError( + "; ".join(self.custody_failures), hold_dir=self.hold_dir + ) + + def release(self, *, raising: bool = True) -> None: + """Put the parent's own entries back at their canonical paths. + + The hold directory is removed only after **every** held entry was + confirmed moved back and nothing remains inside it — a directory whose + contents were not confirmed moved is never deleted. On failure the + hold directory is preserved and an :class:`ArtifactCustodyError` + naming it is raised; ``raising=False`` (the caller is unwinding a + child's own exception, which must stay the diagnosis) records and + logs instead of raising. + """ + failures: list[str] = [] + for source in mounted_artifact_dirs(self.rollout_dir): + if source.name not in self.held: + continue + try: + _move_entries(self.hold_dir / source.name, source) + except Exception as exc: + failures.append( + f"branch could not restore the parent's {source} directory ({exc})" + ) + logger.error( + "branch could not restore the parent's %s directory from " + "%s — the parent's evidence is preserved there, not lost", + source, + self.hold_dir, + exc_info=True, + ) + remaining = ( + sorted( + str(entry.relative_to(self.hold_dir)) + for entry in self.hold_dir.rglob("*") + if entry.is_symlink() or not entry.is_dir() + ) + if self.hold_dir.is_dir() + else [] + ) + if failures or remaining: + if not failures: + # Every move-back "succeeded" yet files remain: unaccounted + # evidence. Deleting it would be exactly the fail-open rmtree + # this class exists to prevent. + failures.append( + f"unexpected entries remain after release: {remaining[:10]}" + ) + message = ( + "; ".join(failures) + + f" — the hold directory is preserved at {self.hold_dir}" + ) + self.custody_failures.extend(failures) + logger.error(message) + if raising: + raise ArtifactCustodyError(message, hold_dir=self.hold_dir) + return + shutil.rmtree(self.hold_dir, ignore_errors=True) diff --git a/src/benchflow/branch_children.py b/src/benchflow/branch_children.py new file mode 100644 index 000000000..d3d9fd589 --- /dev/null +++ b/src/benchflow/branch_children.py @@ -0,0 +1,154 @@ +"""How a branch child executes its delta — the runner side of the engine. + +One module answers "given this child's delta and this fork's boundary, what +actually runs?" (RFC §3.3). There are exactly two execution paths, and the +boundary — not the delta — picks between them: + +* **in place** — the ordinary child at a post-installation boundary: a fresh + agent session over the restored world, driven on the parent Rollout instance + (:func:`make_default_runner`). An ``injected_prompt`` delta binds here as the + child's user-visible continuation prompt. +* **fresh rollout** — every engine-run child of ``env-ready``: that boundary + precedes ``install_agent()``, so the restored world has no agent to connect + to and the child re-runs installation as its own Rollout over the restored + sandbox. The implementation lives in :mod:`benchflow.branch_skill` (its + original reason to exist — the skills ablation — named the module, and six + test files plus the ``run_fresh_child`` monkeypatch seam pin that import + path); this module re-exports its API so delta-execution callers have one + home to import from. + +:func:`select_child_runner` is the one place the choice is made — the branch +transaction loop asks it per child, so the boundary rule ("fresh at env-ready, +in place elsewhere, caller-supplied runners are left alone") cannot fork +between call sites. + +The gates deciding whether a delta may execute at all live in +:mod:`benchflow.branch_policy` and fail closed before anything is quiesced. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING + +from benchflow.branch import UnscoredChildError + +# The fresh-rollout execution path (module of record: benchflow.branch_skill). +from benchflow.branch_skill import ( # noqa: F401 + EXECUTION_FRESH_ROLLOUT, + FRESH_CHILD_LAYER, + FRESH_CHILD_STAGE, + SKILL_DELTA_LAYER, + SKILL_DELTA_STAGE, + BranchEnvironmentImageConflict, + child_skill_config, + fresh_child_skill_mode, + make_fresh_child_runner, + provision_child_environment, + resolve_child_skill_policy, + resolve_environment_ref_delta, + run_fresh_child, +) + +if TYPE_CHECKING: + from pathlib import Path + + from benchflow.branch_delta import BranchDelta + from benchflow.rollout import Rollout + from benchflow.trajectories.tree import RolloutNode + +# The per-child runner: given the child's branch node, run its continuation and +# return the scalar return. No ``int`` index — a caller that needs per-child +# prompts binds them into a closure (see ``select_child_runner``). +ChildRunner = Callable[["RolloutNode"], Awaitable[float]] + + +def make_default_runner( + rollout: Rollout, *, prompts: list[str] | None = None +) -> ChildRunner: + """Build the default per-child runner bound to ``rollout``. + + The default runner re-runs the child from the parent's env checkpoint with + a *fresh agent session* — agent-session snapshot is the unsolved hard part + (``docs/architecture.md``, "The hard part"), so the agent restarts per + child. Each child connects a fresh agent and disconnects it at the end, so + no two children's agents overlap (the next child connects only after the + previous one disconnected). ``verify()`` returning ``None``, an empty dict, + or a dict with a ``None`` reward raises + :class:`~benchflow.branch.UnscoredChildError`: the child ran but was never + scored, and a fabricated ``0.0`` would read as a real failing score. + + ``prompts`` — the child's continuation prompts; ``None`` keeps the + rollout's resolved prompts. An ``injected_prompt`` delta (RFC §3.3) binds + here, so the injection is the child's user-visible first message — never + silently merged into other prompt content (#908). + """ + + async def _runner(child: RolloutNode) -> float: + await rollout.connect() + # Fill the pending branch-child node in place — the continuation Step + # lands on `child` itself, no content-free placeholder. + await rollout.execute(prompts, node=child) + rewards = await rollout.verify() + await rollout.disconnect() + if not rewards or rewards.get("reward") is None: + raise UnscoredChildError( + f"branch child {child.id} produced no verifier reward " + f"(verify() returned {rewards!r})" + _verifier_error_suffix(rollout) + ) + return float(rewards["reward"]) + + return _runner + + +def select_child_runner( + rollout: Rollout, + *, + delta: BranchDelta | None, + default: ChildRunner, + run_child: ChildRunner | None, + fresh_children: bool, + parent: RolloutNode, + run_dir: Path | None, + fresh_runner_factory: Callable[..., ChildRunner] = make_fresh_child_runner, +) -> ChildRunner: + """Pick the runner one child executes under — the boundary rule, per child. + + Every case was validated by :mod:`benchflow.branch_policy` to not combine + with an explicit ``run_child``. A child forked from ``env-ready`` runs as + a fresh Rollout over the just-restored sandbox — that boundary precedes + ``install_agent()``, so there is no agent in the restored world to connect + to and no installed skills to keep; the child re-runs installation for + itself, under the delta's skill_mode when it has one and under the + parent's own recorded mode when it does not. Everywhere else the agent is + already installed and the child runs in place, with an injected-prompt + delta binding the child's continuation prompt into a per-child default + runner — the formalized version of the caller's per-child-prompt closure. + + ``fresh_runner_factory`` defaults to :func:`make_fresh_child_runner`; the + orchestrator passes its own module global through so + ``benchflow.rollout_branch`` remains the patch seam for faking the + fresh-child path. + """ + if fresh_children: + return fresh_runner_factory( + rollout, + delta=delta, + parent=parent, + branch_stage=FRESH_CHILD_STAGE, + run_dir=run_dir, + ) + if run_child is None and delta is not None and delta.injected_prompt is not None: + return make_default_runner(rollout, prompts=[delta.injected_prompt]) + return default + + +def _verifier_error_suffix(rollout: Rollout) -> str: + """`` — `` when the rollout recorded one, else ``''``. + + The verifier's own diagnostic is the difference between "the agent scored + nothing" and "the score never reached the host", so it is carried into the + unscored reason verbatim rather than left in the log. + """ + error = getattr(rollout, "_verifier_error", None) + return f" — {error}" if error else "" diff --git a/src/benchflow/branch_delta.py b/src/benchflow/branch_delta.py new file mode 100644 index 000000000..c5682e482 --- /dev/null +++ b/src/benchflow/branch_delta.py @@ -0,0 +1,114 @@ +"""Per-child branch deltas — the recorded change a branch child runs under. + +A branch child's delta is a tuple over the three run-level variation axes plus +an injected prompt (rollout-branching RFC §3.3). Every member reuses an +existing, content-addressed mechanism: ``environment_ref`` is an S-axis +registry ref, ``config_override`` a C-axis allowlisted patch hashed exactly +like the run-level overlay (#790), ``skill_mode`` the install-time skills +toggle, and ``injected_prompt`` an explicit, recorded first message — never a +silent injection (#908). + +:class:`BranchDelta` is the *schema*: all four fields exist now so the +artifact format is stable. Which fields the branch engine executes is the +engine's contract (:mod:`benchflow.rollout_branch`); provenance hashes raw +content (the prompt text never appears in artifacts, only its digest). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from benchflow._utils.config_override import overlay_hash +from benchflow._utils.content_address import sha256_prefixed +from benchflow.skill_policy import SKILL_MODE_NO_SKILL, SKILL_MODE_WITH_SKILL + +# The skill modes a branch child may flip to. ``self-gen`` is a run-level mode, +# not a per-child ablation axis, so it is not branchable. +_BRANCH_SKILL_MODES = frozenset({SKILL_MODE_NO_SKILL, SKILL_MODE_WITH_SKILL}) + + +class BranchDeltaNotSupported(NotImplementedError): + """A BranchDelta a branch request cannot execute as asked (fail closed). + + Every executable field has preconditions — ``skill_mode`` / + ``config_override`` / ``environment_ref`` run only as fresh child rollouts + from the ``env-ready`` stage snapshot, ``skill_mode`` additionally needs a + ``no-skill`` parent, and ``environment_ref`` needs a same-image, + framework-started-services manifest pair (the outage pattern) — and a + request that does not satisfy them fails closed *before any child runs*, + rather than running a child that silently ignores its delta or measures a + world its restore could not roll back. Defined here, beside the schema, + so both the branch engine (:mod:`benchflow.rollout_branch`) and the + fresh-child module (:mod:`benchflow.branch_skill`) can raise and subclass + it without a circular import; the engine re-exports it. + """ + + +@dataclass(frozen=True) +class BranchDelta: + """The exactly-one-controlled-change a branch child runs under (RFC §3.3). + + All fields default to ``None`` (= inherit the parent's value); an + all-``None`` delta is the zero-delta child, byte-for-byte today's branch + behavior. ``skill_mode`` is validated against the branchable modes + (``no-skill`` / ``with-skill``) at construction — a bad mode fails closed + here, not at child run time. + """ + + environment_ref: str | None = None + config_override: dict[str, Any] | None = None + skill_mode: str | None = None + injected_prompt: str | None = None + + def __post_init__(self) -> None: + if self.skill_mode is not None and self.skill_mode not in _BRANCH_SKILL_MODES: + raise ValueError( + f"skill_mode must be one of {sorted(_BRANCH_SKILL_MODES)} when " + f"set, got {self.skill_mode!r}" + ) + + @property + def is_empty(self) -> bool: + """True iff every field is unset — the zero-delta child.""" + return ( + self.environment_ref is None + and self.config_override is None + and self.skill_mode is None + and self.injected_prompt is None + ) + + def provenance_dict(self) -> dict[str, Any]: + """The delta as recorded in lineage artifacts (RFC §3.4). + + Small literal fields (``environment_ref``, ``skill_mode``) are recorded + verbatim; content-bearing fields are recorded as sha256 content + addresses only — ``config_override`` hashed over its canonical JSON + (``sort_keys=True``, the run-level overlay's exact hash) and + ``injected_prompt`` over its UTF-8 text. No raw prompt text ever + appears in provenance. Unset fields serialize as ``null``. + + A set ``config_override`` additionally records its *top-level keys* + (sorted) as ``config_override_keys`` — the same ``keys`` + ``sha256`` + record the run-level overlay leaves in ``config.json`` (#790), so a + reader can see *which* allowlisted sections the arm patched without + the hash preimage. The key is present only when the field is set, so + every other delta keeps the provenance shape it has always had. + """ + provenance: dict[str, Any] = { + "environment_ref": self.environment_ref, + "config_override_sha256": ( + overlay_hash(self.config_override) + if self.config_override is not None + else None + ), + "skill_mode": self.skill_mode, + "injected_prompt_sha256": ( + sha256_prefixed(self.injected_prompt.encode()) + if self.injected_prompt is not None + else None + ), + } + if self.config_override is not None: + provenance["config_override_keys"] = sorted(self.config_override) + return provenance diff --git a/src/benchflow/branch_lineage.py b/src/benchflow/branch_lineage.py new file mode 100644 index 000000000..caa868f51 --- /dev/null +++ b/src/benchflow/branch_lineage.py @@ -0,0 +1,313 @@ +"""Branch lineage artifacts — ``tree.json``, per-child provenance, stages. + +Today ``RolloutTree`` lives and dies in memory; branching must leave the same +quality of evidence as a linear run. This module serializes the tree to a +deterministic ``tree.json`` in the run folder, writes the stage-snapshot +registry (RFC §3.2) as ``stage_snapshots.json``, and builds the +``kind="benchflow-branch"`` source-provenance dict each branch child carries — +the same seam ``benchflow-continue`` uses. + +Pure writers only: no engine state, no wall-clock timestamps (determinism is a +test guarantee — goldens pin the output byte-for-byte). Failure isolation is +the caller's job: the engine wraps these writes so an artifact error never +corrupts branch results. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from benchflow.branch import UNSCORED_KEY, StageSnapshot +from benchflow.branch_delta import BranchDelta +from benchflow.environment.protocol import StateSnapshot +from benchflow.trajectories.tree import RolloutNode, RolloutTree + +_SCHEMA_VERSION = 1 +_SNAPSHOT_KEY = "snapshot" +_REWARD_KEY = "reward" +_VALUE_KEY = "value" +_DELTA_KEY = "delta" +_EXECUTION_KEY = "delta_execution" + + +def _snapshot_refs(snap: Any) -> dict[str, str | None] | None: + """Serializable per-layer refs of a recorded checkpoint (either shape). + + A legacy bare :class:`StateSnapshot` is an environment-only checkpoint; a + :class:`StageSnapshot` carries one ref per requested layer. ``None`` in + and anything unrecognized out map to ``None`` — the serializer records, + it does not validate. + """ + if isinstance(snap, StateSnapshot): + return {"environment": snap.id, "sandbox": None} + if isinstance(snap, StageSnapshot): + return { + "environment": ( + snap.environment_ref.id if snap.environment_ref is not None else None + ), + "sandbox": snap.sandbox_ref.ref if snap.sandbox_ref is not None else None, + } + return None + + +def _provenance(delta: BranchDelta | dict[str, Any] | None) -> dict[str, Any]: + """A delta's provenance dict; ``None`` records as the zero delta. + + Accepts an already-serialized provenance dict verbatim — the branch + engine records each child's delta provenance on the node at fork time + (``node.state["delta"]``), and the writers here pass it through. + """ + if delta is None: + return BranchDelta().provenance_dict() + if isinstance(delta, BranchDelta): + return delta.provenance_dict() + return delta + + +def serialize_tree( + tree: RolloutTree, + *, + run_dir: Path, + cut_point: dict[str, Any] | None = None, +) -> Path: + """Write ``/tree.json`` — the deterministic lineage artifact. + + Every node serializes from its *own* recorded state: id / parent / stage + tag (the recorded checkpoint's ``stage``) / snapshot refs, plus + ``reward``, ``value``, ``delta`` (the provenance dict the engine attached + at fork time) and ``delta_execution`` (how the engine ran the child — + absent for the ordinary in-place child) when present on the node. Nothing + is inferred from position — a tree with several branch points, or several branch + events at the same parent, serializes each node's provenance exactly as + recorded. ``cut_point`` is recorded at the top level. Output is + deterministic — sorted keys, indented, trailing newline, no wall-clock + timestamps. + """ + nodes_payload: list[dict[str, Any]] = [] + for node in tree.nodes(): + snap = node.state.get(_SNAPSHOT_KEY) + entry: dict[str, Any] = { + "id": node.id, + "parent": node.parent.id if node.parent is not None else None, + "stage": getattr(snap, "stage", None), + "snapshot": _snapshot_refs(snap), + } + if _REWARD_KEY in node.state: + entry["reward"] = float(node.state[_REWARD_KEY]) + if UNSCORED_KEY in node.state: + # A node carries a reward or a reason it has none — never a + # stand-in number. Serializing the reason keeps the lineage + # honest about which children the run actually observed. + entry[UNSCORED_KEY] = str(node.state[UNSCORED_KEY]) + if _VALUE_KEY in node.state: + entry["value"] = float(node.state[_VALUE_KEY]) + if _DELTA_KEY in node.state: + entry["delta"] = node.state[_DELTA_KEY] + if _EXECUTION_KEY in node.state: + entry[_EXECUTION_KEY] = node.state[_EXECUTION_KEY] + nodes_payload.append(entry) + + payload = { + "schema_version": _SCHEMA_VERSION, + "cut_point": cut_point, + "nodes": nodes_payload, + } + path = Path(run_dir) / "tree.json" + path.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n") + return path + + +def _stage_exchange_index(snap: Any) -> int | None: + """The recorded completed-exchange index of a stage capture, or ``None``. + + ``capture_stage`` stores the count in ``snapshot.meta`` (RFC §3.5 — + stage-tagged replay cuts); the serializer records exactly an ``int`` and + maps anything else — absent meta, an unavailable count, a foreign value — + to ``None``, the honest "unknown index". + """ + meta = getattr(snap, "meta", None) + value = meta.get("exchanges_completed") if isinstance(meta, Mapping) else None + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + + +def stage_snapshots_payload( + snapshots: Mapping[str, StageSnapshot], +) -> dict[str, dict[str, Any]]: + """The stage registry as a serializable ``stage -> refs`` mapping (RFC §3.2). + + Each entry records the per-layer roll-back handles, the ``layers`` the + stage actually captured (sorted) — the set a stage branch restores — and + ``exchanges_completed``, the LLM-exchange index recorded at capture time + (``null`` when it was unavailable): the stage→exchange-index data a + stage-named replay cut (``bench eval continue --cut-stage``) resolves. + """ + payload: dict[str, dict[str, Any]] = {} + for stage, snap in snapshots.items(): + refs = _snapshot_refs(snap) or {"environment": None, "sandbox": None} + payload[stage] = { + "environment_ref": refs["environment"], + "sandbox_ref": refs["sandbox"], + "layers": sorted(layer for layer, ref in refs.items() if ref is not None), + "exchanges_completed": _stage_exchange_index(snap), + } + return payload + + +def write_stage_snapshots( + *, run_dir: Path, snapshots: Mapping[str, StageSnapshot] +) -> Path: + """Write ``/stage_snapshots.json`` — the recorded stage registry. + + Deterministic like every other lineage artifact: sorted keys, indented, + trailing newline, no wall-clock timestamps. Rewritten in full on each + capture so a run that dies mid-lifecycle still leaves the stages it did + reach. The caller isolates failures. + """ + return write_stage_snapshots_stages( + run_dir=run_dir, stages=stage_snapshots_payload(snapshots) + ) + + +def write_stage_snapshots_stages( + *, run_dir: Path, stages: Mapping[str, Mapping[str, Any]] +) -> Path: + """Write ``stage_snapshots.json`` from an already-serialized stage mapping. + + The seam :func:`write_stage_snapshots` (capture-time, from the live + registry) and the cleanup-time lifetime rewrite (annotated entries — + ``ephemeral`` / ``exported``, see + :func:`benchflow.branch_policy.finalize_stage_snapshots`) share, so both + writers emit one deterministic schema. + """ + path = Path(run_dir) / "stage_snapshots.json" + payload = { + "schema_version": _SCHEMA_VERSION, + "stages": {stage: dict(entry) for stage, entry in stages.items()}, + } + path.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n") + return path + + +def annotate_stage_snapshots_file( + *, run_dir: Path, annotations: Mapping[str, Mapping[str, Any]] +) -> Path | None: + """Merge lifetime-annotated entries into an existing ``stage_snapshots.json``. + + ``annotations`` maps stage name → the fully annotated entry (refs plus + ``ephemeral`` / ``exported`` / ``export_error``), which *replaces* the + file's entry for that stage — the annotated entry was built from the same + recorded refs, plus the lifetime the annotator just decided. Returns + ``None`` (no-op) when the file does not exist: there is no recorded + artifact to make truthful. A corrupt file raises; the caller isolates + failures, as for every writer in this module. + """ + path = Path(run_dir) / "stage_snapshots.json" + if not path.exists(): + return None + payload = json.loads(path.read_text()) + stages = payload.get("stages") + if not isinstance(stages, dict): + stages = {} + payload["stages"] = stages + for stage, entry in annotations.items(): + stages[stage] = dict(entry) + path.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n") + return path + + +def child_provenance( + parent_ref: str, + *, + branch_stage: str, + snapshot: StageSnapshot | StateSnapshot | None, + delta: BranchDelta | dict[str, Any] | None, + cut_point: dict[str, Any] | None = None, + delta_execution: str | None = None, +) -> dict[str, Any]: + """The ``kind="benchflow-branch"`` source-provenance dict (RFC §3.4). + + The branch twin of the ``kind="benchflow-continue"`` provenance a + continued run carries: which rollout the child forked from, at which stage + boundary, from which snapshot refs, and under which (content-addressed) + delta — either a :class:`BranchDelta` or the provenance dict the engine + recorded on the child node at fork time. ``cut_point`` stays ``None`` + until the replay cut-point API lands. + + ``delta_execution`` names how the engine executed the delta when that is + not the ordinary in-place child — ``"fresh-rollout"`` for a ``skill_mode`` + child, which re-runs installation as its own Rollout over the restored + snapshot. The key is written only when given, so a child run in place + keeps the exact dict shape it has always had. + """ + refs = _snapshot_refs(snapshot) or {"environment": None, "sandbox": None} + provenance: dict[str, Any] = { + "kind": "benchflow-branch", + "parent_rollout": parent_ref, + "branch_stage": branch_stage, + "snapshot_ref": { + "sandbox": refs["sandbox"], + "environment": refs["environment"], + }, + "cut_point": cut_point, + "delta": _provenance(delta), + } + if delta_execution is not None: + provenance[_EXECUTION_KEY] = delta_execution + return provenance + + +def branch_child_dir(run_dir: Path, parent_id: str, child_id: str) -> Path: + """The per-child artifact directory of one branch event (RFC §3.4). + + Namespaced by node id (unique within the tree) so a second branch event — + at the same parent or a different one — can never overwrite an earlier + event's artifacts. A child the engine ran as its own Rollout uses this same + directory as its run directory, so the standard ``config.json`` / + ``result.json`` / trajectory files land beside the branch's own + ``provenance.json``. + """ + return Path(run_dir) / "branches" / parent_id / "children" / child_id + + +def write_branch_artifacts( + *, + run_dir: Path, + tree: RolloutTree, + parent: RolloutNode, + children: Sequence[RolloutNode], +) -> None: + """Write a completed branch's lineage artifacts under ``run_dir``. + + ``tree.json`` at the top level; per child, :func:`branch_child_dir` holding + ``provenance.json`` and — when the child recorded a return — + ``reward.json``. Each child's delta provenance and execution mode are read + from ``child.state`` (attached by the engine at fork time), never inferred + positionally. The caller isolates failures: any exception here must be + caught and logged so an artifact-write error never corrupts branch results. + """ + snapshot = parent.state.get(_SNAPSHOT_KEY) + stage = getattr(snapshot, "stage", None) + branch_stage = stage if stage is not None else f"cursor:{parent.id}" + serialize_tree(tree, run_dir=run_dir) + for child in children: + child_dir = branch_child_dir(run_dir, parent.id, child.id) + child_dir.mkdir(parents=True, exist_ok=True) + provenance = child_provenance( + str(run_dir), + branch_stage=branch_stage, + snapshot=snapshot, + delta=child.state.get(_DELTA_KEY), + delta_execution=child.state.get(_EXECUTION_KEY), + ) + (child_dir / "provenance.json").write_text( + json.dumps(provenance, sort_keys=True, indent=2) + "\n" + ) + if _REWARD_KEY in child.state: + (child_dir / "reward.json").write_text( + json.dumps({"reward": float(child.state[_REWARD_KEY])}) + "\n" + ) diff --git a/src/benchflow/branch_policy.py b/src/benchflow/branch_policy.py new file mode 100644 index 000000000..6bbc6a610 --- /dev/null +++ b/src/benchflow/branch_policy.py @@ -0,0 +1,843 @@ +"""Stage and snapshot policy — what a branch may capture and fork, and when. + +The fail-closed half of the branch engine (RFC §3.1–§3.3): every rule that can +be decided *before* anything is quiesced, checkpointed, or run lives here, so a +fork that cannot execute soundly dies with nothing paid for. Three families: + +* **layers** — which checkpoint layers exist, that a requested set is + non-empty (:func:`resolve_layers`), and that the live planes can actually + take each requested layer (:func:`gate_layers`); +* **stages** — the one capture path for a stage boundary + (:func:`capture_stage`, called by the lifecycle at the boundaries named in + ``RolloutConfig.snapshot_stages``) and the resolution of a recorded stage + into a branch's roll-back point (:func:`recorded_stage_checkpoint`); +* **deltas** — the whole-vector gate (:func:`validate_deltas`) mirroring what + the delta will need at execution time: the boundary it can execute from, the + layer that boundary's snapshot must carry, the parent precondition a + ``skill_mode`` delta cannot see from the delta alone, and the + caller-supplied-runner conflicts. + +Execution lives elsewhere: :mod:`benchflow.branch_transaction` runs the fork +this module admitted, and :mod:`benchflow.branch_children` picks each child's +execution path. +""" + +from __future__ import annotations + +import dataclasses +import json +import logging +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from benchflow._utils.config_override import validate_overlay +from benchflow.branch import StageSnapshot +from benchflow.branch import checkpoint_composed as _checkpoint_composed +from benchflow.branch_children import ( + FRESH_CHILD_LAYER, + FRESH_CHILD_STAGE, + SKILL_DELTA_LAYER, + SKILL_DELTA_STAGE, + resolve_child_skill_policy, + resolve_environment_ref_delta, +) +from benchflow.branch_delta import BranchDelta, BranchDeltaNotSupported +from benchflow.branch_lineage import write_stage_snapshots +from benchflow.branch_stage import ( + BranchStageNotCaptured, + captured_stages, + validate_stage, +) +from benchflow.skill_policy import SKILL_MODE_NO_SKILL + +if TYPE_CHECKING: + from benchflow.branch_children import ChildRunner + from benchflow.rollout import Rollout + from benchflow.trajectories.tree import RolloutNode + +logger = logging.getLogger(__name__) + +# The checkpoint layers a branch can compose (RFC §3.1). Agent-session is +# layer three and explicitly out of scope for v1. +_SNAPSHOT_LAYERS = frozenset({"environment", "sandbox"}) + +# The BranchDelta fields the engine executes vs records-only. +# ``injected_prompt`` runs in place on the parent's rollout; ``skill_mode``, +# ``config_override`` and ``environment_ref`` run the child as a fresh rollout +# from the env-ready snapshot (RFC §3.3, gated by +# :func:`_validate_skill_delta` / :func:`_validate_config_delta` / +# :func:`_validate_environment_delta`). Any remaining field exists in the +# schema and provenance so the artifact format is stable, and fails closed +# here. Derived from the schema (all BranchDelta fields minus the executable +# set) so a future BranchDelta field is unsupported-by-default — it fails +# closed here instead of being silently ignored. Sorted for deterministic +# errors; empty today, and load-bearing the day a field is added. +_EXECUTABLE_DELTA_FIELDS = frozenset( + {"injected_prompt", "skill_mode", "config_override", "environment_ref"} +) +_UNSUPPORTED_DELTA_FIELDS = tuple( + sorted({f.name for f in dataclasses.fields(BranchDelta)} - _EXECUTABLE_DELTA_FIELDS) +) + + +class BranchParentSkillModeConflict(BranchDeltaNotSupported): + """A ``skill_mode`` delta forked from a parent carrying a baked-in pack. + + The precondition the other skill gates cannot see: the *parent's own* + skill mode. A ``with-skill`` parent's ``setup()`` injects the pack into + the Dockerfile it builds, so the pack is in the image — and the + ``env-ready`` snapshot is a commit of that image. A ``no-skill`` child + restoring it finds ``/skills`` already there, ``deploy_skills`` correctly + deploys nothing on top, and the arm is reported as ``no-skill`` while the + agent runs *with* the parent's skills. Nothing downstream can detect that: + the child's config, provenance and ablation row all say ``no-skill``. + + So the delta is refused unless the snapshot's world provably carries no + pack — i.e. the parent's own recorded skill mode is ``no-skill``, the one + mode whose skill policy resolves no host directory and therefore never + reaches ``inject_skills_into_dockerfile``. Each arm then deploys its own + skills (or none) at install time, over a pack-free image. + """ + + +class BranchChildExecutionNotSupported(NotImplementedError): + """The engine cannot run this fork's children safely (fail closed). + + Raised when the *boundary*, not the delta, makes execution unsound: an + ``env-ready`` fork whose snapshot omits the container layer. Every child of + that stage re-runs ``install_agent()`` for itself (see + :mod:`benchflow.branch_skill`), and without the container layer the restore + cannot undo the previous child's installation — the second child would + install on top of the first and report the result as a controlled + comparison. + """ + + +def resolve_layers(snapshot_layers: Iterable[str], *, subject: str) -> frozenset[str]: + """Validate a layer set — the one gate every snapshot path uses. + + ``subject`` names the operation in the diagnostic ("branch", "stage + snapshot 'env-ready'", "branch_at_stage('pre-verify'): the recorded stage + snapshot"). + + Called on *requested* layers by the cursor branch and by + :func:`capture_stage`, and on the layers a stage branch *derived* from its + recorded snapshot — the empty set is the same fail-open either way, so it + is rejected in one place rather than trusted because it came from a + checkpoint that was already taken. + """ + layers = frozenset(snapshot_layers) + unknown = layers - _SNAPSHOT_LAYERS + if unknown: + raise ValueError( + f"unknown snapshot_layers {sorted(unknown)!r} — allowed layers " + f"are {sorted(_SNAPSHOT_LAYERS)!r}" + ) + if not layers: + raise ValueError( + f"{subject} needs at least one layer — a fork whose checkpoint " + "captures nothing has no roll-back point, so nothing is restored " + "between children and every child runs in the world the previous " + "one left" + ) + return layers + + +def gate_layers( + rollout: Rollout, + layers: frozenset[str], + *, + subject: str, + requested: str, + gate_sandbox: bool = False, +) -> tuple[Any, Any]: + """Fail closed when a requested layer's plane or capability is missing. + + Returns the ``(environment, sandbox)`` live objects for the requested + layers — ``None`` for a layer that was not requested, which is exactly the + "layer not requested" signal :func:`checkpoint_composed` / + :func:`restore_composed` take. Every diagnostic names ``subject``, so a + stage capture says *which stage* could not be taken (RFC §3.2) instead of + the generic branch message. + + The container layer keeps its #384 semantics: the Branch lifecycle + composes container ⊃ environment-state ⊃ agent-session, and restoring only + environment state can produce inconsistent state for runs that mutate + process/service state the Environment manifest does not capture — so a + missing ``supports_snapshot`` never silently degrades. + """ + if "environment" in layers and getattr(rollout, "_environment", None) is None: + raise RuntimeError( + f"{subject} needs the Environment plane — there is no world to " + "snapshot. Pass RolloutConfig(environment_manifest=...)." + ) + sandbox = getattr(rollout, "_env", None) + if (gate_sandbox or "sandbox" in layers) and not getattr( + sandbox, "supports_snapshot", False + ): + sandbox_name = type(sandbox).__name__ if sandbox else "" + raise RuntimeError( + f"{subject} cannot run with {requested}: the active " + f"sandbox {sandbox_name!r} does not implement container-level " + "snapshot/restore. Use a provider whose Sandbox satisfies the " + "checkpoint contract (DockerSandbox or DaytonaSandbox in direct " + "mode), or drop the sandbox layer if Environment-state " + "checkpoint is sufficient for this run." + ) + return ( + rollout._environment if "environment" in layers else None, + sandbox if "sandbox" in layers else None, + ) + + +def _validate_skill_delta( + rollout: Rollout, + skill_mode: str, + *, + index: int, + at_stage: str | None, + layers: frozenset[str], + run_child: ChildRunner | None, +) -> None: + """Gate a ``skill_mode`` delta before anything is quiesced (RFC §3.3). + + Four preconditions, each a way the delta would otherwise measure nothing: + + * **the stage** — skills are deployed by ``install_agent()``, so only the + ``env-ready`` boundary (captured before it) can vary them; a cursor + branch forks a world that already resolved the question. + * **the container layer** — skills live in the container filesystem, so an + environment-state-only checkpoint cannot roll the parent's pack back and + a ``no-skill`` child would still find it mounted. + * **the parent's own skill mode** — the layer above rolls the container + back *to the parent's env-ready image*, and a ``with-skill`` parent + baked the pack into that image at build time. Rolling back to it hands + a ``no-skill`` child the pack it is supposed to be running without, and + every artifact still labels the arm ``no-skill``. Only a parent whose + recorded mode is ``no-skill`` provably carries no pack, so anything else + raises :class:`BranchParentSkillModeConflict`. + * **the runner** — a caller-supplied ``run_child`` owns the child's + execution, so it, not the engine, would decide what the child installs. + + The task's own skill policy is resolved here too, so ``with-skill`` + against a task shipping no bundled pack raises the same typed error + ``setup()`` would raise, before the branch restored anything. + """ + where = f"stage {at_stage!r}" if at_stage is not None else "the cursor" + if at_stage != SKILL_DELTA_STAGE: + raise BranchDeltaNotSupported( + f"deltas[{index}].skill_mode={skill_mode!r} executes only at " + f"the {SKILL_DELTA_STAGE!r} stage boundary, but this branch forks " + f"from {where}: skills are deployed by install_agent(), which has " + "already run by then. Capture the boundary with " + "RolloutConfig(snapshot_stages={'env-ready'}, " + "snapshot_layers={'environment', 'sandbox'}) and fork with " + "branch_at_stage('env-ready', ...) — the child then runs as a " + "fresh rollout over the restored sandbox (use_prebuilt_env)." + ) + if SKILL_DELTA_LAYER not in layers: + raise BranchDeltaNotSupported( + f"deltas[{index}].skill_mode={skill_mode!r} needs the " + f"{SKILL_DELTA_LAYER!r} layer in the {SKILL_DELTA_STAGE!r} " + f"snapshot, which captured layers={sorted(layers)!r}: skills are " + "deployed into the container filesystem, and an " + "environment-state-only checkpoint cannot roll them back — the " + "child would re-install against the parent's skills. Re-capture " + "the stage with snapshot_layers={'environment', 'sandbox'} (or " + "{'sandbox'} alone for a stateless environment)." + ) + parent_mode = rollout._config.recorded_skill_mode + if parent_mode != SKILL_MODE_NO_SKILL: + raise BranchParentSkillModeConflict( + f"deltas[{index}].skill_mode={skill_mode!r} cannot fork a parent " + f"whose own skill mode is {parent_mode!r}: only a " + f"{SKILL_MODE_NO_SKILL!r} parent provably bakes no skill pack " + f"into the image its setup() builds, and the " + f"{SKILL_DELTA_STAGE!r} snapshot every child restores is a commit " + "of that image. A 'no-skill' child of a parent that did bake one " + "restores the pack, deploys nothing on top of it, runs *with* the " + "parent's skills, and is still recorded as 'no-skill' in its " + "config, its provenance and its ablation row. Run the parent with " + f"skill_mode={SKILL_MODE_NO_SKILL!r} and let each arm's delta " + "deploy its own skills at install time." + ) + if run_child is not None: + raise ValueError( + f"deltas[{index}].skill_mode cannot be combined with an explicit " + "run_child — the caller's runner owns the child's execution, so " + "the engine cannot re-run install_agent() under the switched " + "mode. Drop run_child so the fresh-rollout runner executes the " + "delta, or switch skills inside your own runner." + ) + resolve_child_skill_policy(rollout._config, skill_mode) + + +def _validate_config_delta( + overlay: dict[str, Any], + *, + index: int, + at_stage: str | None, + run_child: ChildRunner | None, +) -> None: + """Gate a ``config_override`` delta before anything is quiesced (RFC §3.3). + + Three preconditions, each a way the delta would otherwise measure nothing + (or worse, measure the wrong thing): + + * **the stage** — the C-axis overlay is deep-merged into the task's + resolved config by ``setup()``, which only a fresh child rollout re-runs; + by any later boundary the parent's config has already been consumed + (its timeout enforced, its prompts resolved, its sandbox built), so a + child there would record the override and run without it. + * **the allowlist** — the same fail-closed allowlist the run-level overlay + passes through (#790): scorer-touching sections (``verifier`` / + ``reward`` / ``solution`` / ``oracle``) are rejected *here*, before the + parent is quiesced, not at child setup after a snapshot was restored. + * **the runner** — a caller-supplied ``run_child`` owns the child's + execution, so the engine cannot run the child under the merged config. + """ + where = f"stage {at_stage!r}" if at_stage is not None else "the cursor" + if at_stage != FRESH_CHILD_STAGE: + raise BranchDeltaNotSupported( + f"deltas[{index}].config_override executes only at the " + f"{FRESH_CHILD_STAGE!r} stage boundary, but this branch forks from " + f"{where}: the overlay is deep-merged into the task's resolved " + "config by the child's own setup(), which only a fresh child " + "rollout (use_prebuilt_env) re-runs — by any later boundary the " + "parent's config has already been consumed. Capture the boundary " + "with RolloutConfig(snapshot_stages={'env-ready'}, " + "snapshot_layers={'environment', 'sandbox'}) and fork with " + "branch_at_stage('env-ready', ...)." + ) + if run_child is not None: + raise ValueError( + f"deltas[{index}].config_override cannot be combined with an " + "explicit run_child — the caller's runner owns the child's " + "execution, so the engine cannot run setup() under the merged " + "config. Drop run_child so the fresh-rollout runner executes the " + "delta, or apply the override inside your own runner." + ) + try: + validate_overlay(overlay) + except ValueError as exc: + raise ValueError(f"deltas[{index}].config_override: {exc}") from None + + +def _validate_environment_delta( + rollout: Rollout, + environment_ref: str, + *, + index: int, + at_stage: str | None, + run_child: ChildRunner | None, +) -> None: + """Gate an ``environment_ref`` delta before anything is quiesced (RFC §3.3). + + The stage and runner gates mirror the other fresh-child deltas: the child + manifest binds at *provision* time, and provisioning happens only on the + fresh-child path a ``env-ready`` fork runs (an in-place child at a later + boundary would keep the parent's provisioned services no matter what its + delta records). The delta's own content gates — a manifest-bound parent, + a resolvable ref, the same image, framework-started services on both + sides — live in :func:`~benchflow.branch_skill.resolve_environment_ref_delta`, + which the child runner re-derives its manifest through, so validation and + execution cannot drift. + """ + where = f"stage {at_stage!r}" if at_stage is not None else "the cursor" + if at_stage != FRESH_CHILD_STAGE: + raise BranchDeltaNotSupported( + f"deltas[{index}].environment_ref={environment_ref!r} executes " + f"only at the {FRESH_CHILD_STAGE!r} stage boundary, but this " + f"branch forks from {where}: the manifest's environment plane is " + "provisioned over the restored sandbox by the fresh child rollout " + "(use_prebuilt_env), and at any later boundary the parent's " + "provisioned services survive the fork — the delta would be " + "recorded but not enforced. Capture the boundary with " + "RolloutConfig(snapshot_stages={'env-ready'}, " + "snapshot_layers={'environment', 'sandbox'}) and fork with " + "branch_at_stage('env-ready', ...)." + ) + if run_child is not None: + raise ValueError( + f"deltas[{index}].environment_ref cannot be combined with an " + "explicit run_child — the caller's runner owns the child's " + "execution, so the engine cannot provision the child manifest's " + "services. Drop run_child so the fresh-rollout runner executes " + "the delta, or provision inside your own runner." + ) + resolve_environment_ref_delta( + rollout._config.environment_manifest, + environment_ref, + subject=f"deltas[{index}].environment_ref", + ) + + +def validate_deltas( + rollout: Rollout, + deltas: Sequence[BranchDelta | None] | None, + *, + n: int, + at_stage: str | None, + layers: frozenset[str], + run_child: ChildRunner | None, +) -> None: + """Gate the whole delta vector before anything runs (RFC §3.3). + + One entry per child, validated in order so a bad or not-yet-executable + delta fails closed with nothing quiesced, checkpointed, or run — and with + a diagnostic naming its index. A field outside the executable set fails + unsupported-by-default; each executable field runs its own gate above; an + ``injected_prompt`` cannot combine with a caller-supplied ``run_child`` + because that runner owns the child's prompts. + """ + if deltas is None: + return + if len(deltas) != n: + raise ValueError( + f"deltas must carry exactly one entry per child: got " + f"{len(deltas)} deltas for n={n}" + ) + for index, delta in enumerate(deltas): + if delta is None: + continue + for field_name in _UNSUPPORTED_DELTA_FIELDS: + if getattr(delta, field_name) is not None: + raise BranchDeltaNotSupported( + f"deltas[{index}].{field_name} is set, but the branch " + "engine executes injected_prompt and skill_mode only " + f"— {field_name!r} runs in the rollout-branching RFC " + "follow-on (child-as-fresh-rollout via " + "use_prebuilt_env). The field is already recorded in " + "the schema and provenance; drop it to branch today." + ) + if delta.config_override is not None: + _validate_config_delta( + delta.config_override, + index=index, + at_stage=at_stage, + run_child=run_child, + ) + if delta.environment_ref is not None: + _validate_environment_delta( + rollout, + delta.environment_ref, + index=index, + at_stage=at_stage, + run_child=run_child, + ) + if delta.skill_mode is not None: + _validate_skill_delta( + rollout, + delta.skill_mode, + index=index, + at_stage=at_stage, + layers=layers, + run_child=run_child, + ) + if delta.injected_prompt is not None and run_child is not None: + raise ValueError( + f"deltas[{index}].injected_prompt cannot be combined " + "with an explicit run_child — the caller's runner owns " + "the child's prompts. Bind the prompt into the " + "run_child closure, or drop run_child so the default " + "runner delivers it." + ) + + +def runs_fresh_children(at_stage: str | None, run_child: ChildRunner | None) -> bool: + """Whether the engine will run this fork's children as fresh rollouts. + + True for every engine-run child of ``env-ready``, whatever its delta. That + boundary precedes ``install_agent()``, so the restored world has no agent + binary, no sandbox user, no seeded verifier workspace, no path lockdown and + no skill pack: an in-place child there connects to something the restore + deleted, or — when the agent happens to survive in the base image — scores + a world missing everything installation deploys and reports it as an + ordinary child under one recorded delta. A caller-supplied ``run_child`` + owns its own execution and is left alone. + """ + return run_child is None and at_stage == FRESH_CHILD_STAGE + + +def validate_fresh_children(layers: frozenset[str], *, at_stage: str) -> None: + """Gate a fresh-child fork before anything is quiesced. + + One precondition beyond the per-delta gates: the stage snapshot must carry + the container layer. Every child re-runs ``install_agent()``, which writes + to the container filesystem, so an environment-state-only checkpoint cannot + put the world back between children — child *k+1* would install on top of + child *k*'s agent, user, lockdown and skill pack. The skill-delta gate says + the same thing in skill terms and fires first when a skill delta is + present; this covers the children that carry no skill delta at all. + """ + if FRESH_CHILD_LAYER not in layers: + raise BranchChildExecutionNotSupported( + f"a branch at the {at_stage!r} stage boundary runs every child as " + f"a fresh rollout, which needs the {FRESH_CHILD_LAYER!r} layer in " + f"that snapshot; it captured layers={sorted(layers)!r}. " + f"{at_stage!r} precedes install_agent(), so each child installs " + "the agent (and its skills, user and lockdown) for itself, and an " + "environment-state-only checkpoint cannot roll that back between " + "children. Re-capture the stage with " + "snapshot_layers={'environment', 'sandbox'} (or {'sandbox'} alone " + "for a stateless environment), or supply your own run_child if " + "you are executing the children yourself." + ) + + +async def completed_exchange_count(rollout: Any) -> int | None: + """The LLM exchanges completed by ``rollout`` at this instant, best-effort. + + The stage half of the RFC §3.5 stage-tagged replay cut: every stage + capture records which ``llm_trajectory.jsonl`` prefix had completed when + the boundary closed, so ``bench eval continue --cut-stage`` can replay + exactly that prefix later. The count comes from the usage gateway's live + capture (``LiveUsageGateway.live_exchange_count`` — drained to the + gateway log's end, so it cannot silently undercount), reached duck-typed + because the branch engine, like the rollout kernel, must not import the + concrete provider plane; the name is agreed in + ``contracts.planes.LiveUsageGateway`` and statically asserted by + ``benchflow.providers.litellm_runtime``. Any absence — no runtime + connected yet (``env-ready``), a gateway without the accessor (Daytona's + stop()-time import), the drain failing or timing out — degrades to + ``None``, recorded honestly as an unknown index, never a guess. + """ + runtime = getattr(rollout, "_usage_runtime", None) + counter = getattr(getattr(runtime, "server", None), "live_exchange_count", None) + if counter is None: + return None + try: + count = await counter() + except Exception: + logger.debug("live exchange count unavailable at stage capture", exc_info=True) + return None + return count if isinstance(count, int) and not isinstance(count, bool) else None + + +async def capture_stage( + rollout: Rollout, + stage: str, + *, + snapshot_layers: Iterable[str] = frozenset({"environment"}), +) -> StageSnapshot: + """Take the composed stage snapshot at ``rollout``'s cursor (RFC §3.2). + + The stage-boundary policy's one capture path: validate the stage name and + the requested layers, fail closed on a missing plane or capability (the + diagnostic names the stage), compose the checkpoint on the cursor node — + so the stage tag flows into ``tree.json`` through the node's own recorded + state — and register it under ``stage`` on the rollout. The registry is + re-serialized to ``stage_snapshots.json`` after every capture, with write + failures logged and isolated from the rollout. + + Each capture also records the completed-LLM-exchange index of the moment + (:func:`completed_exchange_count`, taken immediately *before* the layers + snapshot so the index names the world the snapshot froze, not one that + kept moving while layers committed) into ``snapshot.meta`` and therefore + into ``stage_snapshots.json`` — the stage→exchange-index data a + stage-named replay cut resolves. ``None`` (count unavailable) is recorded + as an honest null. + + Callers are the lifecycle's own boundaries (``start()``/``verify()``, gated + on ``RolloutConfig.snapshot_stages``) and ``Rollout.mark_stage`` for the + boundaries only the caller can see. Two stages captured without an + intervening Step land on the same node (``pre-verify`` and ``post-verify`` + of a rollout that verifies once, say): the registry keeps both, and the + node's own tag is the most recent — a node has one checkpoint, and the + branch that adopts it re-tags it with the stage it branched from. + """ + validate_stage(stage) + subject = f"stage snapshot {stage!r}" + layers = resolve_layers(snapshot_layers, subject=subject) + snap_env, snap_sandbox = gate_layers( + rollout, + layers, + subject=subject, + requested=f"snapshot_layers={sorted(layers)!r}", + ) + exchanges_completed = await completed_exchange_count(rollout) + try: + snapshot = await _checkpoint_composed( + rollout._cursor, + environment=snap_env, + sandbox=snap_sandbox, + stage=stage, + ) + except Exception as exc: + exc.add_note( + f"stage snapshot {stage!r} (snapshot_layers={sorted(layers)!r}) " + "failed during checkpoint — the rollout fails closed; no partial " + "checkpoint was recorded on the node." + ) + raise + snapshot.meta["exchanges_completed"] = exchanges_completed + rollout._stage_snapshots[stage] = snapshot + rollout._stage_nodes[stage] = rollout._cursor + + run_dir = getattr(rollout, "_rollout_dir", None) + if run_dir is not None: + try: + write_stage_snapshots(run_dir=run_dir, snapshots=rollout._stage_snapshots) + except Exception: + logger.warning( + "stage snapshot artifact write under %s failed — the recorded " + "stage registry is unaffected", + run_dir, + exc_info=True, + ) + return snapshot + + +def recorded_stage_checkpoint( + rollout: Rollout, stage: str, snapshot_layers: Iterable[str] | None +) -> tuple[StageSnapshot, RolloutNode, frozenset[str]]: + """Resolve a recorded stage into ``(snapshot, branch node, layers)``. + + A stage that was never captured fails closed with + :class:`~benchflow.branch_stage.BranchStageNotCaptured` naming what *was* + captured — degrading to a checkpoint at the cursor would fork a different + world state than the caller asked for and mislabel it in provenance. The + layers are derived from the recorded snapshot; an explicit + ``snapshot_layers`` that disagrees is a caller bug, not a request to + re-snapshot. + + Derived is not the same as trusted: a snapshot carrying *neither* ref + derives the empty set, which the layer gates below then have nothing to + check — the fork would restore nothing between children and still record a + clean stage fork. So the derived set goes through :func:`resolve_layers`, + the same non-empty gate the cursor branch runs on its requested set, and + it runs before the disagreement check: an empty capture is broken however + it is described, and ``snapshot_layers=set()`` would otherwise "agree" + with it. + """ + validate_stage(stage) + registry: dict[str, StageSnapshot] = getattr(rollout, "_stage_snapshots", {}) + snapshot = registry.get(stage) + if snapshot is None: + raise BranchStageNotCaptured( + f"no snapshot recorded at stage {stage!r} — this rollout captured " + f"{captured_stages(registry)!r}. Request the stage up front with " + "RolloutConfig(snapshot_stages={...}), or record it at the cut " + "point with Rollout.mark_stage()." + ) + layers = resolve_layers( + ( + layer + for layer, ref in ( + ("environment", snapshot.environment_ref), + ("sandbox", snapshot.sandbox_ref), + ) + if ref is not None + ), + subject=f"branch_at_stage({stage!r}): the recorded stage snapshot", + ) + if snapshot_layers is not None and frozenset(snapshot_layers) != layers: + raise ValueError( + f"branch_at_stage({stage!r}, snapshot_layers=" + f"{sorted(frozenset(snapshot_layers))!r}) disagrees with the " + f"layers stage {stage!r} actually captured ({sorted(layers)!r}) — " + "a stage branch restores exactly what the stage snapshotted; " + "re-run with the layers the capture used, or capture the stage " + "with the layers you want." + ) + node = getattr(rollout, "_stage_nodes", {}).get(stage, rollout._cursor) + return snapshot, node, layers + + +# Snapshot lifetime (RFC §3.6) — the retention policy for what a capture +# committed. A `bf-snap-*` image dies with the rollout's `compose down --rmi +# all`, so every recorded ref must either be exported before cleanup or be +# marked ephemeral in the artifact: a ref `docker image inspect` can no longer +# resolve must never read as restorable. `bench eval ablate --keep-snapshots` +# and `bench eval run --keep-snapshots` share this machinery, so +# `ablation.json` and `stage_snapshots.json` teach readers one schema. + + +class SnapshotExportUnsupported(RuntimeError): + """The active sandbox backend cannot ``docker save`` a snapshot image.""" + + +def _snapshot_tar_name(ref: str) -> str: + """A filesystem-safe tar basename for a snapshot image ref.""" + import re + + return re.sub(r"[^A-Za-z0-9._-]+", "-", ref).strip("-") or "snapshot" + + +def _file_sha256(path: Path) -> str: + """Streaming ``sha256:``-prefixed digest of a file (tars can be large).""" + import hashlib + + digest = hashlib.sha256() + with Path(path).open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return f"sha256:{digest.hexdigest()}" + + +def image_id_from_export_tar(tar_path: Path) -> str | None: + """The image id a ``docker save`` tar will load as, or ``None``. + + ``docker save`` records the image's config digest in the tar's + ``manifest.json`` (``Config`` is ``blobs/sha256/`` in current + format, ``.json`` in the legacy one); that digest *is* the id + ``docker image inspect`` reports after ``docker load``, so recording it + at export time lets the import path verify it restored exactly the image + the run recorded. Best-effort by design: an unreadable/foreign tar + returns ``None`` — the honest "unverifiable", never a guess — and the + import path then verifies only that the recorded ref resolves. + """ + import re + import tarfile + + try: + with tarfile.open(tar_path) as tar: + handle = tar.extractfile("manifest.json") + if handle is None: + return None + manifest = json.loads(handle.read()) + config = manifest[0]["Config"] + match = re.search(r"([0-9a-f]{64})", str(config)) + return None if match is None else f"sha256:{match.group(1)}" + except Exception: + logger.debug("no image id readable from %s", tar_path, exc_info=True) + return None + + +async def export_stage_snapshot( + sandbox: Any, *, sandbox_ref: str, out_dir: Path +) -> dict[str, Any]: + """``docker save`` a stage snapshot's sandbox image into ``out_dir``. + + The durable half of ``--keep-snapshots`` (RFC §3.6): the tar lands at + ``/snapshots/.tar`` and the returned record carries its + path, content sha256 and image id — enough for a reader to verify, + ``docker load`` and identity-check the world later + (:func:`benchflow.snapshot_import.import_stage_snapshots`). Raises + :class:`SnapshotExportUnsupported` when the sandbox backend cannot export + (the caller records the failure; the artifact stays truthful). + """ + export = getattr(sandbox, "export_image", None) + if export is None: + backend = "no sandbox attached" if sandbox is None else type(sandbox).__name__ + raise SnapshotExportUnsupported( + f"--keep-snapshots: the sandbox backend ({backend}) does not " + "support exporting snapshot images (export_image); the docker " + "backend does" + ) + snapshots_dir = Path(out_dir) / "snapshots" + snapshots_dir.mkdir(parents=True, exist_ok=True) + tar_path = snapshots_dir / f"{_snapshot_tar_name(sandbox_ref)}.tar" + await export(sandbox_ref, tar_path) + return { + "path": str(tar_path), + "sha256": _file_sha256(tar_path), + "image_id": image_id_from_export_tar(tar_path), + } + + +async def annotate_stage_snapshot_lifetime( + entry: dict[str, Any], *, sandbox: Any, keep: bool, out_dir: Path +) -> None: + """Stamp one serialized stage entry with its snapshot's lifetime. + + Mutates ``entry`` (a ``stage_snapshots_payload`` / + ``AblationReport.stage_snapshot`` entry) in place. With ``keep`` the + sandbox-layer image is exported to ``/snapshots/.tar`` and + the entry records ``ephemeral: false`` plus the tar's path, sha256 and + image id under ``exported``; without it — or when the export fails, + recorded as ``export_error`` — the entry records ``ephemeral: true, + exported: null`` so a reader knows the ref no longer resolves. Never + raises: the run's rewards must survive a failed export. + """ + entry["ephemeral"] = True + entry["exported"] = None + sandbox_ref = entry.get("sandbox_ref") + if not keep: + return + if sandbox_ref is None: + entry["export_error"] = ( + "--keep-snapshots: this stage recorded no sandbox-layer image to export" + ) + return + try: + exported = await export_stage_snapshot( + sandbox, sandbox_ref=sandbox_ref, out_dir=Path(out_dir) + ) + except Exception as exc: + from benchflow._utils.text import describe_exception + + entry["export_error"] = describe_exception(exc) + logger.error( + "--keep-snapshots could not export %s — the artifact records the " + "snapshot as ephemeral: %s", + sandbox_ref, + entry["export_error"], + ) + else: + entry["ephemeral"] = False + entry["exported"] = exported + + +async def finalize_stage_snapshots(rollout: Rollout) -> None: + """Make ``stage_snapshots.json`` truthful before cleanup destroys images. + + Called by ``Rollout.cleanup()`` immediately before the sandbox is stopped + with ``delete=True`` — the ``compose down --rmi all`` that destroys every + committed ``bf-snap-*`` image. Each recorded stage entry is annotated + with its lifetime (:func:`annotate_stage_snapshot_lifetime`): + ``RolloutConfig.keep_snapshots`` exports the images into + ``/snapshots/`` first and records the tars; otherwise the refs + are marked ``ephemeral: true, exported: null``, so the artifact never + shows a bare ref that no longer resolves (PR #1046 second review). An + entry an earlier writer already annotated — ``bench eval ablate + --keep-snapshots`` exports the branched stage into its own out-dir before + this runs — is preserved, not clobbered. Never raises: retention must not + block teardown. + """ + run_dir = getattr(rollout, "_rollout_dir", None) + try: + registry: dict[str, StageSnapshot] = getattr(rollout, "_stage_snapshots", {}) + if not registry or run_dir is None: + return + from benchflow.branch_lineage import ( + stage_snapshots_payload, + write_stage_snapshots_stages, + ) + + keep = bool(getattr(getattr(rollout, "_config", None), "keep_snapshots", False)) + path = Path(run_dir) / "stage_snapshots.json" + existing: dict[str, Any] = {} + if path.exists(): + try: + recorded = json.loads(path.read_text()).get("stages") + if isinstance(recorded, dict): + existing = recorded + except (json.JSONDecodeError, OSError): + logger.warning( + "unreadable stage_snapshots.json under %s — rewriting it " + "from the recorded registry", + run_dir, + exc_info=True, + ) + stages: dict[str, dict[str, Any]] = {} + for stage, entry in stage_snapshots_payload(registry).items(): + recorded_entry = existing.get(stage) + if isinstance(recorded_entry, dict) and "ephemeral" in recorded_entry: + stages[stage] = recorded_entry + continue + await annotate_stage_snapshot_lifetime( + entry, + sandbox=getattr(rollout, "_env", None), + keep=keep, + out_dir=Path(run_dir), + ) + stages[stage] = entry + write_stage_snapshots_stages(run_dir=Path(run_dir), stages=stages) + except Exception: + logger.warning( + "stage-snapshot lifetime finalization under %s failed — cleanup continues", + run_dir, + exc_info=True, + ) diff --git a/src/benchflow/branch_report.py b/src/benchflow/branch_report.py new file mode 100644 index 000000000..01830b52a --- /dev/null +++ b/src/benchflow/branch_report.py @@ -0,0 +1,562 @@ +"""Branch reporting — the ablation report model and its attribution. + +Everything that turns a finished (or half-finished) fork into something a +reader can hold: the report half of ``bench eval ablate`` that used to live +inside :mod:`benchflow.ablate` — the :class:`ArmOutcome` / +:class:`AblationReport` model, the pairing of arms with the child nodes the +engine forked for them (:func:`outcomes_for_arms`), the per-test mining from +each child's own verifier artifacts, and the attribution that turns rewards +into one-line observational verdicts (:func:`attribute`, +:func:`sub_test_attribution`). The fork's own lineage artifacts are adjacent +but deliberately elsewhere: the serializers live in +:mod:`benchflow.branch_lineage`, and the engine's failure-isolated lineage +write stays in :mod:`benchflow.rollout_branch`, whose module namespace is the +patch seam the lineage-isolation tests pin. + +The split follows the data's direction: :mod:`benchflow.ablate` owns the +*request* side (driving the parent and the fork, with arm specs and +pre-flight validation in :mod:`benchflow.ablate_arms`); this module owns the +*result* side, and nothing here reaches back — ``ablate`` imports from +``branch_report``, never the reverse, and the engine-facing lineage writer +knows nothing about arms. + +Determinism is the report's contract: arms keep request order, every delta is +the engine's own content-addressed provenance dict, test names sort, and no +wall-clock *timestamp* is recorded anywhere (per-arm ``wall_clock_sec`` is a +measured duration — part of an arm's cost, not a stamp). +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from benchflow.branch import UNSCORED_KEY +from benchflow.branch_transaction import CHILD_WALL_CLOCK_KEY +from benchflow.skill_policy import SKILL_MODE_NO_SKILL, SKILL_MODE_WITH_SKILL + +if TYPE_CHECKING: + from collections.abc import Sequence + + from benchflow.ablate_arms import AblationArm + from benchflow.environment.manifest import ManifestBinding + from benchflow.trajectories.tree import RolloutNode, RolloutTree + +logger = logging.getLogger(__name__) + +SCHEMA_VERSION = 1 + +#: The report an ablation writes into its output directory. +REPORT_FILENAME = "ablation.json" + +#: Reward at or above which a rollout counts as a pass — the framework's +#: binary convention (``bench eval metrics`` and ``bench review`` both read +#: 1.0 as passing). +PASS_REWARD = 1.0 + +STATUS_PASS = "pass" +STATUS_FAIL = "fail" +STATUS_ERROR = "error" +STATUS_SKIPPED = "skipped" + +#: Reference name used when an arm has no counterpart arm to compare against. +REFERENCE_PARENT = "parent" + + +# The ablation report model + + +def environment_stamp(binding: ManifestBinding | None) -> dict[str, Any] | None: + """The report's record of a bound environment, deterministic by design. + + Name, the ref exactly as the caller wrote it (flag value, arm spec, or + ``task.md`` declaration — never a machine-local resolved path), the + manifest's ``sha256`` content address, and the image(s) it names. This is + the report-side answer to "which environment did this ablation compare + arms in", readable without the registry. + """ + if binding is None: + return None + manifest = binding.manifest + return { + "name": manifest.name, + "ref": binding.ref, + "env_hash": binding.env_hash, + "image": manifest.image, + "base_image": manifest.base_image, + } + + +@dataclass +class ArmOutcome: + """What one arm did: its reward, its cost, and its recorded delta. + + ``tests`` is the arm's per-test outcome map (``{test name -> status}``) + mined from its own verifier artifacts, or ``None`` when that verifier + reported none. ``None`` and ``{}`` are not the same thing here and the + distinction is load-bearing: a missing map means *not observed*, and no + sub-test claim may be made about that arm. + + ``environment`` is the :func:`environment_stamp` of the manifest this + arm's delta swapped in — set only for an ``env:`` arm, absent for every + arm that inherits the parent's bound environment (which the report stamps + once at the top level). + """ + + name: str + kind: str + delta: dict[str, Any] = field(default_factory=dict) + source: str | None = None + reward: float | None = None + wall_clock_sec: float | None = None + delta_execution: str | None = None + node_id: str | None = None + artifacts: str | None = None + tests: dict[str, str] | None = None + error: str | None = None + reference: str | None = None + verdict: str = "" + environment: dict[str, Any] | None = None + + @property + def status(self) -> str: + """``pass`` / ``fail`` / ``error`` / ``skipped`` — derived in one place. + + ``skipped`` is an arm that never ran because an earlier arm errored: + the branch engine runs children sequentially and propagates a child + failure, so the arms after it have no world to report on. They are + reported as skipped, never as a zero reward. + """ + if self.reward is not None: + return STATUS_PASS if self.reward >= PASS_REWARD else STATUS_FAIL + if self.error is not None: + return STATUS_ERROR + return STATUS_SKIPPED + + def to_dict(self) -> dict[str, Any]: + return { + "name": self.name, + "kind": self.kind, + "delta": self.delta, + "delta_execution": self.delta_execution, + "source": self.source, + "reward": self.reward, + "status": self.status, + "wall_clock_sec": self.wall_clock_sec, + "node": self.node_id, + "artifacts": self.artifacts, + "tests": None if self.tests is None else dict(sorted(self.tests.items())), + "error": self.error, + "reference": self.reference, + "verdict": self.verdict, + "environment": self.environment, + } + + +@dataclass +class AblationReport: + """The ablation's result: the parent's own run plus one row per arm. + + ``environment`` is the :func:`environment_stamp` of the world the parent + (and therefore every arm that carries no environment delta) ran in — + ``None`` when no manifest was bound. ``stage_snapshot`` is the branched + stage's recorded snapshot refs (the committed sandbox image ref and the + environment snapshot id, plus the layers captured), the handles a reader + needs to restore that world and re-branch it by hand later — annotated by + :func:`~benchflow.ablate.retain_stage_snapshot` with their lifetime: + ``ephemeral: true`` with ``exported: null`` when cleanup destroyed the + image (the default), or ``ephemeral: false`` with the exported tar's path + and sha256 under ``--keep-snapshots``. + """ + + task_id: str + task_path: str + stage: str + snapshot_layers: list[str] + agent: str + model: str | None + sandbox: str + arms: list[ArmOutcome] + parent_reward: float | None = None + parent_error: str | None = None + parent_run_dir: str | None = None + value: float | None = None + error: str | None = None + environment: dict[str, Any] | None = None + stage_snapshot: dict[str, Any] | None = None + + @property + def has_errors(self) -> bool: + """Whether any arm failed to produce a reward. + + The *parent*'s own error is deliberately not counted: attributing a + failed run is the reason this command exists (RFC §1), so a parent that + failed while every arm scored is a complete ablation, not a failed one. + """ + return self.error is not None or any( + arm.status in (STATUS_ERROR, STATUS_SKIPPED) for arm in self.arms + ) + + def to_dict(self) -> dict[str, Any]: + """The ``ablation.json`` payload. + + Deterministic given the same rewards: arms keep request order, every + delta is the engine's own content-addressed provenance dict, test names + sort, and no wall-clock *timestamp* is recorded anywhere. Per-arm + ``wall_clock_sec`` is a measured duration — the one field that varies + between identical runs, carried because an ablation's cost per arm is + part of its result. + + ``test_attribution`` is derived, not stored: it is a reading of the + arms' own ``tests`` maps, so the section can never disagree with the + rows it summarizes. + """ + return { + "schema_version": SCHEMA_VERSION, + "task": {"id": self.task_id, "path": self.task_path}, + "stage": self.stage, + "snapshot_layers": sorted(self.snapshot_layers), + "stage_snapshot": self.stage_snapshot, + "agent": self.agent, + "model": self.model, + "sandbox": self.sandbox, + "environment": self.environment, + "parent": { + "reward": self.parent_reward, + "error": self.parent_error, + "run_dir": self.parent_run_dir, + }, + "value": self.value, + "error": self.error, + "arms": [arm.to_dict() for arm in self.arms], + "test_attribution": sub_test_attribution(self.arms), + } + + +def write_ablation_report(report: AblationReport, out_dir: Path) -> Path: + """Write ``/ablation.json`` and return its path. + + Deterministic like every other branch artifact: sorted keys, indented, + trailing newline. + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + path = out_dir / REPORT_FILENAME + path.write_text( + json.dumps(report.to_dict(), sort_keys=True, indent=2) + "\n", encoding="utf-8" + ) + return path + + +# Reading the fork back: arms -> child nodes -> outcomes + + +def branch_children_of(tree: RolloutTree) -> list[RolloutNode]: + """The branch children of this run, in fork order. + + A branch child is the only node the engine records a ``delta`` on (at fork + time, RFC §3.4), and ``RolloutTree.nodes()`` yields pre-order from the root + — so this is fork order without reading engine private state or guessing at + positions. + """ + return [node for node in tree.nodes() if "delta" in node.state] + + +def _child_test_outcomes(child_dir: Path) -> dict[str, str] | None: + """One branch child's per-test outcomes, whichever way the engine ran it. + + The child kinds keep their verifier output in different places — a + fresh-rollout child in its own run directory, an in-place child in the + ``mounted/`` archive of what it wrote to the parent's shared bind mounts — + and :func:`~benchflow.branch_artifacts.child_artifact_roots` is the one + place that knows both. Each candidate is read with the CLI's existing CTRF + reader (one parser, no branch-specific copy), and the first root that + reports outcomes wins; ``None`` when neither does, which is a real + observation ("this arm's verifier emitted no per-test data") and keeps the + arm out of every sub-test comparison. + + Without the ``mounted/`` fallback an in-place ablation (``--at-stage + pre-verify`` / ``post-verify``) fell back to scalar-only attribution while + the per-test data sat on disk one directory down — the exact difference + the sub-test section exists to expose. + """ + from benchflow.branch_artifacts import child_artifact_roots + from benchflow.cli._failure_evidence import artifact_test_outcomes + + for root in child_artifact_roots(child_dir): + outcomes = artifact_test_outcomes(root) + if outcomes is not None: + return outcomes + return None + + +def outcomes_for_arms( + arms: Sequence[AblationArm], + children: Sequence[RolloutNode], + *, + run_dir: Path | None, + branch_error: str | None, + environment_stamps: dict[str, dict[str, Any]] | None = None, +) -> list[ArmOutcome]: + """Pair each arm with the child node the engine forked for it. + + Children are created one at a time and each runs before the next is + attached, so a failing arm leaves a node with no recorded reward and the + arms after it leave no node at all: the arm that raised carries the branch + error, the rest report as skipped. + + An arm whose child ran but was never scored (the engine recorded + :data:`~benchflow.branch.UNSCORED_KEY` on the node) is an arm *error* + carrying the engine's reason — never a reward. A missing score is not an + observation, and an ablation that reports one as ``0.00`` invents its own + evidence. + + Each child's per-test outcomes are mined from its own verifier artifacts + here, where the child directory is known — see :func:`_child_test_outcomes` + for the two places a child's verifier output can live. Report-time reads + only: the engine stays file-free, and a child whose verifier emitted no + CTRF report keeps ``tests = None``. + """ + from benchflow.branch_lineage import branch_child_dir + + outcomes: list[ArmOutcome] = [] + for index, arm in enumerate(arms): + node = children[index] if index < len(children) else None + outcome = ArmOutcome( + name=arm.name, + kind=arm.kind, + delta=arm.delta.provenance_dict(), + source=arm.source, + # The swapped-in environment an env: arm declares — request-level + # provenance, stamped whether or not the arm got to run. + environment=(environment_stamps or {}).get(arm.name), + ) + if node is not None: + outcome.node_id = node.id + outcome.delta_execution = node.state.get("delta_execution") + outcome.wall_clock_sec = node.state.get(CHILD_WALL_CLOCK_KEY) + if run_dir is not None and node.parent is not None: + child_dir = branch_child_dir(run_dir, node.parent.id, node.id) + outcome.artifacts = str(child_dir) + outcome.tests = _child_test_outcomes(child_dir) + reward = node.state.get("reward") + if reward is None: + outcome.error = ( + node.state.get(UNSCORED_KEY) + or branch_error + or "the branch ended without a reward" + ) + else: + outcome.reward = float(reward) + outcomes.append(outcome) + return outcomes + + +# Attribution + + +#: How many differing test names a one-line verdict spells out before it rolls +#: the rest up as ``(+N more)`` — a verdict is a line, not a list. +_VERDICT_TEST_NAMES = 3 + + +def _score(reward: float) -> str: + return f"{reward:.2f}" + + +def _tested(arms: Sequence[ArmOutcome]) -> list[ArmOutcome]: + """The arms whose verifier actually reported per-test outcomes. + + An arm with ``tests is None`` was not observed at this granularity, so it + is excluded from every sub-test comparison rather than being treated as an + arm where nothing ran — otherwise one CTRF-less arm would make every test + of its counterpart look like a difference. + """ + return [arm for arm in arms if arm.tests is not None] + + +def _name_list(names: Sequence[str]) -> str: + """Up to :data:`_VERDICT_TEST_NAMES` names, the rest rolled up as a count.""" + shown = ", ".join(names[:_VERDICT_TEST_NAMES]) + extra = len(names) - _VERDICT_TEST_NAMES + return f"{shown} (+{extra} more)" if extra > 0 else shown + + +def differing_tests(arms: Sequence[ArmOutcome]) -> list[dict[str, Any]]: + """The tests whose outcome is not identical across the arms that reported + one, sorted by test name. + + This is the sub-outcome an ablation's scalar reward can hide: two arms can + both score 0.00 while one passes ``test_a`` and fails ``test_b`` and the + other does the reverse. Each entry is ``{"test": name, "outcomes": {arm -> + status}}``, with ``None`` for an arm whose report does not name that test + at all — a real difference (the test was reported for one arm and not the + other), stated as a missing observation rather than invented as a failure. + + Comparison needs at least two observed arms; with fewer there is nothing to + differ *from* and the result is empty. + """ + tested = _tested(arms) + if len(tested) < 2: + return [] + differences: list[dict[str, Any]] = [] + for name in sorted({test for arm in tested for test in arm.tests or {}}): + outcomes = {arm.name: (arm.tests or {}).get(name) for arm in tested} + if len(set(outcomes.values())) > 1: + differences.append( + {"test": name, "outcomes": dict(sorted(outcomes.items()))} + ) + return differences + + +def _scalar_tie(arms: Sequence[ArmOutcome]) -> bool: + """Whether every arm that produced a reward produced the *same* reward. + + The condition under which the per-arm verdicts read "no difference" — and + therefore the condition the per-test section exists to qualify. + """ + rewards = [arm.reward for arm in arms if arm.reward is not None] + return len(rewards) >= 2 and len(set(rewards)) == 1 + + +def sub_test_attribution(arms: Sequence[ArmOutcome]) -> dict[str, Any]: + """The report's ``test_attribution`` section: which sub-outcomes differ. + + Always present, and always honest about its own coverage: it names the + arms that reported per-test outcomes *and* the arms that did not, so a + reader can tell "the tests tie" from "no test data was mined". Only the + differing tests are listed with their per-arm outcomes; tying tests are + noise for attribution and appear as names only. Deterministic — sorted + names throughout, no wall-clock anywhere. + """ + tested = _tested(arms) + differing = differing_tests(arms) + differing_names = [entry["test"] for entry in differing] + differing_set = set(differing_names) + tying = [ + name + for name in sorted({test for arm in tested for test in arm.tests or {}}) + if name not in differing_set + ] + scalar_tie = _scalar_tie(arms) + if len(tested) < 2: + summary = ( + f"scalar-only attribution — {len(tested)} of {len(arms)} arms " + "reported per-test outcomes, so no sub-test comparison was made" + ) + elif not differing: + summary = ( + f"{len(tested)} arms reported per-test outcomes and all " + f"{len(tying)} tie — no sub-test difference in this comparison" + ) + elif scalar_tie: + summary = ( + f"scalar rewards tie, but {len(differing_names)} sub-test " + f"outcome(s) differ: {_name_list(differing_names)}" + ) + else: + summary = ( + f"{len(differing_names)} sub-test outcome(s) differ: " + f"{_name_list(differing_names)}" + ) + return { + "arms_with_tests": [arm.name for arm in tested], + "arms_without_tests": [arm.name for arm in arms if arm.tests is None], + "differing_tests": differing, + "tying_tests": tying, + "scalar_tie": scalar_tie, + "summary": summary, + } + + +def _reference_for( + arm: ArmOutcome, by_name: dict[str, ArmOutcome], parent_reward: float | None +) -> tuple[str | None, float | None]: + """The single arm (or the parent) this arm's verdict compares against. + + The two skills arms are each other's counterpart — that pair *is* the + ablation. Any other arm compares against the parent's own linear reward, + the only other observation the run produced; when neither exists the + verdict states the reward and claims nothing. + """ + counterpart = { + SKILL_MODE_WITH_SKILL: SKILL_MODE_NO_SKILL, + SKILL_MODE_NO_SKILL: SKILL_MODE_WITH_SKILL, + }.get(arm.name) + if counterpart is not None: + other = by_name.get(counterpart) + if other is not None and other.reward is not None: + return counterpart, other.reward + if parent_reward is not None: + return REFERENCE_PARENT, parent_reward + return None, None + + +def attribute( + outcomes: Sequence[ArmOutcome], *, parent_reward: float | None, stage: str +) -> None: + """Fill in each arm's reference and one-line verdict, in place. + + The verdict is an observation, not a causal claim: it names the two rewards + it compares, the stage they were forked from, and the fact that it rests on + one run per arm. Nothing is inferred about boundaries that were not forked + — localizing a failure to a stage takes a second ablation at a second + boundary (RFC §5, T3). + + One qualification the scalar cannot make on its own: when two arms tie on + reward but their verifiers disagree per test, the verdict says so and names + the tests instead of reading "no difference in this comparison". A binary + reward can net two opposite-signed sub-outcomes to exactly zero, and a tool + that printed "no difference" there would be true about the reward and false + about the behavior. The unqualified wording survives only where the + sub-tests were observed and tie, or were never observed at all. + """ + by_name = {outcome.name: outcome for outcome in outcomes} + for arm in outcomes: + reward = arm.reward + if arm.status == STATUS_ERROR: + arm.verdict = "errored before scoring — no reward to attribute" + continue + if reward is None: + arm.verdict = "not run — an earlier arm errored" + continue + ref_name, ref_reward = _reference_for(arm, by_name, parent_reward) + arm.reference = ref_name + if ref_name is None or ref_reward is None: + arm.verdict = ( + f"scores {_score(reward)} at {stage} — no counterpart arm or " + "parent reward to compare against" + ) + elif reward == ref_reward: + other = by_name.get(ref_name) + differing = ( + [entry["test"] for entry in differing_tests([arm, other])] + if other is not None + else [] + ) + tail = ( + f"scalar rewards tie, but {len(differing)} sub-test outcome(s) " + f"differ: {_name_list(differing)}" + if differing + else "no difference in this comparison" + ) + arm.verdict = ( + f"matches {ref_name} at {stage} (both {_score(reward)}) — {tail}" + ) + elif (reward >= PASS_REWARD) != (ref_reward >= PASS_REWARD): + passes = reward >= PASS_REWARD + arm.verdict = ( + f"{'passes' if passes else 'fails'} ({_score(reward)}) where " + f"{ref_name} {'fails' if passes else 'passes'} " + f"({_score(ref_reward)}) at {stage} — this delta decides the " + f"outcome when applied at {stage} (1 run per arm)" + ) + else: + direction = "higher" if reward > ref_reward else "lower" + arm.verdict = ( + f"scores {_score(reward)} vs {ref_name} {_score(ref_reward)} at " + f"{stage} — {direction} reward in this single comparison" + ) diff --git a/src/benchflow/branch_result.py b/src/benchflow/branch_result.py new file mode 100644 index 000000000..0bfec9108 --- /dev/null +++ b/src/benchflow/branch_result.py @@ -0,0 +1,247 @@ +"""A first-class ``result.json`` for in-place branch children (RFC §3.4). + +A branch child forked from ``env-ready`` runs as its own Rollout, so it leaves +the standard run artifacts — ``config.json`` / ``result.json`` / +``timing.json`` / trajectory — in its child directory for free. An *in-place* +child (a ``pre-verify`` / ``post-verify`` stage branch, or a cursor branch) +continues the **parent** Rollout instance instead: until this module, its +directory held only ``provenance.json`` / ``reward.json`` plus the ``mounted/`` +archive, and "what happened in this arm" could only be answered by +cross-reading ``tree.json``. This module closes that gap so every child +directory is self-describing, whichever way the engine ran the child. + +The mechanism leans on the isolation the engine already provides. The parent's +linear and result-bearing state is captured before the fork and restored +around every child (:class:`~benchflow.rollout_branch._LinearState`), so while +a child runs, the shared instance *is* that child's state. Two functions split +the work: + +* :func:`scope_child_result_state` — called by the engine after it restores + the parent's state onto the instance and before the child runs. It zeroes + the result-bearing fields the child's own phases write (``_timing``, + ``_rewards``, ``_verifier_error``, the diagnostics collector, …), so that + what those fields hold *after* the child is exactly the child's own — never + the parent's value showing through a field the child happened not to write. + The parent's values come back with the engine's next restore, exactly as + they always have. +* :func:`write_in_place_child_result` — called after the child completed + (scored or unscored; a child that raised ends the fork and is evidenced by + its ``mounted/`` archive instead). It builds the child's result through the + same :func:`~benchflow.rollout._results._build_rollout_result` the linear + run and the fresh-rollout child use — one result writer, no branch-specific + copy — from the child's own slice of the shared state: the trajectory steps + and prompts *appended since the fork baseline*, the scoped timing dict, the + child's own rewards and verifier error. + +The honesty rule is the module's contract: a field an in-place child does not +genuinely produce is ``null``/absent, **never copied from the parent**. That +is why token usage reports ``usage_source="unavailable"`` (per-child usage is +not attributable on a shared instance), why ``error`` is ``None`` unless the +child itself recorded one, and why the child's ``rewards`` fall back to the +reward the engine recorded on the child node — an observation the child did +produce — rather than to the parent's rewards dict. + +Failure-isolated like every branch artifact writer: a result that cannot be +built is logged and skipped, and never costs the reward the fork was run to +measure. +""" + +from __future__ import annotations + +import copy +import logging +from typing import TYPE_CHECKING + +from benchflow.branch_lineage import branch_child_dir, child_provenance + +if TYPE_CHECKING: + from datetime import datetime + from pathlib import Path + + from benchflow.rollout import Rollout + from benchflow.trajectories.tree import RolloutNode + +logger = logging.getLogger(__name__) + +#: The result-bearing fields :func:`scope_child_result_state` zeroes to +#: ``None``. Each is either assigned outright by a child phase (``_rewards`` / +#: ``_verifier_error`` by ``verify()``) or never written by a child at all +#: (``_error`` / ``_export_error`` / ``_evolved_skills`` belong to ``run()`` +#: and ``cleanup()``, which no in-place child re-runs) — so after the child, +#: a non-``None`` value is the child's own and ``None`` is the honest +#: "not produced by this child". +_SCOPED_NONE_FIELDS: tuple[str, ...] = ( + "_rewards", + "_verifier_error", + "_error", + "_export_error", + "_evolved_skills", +) + + +def scope_child_result_state(rollout: Rollout) -> None: + """Zero the result-bearing state an in-place child reports through. + + Called between the engine's per-child ``restore_onto`` (which put the + *parent's* values back) and the child runner. Only attributes the instance + already has are touched — the ``_LinearState`` convention: what is absent + stays absent — and the parent's own values are restored by the engine + before the next child and at the end of the fork, so nothing here outlives + the child it scopes. + """ + for name in _SCOPED_NONE_FIELDS: + if hasattr(rollout, name): + setattr(rollout, name, None) + if hasattr(rollout, "_timing"): + rollout._timing = {} + if hasattr(rollout, "_diagnostics"): + from benchflow.diagnostics import RolloutDiagnostics + + rollout._diagnostics = RolloutDiagnostics() + + +def write_in_place_child_result( + rollout: Rollout, + *, + parent: RolloutNode, + child: RolloutNode, + run_dir: Path, + started_at: datetime, + base_trajectory_len: int, + base_prompt_count: int, + base_tool_calls: int, +) -> None: + """Write the completed in-place child's own result artifacts, best-effort. + + The ``base_*`` arguments are the fork baseline — the lengths/counts of the + parent's captured linear state — so the child's trajectory, prompts and + tool-call count are the *delta* its continuation appended, never the + parent's history re-labelled as the child's. Everything else is read off + the instance, which :func:`scope_child_result_state` made the child's own. + + Never raises: an artifact failure is logged and the child's reward (already + recorded on its node) stands. + """ + try: + _write_result( + rollout, + parent=parent, + child=child, + run_dir=run_dir, + started_at=started_at, + base_trajectory_len=base_trajectory_len, + base_prompt_count=base_prompt_count, + base_tool_calls=base_tool_calls, + ) + except Exception: + logger.warning( + "in-place branch child %s result artifacts failed to build — the " + "child's reward is unaffected", + child.id, + exc_info=True, + ) + + +def _write_result( + rollout: Rollout, + *, + parent: RolloutNode, + child: RolloutNode, + run_dir: Path, + started_at: datetime, + base_trajectory_len: int, + base_prompt_count: int, + base_tool_calls: int, +) -> None: + """Build and write the child's result set — see the module docstring.""" + config = getattr(rollout, "_config", None) + if config is None: + # A harness-built stand-in with no config cannot honestly describe a + # run; there is nothing to synthesize a result *about*. + logger.debug( + "in-place branch child %s has no rollout config; skipping result synthesis", + child.id, + ) + return + from benchflow.rollout._results import _build_rollout_result + from benchflow.skill_policy import resolve_task_skill_policy + + child_dir = branch_child_dir(run_dir, parent.id, child.id) + + # The child's own continuation: everything appended past the fork baseline. + trajectory = list(getattr(rollout, "_trajectory", []) or [])[base_trajectory_len:] + prompts = list(getattr(rollout, "_executed_prompts", []) or [])[base_prompt_count:] + n_tool_calls = max( + 0, getattr(rollout, "_n_tool_calls", base_tool_calls) - base_tool_calls + ) + + # The child's own score. ``_rewards`` was scoped to None before the child + # ran, so a dict here is what *its* verify() produced; the fallback is the + # scalar the engine recorded on the child node (a caller-run child may + # score without writing ``_rewards``). A child the engine recorded as + # *unscored* publishes ``None`` outright — the node's unscored marker is + # the engine's verdict that no score was observed, and a leftover ``{}`` + # from a verifier that returned nothing must not read as an observation. + from benchflow.branch import UNSCORED_KEY + + if UNSCORED_KEY in child.state: + rewards = None + else: + rewards = copy.deepcopy(getattr(rollout, "_rewards", None)) + if rewards is None and "reward" in child.state: + rewards = {"reward": float(child.state["reward"])} + + # The same branch provenance the child's provenance.json carries, so + # result.json names its parent, stage and snapshot refs on its own. + snapshot = parent.state.get("snapshot") + stage = getattr(snapshot, "stage", None) + provenance = child_provenance( + str(run_dir), + branch_stage=stage if stage is not None else f"cursor:{parent.id}", + snapshot=snapshot, + delta=child.state.get("delta"), + delta_execution=child.state.get("delta_execution"), + ) + + # The world the in-place child genuinely ran in is the parent's: shared + # sandbox, shared skill deployment. The parent's resolved policy (or the + # same resolution _build_result would run) is therefore the child's own. + skill_policy = getattr(rollout, "_task_skill_policy", None) + if skill_policy is None: + skill_policy = resolve_task_skill_policy( + task_path=config.task_path, + skill_mode=config.recorded_skill_mode, + runtime_skills_dir=config.skills_dir, + declared_sandbox_skills_dir=None, + ) + sandbox_id_fn = getattr(rollout, "_current_sandbox_id", None) + _build_rollout_result( + child_dir, + task_name=config.task_path.name, + rollout_name=child.id, + agent=config.primary_agent, + agent_name=getattr(rollout, "_agent_name", "") or "", + model=config.primary_model, + n_tool_calls=n_tool_calls, + prompts=prompts, + error=getattr(rollout, "_error", None), + verifier_error=getattr(rollout, "_verifier_error", None), + export_error=getattr(rollout, "_export_error", None), + trajectory=trajectory, + partial_trajectory=bool(getattr(rollout, "_partial_trajectory", False)), + trajectory_source=getattr(rollout, "_trajectory_source", None), + rewards=rewards, + started_at=started_at, + timing=dict(getattr(rollout, "_timing", {}) or {}), + scenes=config.effective_scenes, + evolved_skills=getattr(rollout, "_evolved_skills", None), + source_provenance=provenance, + dataset=config.dataset, + task_digest=config.task_digest, + diagnostics=getattr(rollout, "_diagnostics", None), + skill_policy=skill_policy, + sandbox_id=sandbox_id_fn() if callable(sandbox_id_fn) else None, + # No usage kwargs on purpose: per-child token usage is not + # attributable on the shared instance, so the result honestly reports + # usage_source="unavailable" instead of inheriting the parent's. + ) diff --git a/src/benchflow/branch_skill.py b/src/benchflow/branch_skill.py new file mode 100644 index 000000000..d372eb142 --- /dev/null +++ b/src/benchflow/branch_skill.py @@ -0,0 +1,527 @@ +"""Executing an ``env-ready`` branch child as a fresh child rollout (RFC §3.3). + +Skills are deployed by ``install_agent()``, so a branch taken at the cursor — +after installation — cannot vary them: the container already carries the pack +(or already lacks it), and restoring environment *state* cannot put a +filesystem back. The ``env-ready`` stage snapshot (RFC §3.2) is the seam that +makes the ablation honest: it is taken at the end of ``start()``, **before** +``install_agent()``, so restoring it yields a world with no agent and no skills +in it yet. + +That seam cuts both ways, and it is a property of the **stage**, not of the +delta: a world restored to ``env-ready`` has no agent binary, no sandbox user, +no seeded verifier workspace, no path lockdown and no skill pack, because none +of them exist until ``install_agent()`` runs. An *in-place* child forked there +— the default runner's ``connect()`` → ``execute()`` → ``verify()`` on the +parent instance — therefore either dies launching an agent the restore just +deleted, or (when the agent survives in the base image) runs and scores inside +a world missing everything installation deploys, and reports that as an +ordinary child carrying one recorded delta. So **every** engine-run child of +``env-ready`` runs through this module, whatever its delta: with a +``skill_mode`` delta the mode is switched, without one the parent's own +recorded mode is re-installed unchanged. + +This module owns the child side of that path. It derives the child's +:class:`~benchflow.rollout.RolloutConfig` from the parent's with the skill mode +set — the existing skill-policy resolution then does the real work +(``no-skill`` strips the task-bundled pack out of the staged build context and +rewrites the Dockerfile ``COPY`` lines, ``with-skill`` mounts it) — and runs +the child as a first-class Rollout over the already-restored sandbox through +the ``use_prebuilt_env`` seam (#388): ``setup`` → ``install_agent`` → +``connect`` → ``execute`` → ``verify`` → ``cleanup``. ``start()`` is +deliberately skipped: the sandbox is up and the environment plane provisioned, +which is exactly what ``env-ready`` means, and most backends' ``start()`` is +not idempotent. + +The gates deciding whether such a child may run at all — the stage, the +snapshot layer it requires, the caller-supplied-runner conflict — belong to the +branch engine (:mod:`benchflow.rollout_branch`), which fails closed before +anything is quiesced, restored, or run. +""" + +from __future__ import annotations + +import dataclasses +import logging +from collections.abc import Awaitable, Callable +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from benchflow._utils.config_override import deep_merge +from benchflow.branch import UnscoredChildError +from benchflow.branch_delta import BranchDeltaNotSupported +from benchflow.branch_lineage import branch_child_dir, child_provenance +from benchflow.branch_stage import STAGE_ENV_READY +from benchflow.skill_policy import ( + SKILL_MODE_WITH_SKILL, + TaskSkillPolicy, + resolve_task_skill_policy, +) + +if TYPE_CHECKING: + from benchflow.branch_delta import BranchDelta + from benchflow.environment.manifest import EnvironmentManifest + from benchflow.rollout import Rollout, RolloutConfig + from benchflow.trajectories.tree import RolloutNode + +logger = logging.getLogger(__name__) + +#: The branch point whose children must run as fresh rollouts: it is the one +#: recorded boundary that precedes ``install_agent()``, so a child restored to +#: it has no agent and no skills and must re-run installation for itself. +#: Also the only branch point a ``skill_mode`` delta can execute from. +FRESH_CHILD_STAGE = STAGE_ENV_READY + +#: The checkpoint layer that snapshot must carry. Everything +#: ``install_agent()`` deploys — the agent binary, the sandbox user, the skill +#: pack, the path lockdown — lives in the container filesystem, so an +#: environment-state-only checkpoint cannot roll it back: a child restored +#: without this layer would re-install on top of the parent's install (and a +#: ``no-skill`` child would still find the parent's pack mounted), quietly +#: measuring nothing. +FRESH_CHILD_LAYER = "sandbox" + +#: Back-compat aliases. Skills were the original reason this boundary exists, +#: and the skill-specific diagnostics still name it that way. +SKILL_DELTA_STAGE = FRESH_CHILD_STAGE +SKILL_DELTA_LAYER = FRESH_CHILD_LAYER + +#: ``delta_execution`` in lineage: the child was run as its own Rollout rather +#: than in place on the parent instance. Absent = the ordinary in-place child. +EXECUTION_FRESH_ROLLOUT = "fresh-rollout" + + +class BranchEnvironmentImageConflict(BranchDeltaNotSupported): + """An ``environment_ref`` delta whose manifest changes the *image*. + + Branching from a snapshot restores the **parent's** container — a commit + of the image the parent's manifest named. A child manifest that names a + different image (or per-task base image) is asking for a world that + snapshot never contained: honoring it would need a rebuild-and-reprovision + path, which contradicts branching from a snapshot (the arms would no + longer share a byte-identical starting world). Fails closed naming both + images; an image-changing environment comparison is two independent runs, + not a branch. + """ + + +def resolve_environment_ref_delta( + parent_manifest: EnvironmentManifest | None, + environment_ref: str, + *, + subject: str = "environment_ref delta", +) -> EnvironmentManifest: + """Resolve an ``environment_ref`` delta into the child's manifest (RFC §3.3). + + The executable slice is the documented tool-outage pattern (``env0@prod`` + vs ``env0@outage``): the *same* image with a *different + framework-started service set*. The env-ready snapshot restores the + parent's container, whose entrypoint starts nothing when + ``owns_lifecycle = false`` — so the running services of the restored world + are exactly what the framework provisions, and provisioning the child + manifest's set over the restore executes the delta soundly. Everything + outside that slice fails closed here, before anything is quiesced: + + * a parent with **no bound manifest** — there is no Environment plane to + swap, and the child would run a world the snapshot never contained; + * a ref that does not resolve (registry spec or manifest path, through the + same :func:`~benchflow.environment.manifest.load_manifest` every run + binds with, so the delta stays content-addressed); + * a manifest that changes the **image** + (:class:`BranchEnvironmentImageConflict` — see its docstring); + * an **entrypoint-owned lifecycle** on either side: the restored + container's entrypoint (re)starts whatever the image bakes in, so the + framework can neither subtract a service from an ``owns_lifecycle`` + parent nor hand lifecycle ownership to a same-image child — the delta + would be recorded but not enforced. + + ``subject`` names the caller's delta in every diagnostic + (``deltas[1].environment_ref``, ``arm 'env:env0@outage'``). + """ + from benchflow.environment.manifest import load_manifest + + if parent_manifest is None: + raise BranchDeltaNotSupported( + f"{subject}={environment_ref!r} cannot execute: the parent bound " + "no environment manifest, so there is no Environment plane to " + "swap — the child would run against a world its snapshot never " + "contained. Bind the parent's world with " + "RolloutConfig(environment_manifest=...) (or declare it in " + "task.md) and make the delta a variation of it." + ) + try: + child_manifest = load_manifest(environment_ref) + except Exception as exc: + raise BranchDeltaNotSupported( + f"{subject}={environment_ref!r} does not resolve to an " + f"environment manifest: {exc}" + ) from exc + if (child_manifest.image, child_manifest.base_image) != ( + parent_manifest.image, + parent_manifest.base_image, + ): + raise BranchEnvironmentImageConflict( + f"{subject}={environment_ref!r} changes the environment image " + f"(parent manifest {parent_manifest.name!r} runs " + f"image={parent_manifest.image!r} / " + f"base_image={parent_manifest.base_image!r}; child manifest " + f"{child_manifest.name!r} names image={child_manifest.image!r} / " + f"base_image={child_manifest.base_image!r}). A branch child " + "restores the parent's committed container, so an image-changing " + "environment delta needs a rebuild path — which contradicts " + "branching from a snapshot. Run the two environments as " + "independent evaluations instead." + ) + if parent_manifest.owns_lifecycle or child_manifest.owns_lifecycle: + raise BranchDeltaNotSupported( + f"{subject}={environment_ref!r} needs framework-started services " + "on both sides (owns_lifecycle = false), got parent " + f"owns_lifecycle={parent_manifest.owns_lifecycle} / child " + f"owns_lifecycle={child_manifest.owns_lifecycle}: the restored " + "container's entrypoint (re)starts whatever the image bakes in, " + "so the framework cannot vary an entrypoint-owned service set — " + "the delta would be recorded in provenance but not enforced in " + "the world. Only the service-level slice (same image, different " + "[[environment.services]]) is executable from a snapshot." + ) + return child_manifest + + +async def provision_child_environment( + child_rollout: Rollout, manifest: EnvironmentManifest +) -> None: + """Provision ``manifest``'s environment plane over the child's sandbox. + + The fresh-child counterpart of the ``start()`` provisioning step the child + deliberately skips: the engine has already restored the parent's container + (killing every framework-started service with it) and rolled back the + declared environment state, so what runs *now* decides the child's service + topology. Building the plane from the **child's own** manifest is what + executes a service-level ``environment_ref`` delta — and re-provisioning + the parent's manifest for every other child is what keeps the control arm + honest: without it, a zero-delta arm of a framework-started environment + scores a world with no services at all and calls itself the baseline. + + Mirrors ``start()`` exactly: provision, then gate on readiness before the + agent is installed. The provisioned plane is attached as the child's own + ``_environment`` (before the readiness gate, so ``cleanup()`` tears its + services down even when the gate fails). + """ + environment = child_rollout._planes.manifest_environment( + manifest, sandbox=child_rollout.env + ) + child_rollout._environment = environment + await environment.provision(ctx={"task_id": child_rollout._config.task_path.name}) + probe = await environment.readiness() + if not probe.ready: + raise RuntimeError( + f"environment plane not ready for branch child " + f"{child_rollout._config.rollout_name}: {probe.error} " + f"(checked: {probe.checked})" + ) + logger.info( + "branch child environment '%s' ready (%d probe(s))", + manifest.name, + len(probe.checked), + ) + + +def fresh_child_skill_mode(config: RolloutConfig, skill_mode: str | None) -> str: + """The skill mode a fresh ``env-ready`` child installs under. + + A ``skill_mode`` delta names it; every other child re-installs the mode the + parent itself recorded (``artifact_skill_mode or skill_mode``). Falling + back to the *recorded* mode rather than the raw field is what keeps a + non-skill child a genuine zero-delta-on-skills child: a + ``from_legacy``-built parent carries its effective mode in + ``artifact_skill_mode``, and re-installing from ``skill_mode`` alone would + hand the child a different world than the parent ran in and label it "no + delta". + """ + return skill_mode if skill_mode is not None else config.recorded_skill_mode + + +def resolve_child_skill_policy( + config: RolloutConfig, skill_mode: str +) -> TaskSkillPolicy: + """Resolve the child's skill policy — the ablation's precondition. + + Called at delta-validation time, before anything is quiesced, so that a + ``with-skill`` delta against a task shipping no bundled pack (or a runtime + ``skills_dir`` that has since gone missing) fails closed with the same + typed error ``setup()`` would raise later — after the branch had already + restored a snapshot and started running children. + """ + return resolve_task_skill_policy( + task_path=config.task_path, + skill_mode=skill_mode, + runtime_skills_dir=( + config.skills_dir if skill_mode == SKILL_MODE_WITH_SKILL else None + ), + declared_sandbox_skills_dir=None, + ) + + +def child_run_location( + config: RolloutConfig, + *, + run_dir: Path | None, + parent_id: str, + child_id: str, +) -> tuple[Path, str, str]: + """Where the child rollout writes its own artifacts, as ``_init_rollout`` args. + + A branch child is a first-class rollout (RFC §3.4), so its run directory + *is* the per-child artifact directory the lineage writer already owns + (:func:`~benchflow.branch_lineage.branch_child_dir`): the standard + ``config.json`` / ``result.json`` / trajectory files land beside the + branch's own ``provenance.json``. ``_init_rollout`` composes the run + directory as ``jobs_dir / job_name / rollout_name``, so the target + directory is split back into exactly those three parts. A parent with no + run directory (a branch-first rollout that never ran ``setup()``) falls + back to its configured jobs directory under a node-derived name. + """ + target = ( + branch_child_dir(Path(run_dir), parent_id, child_id) + if run_dir is not None + else Path(config.jobs_dir) + / (config.job_name or "branches") + / f"branch-{child_id}" + ) + return target.parent.parent, target.parent.name, target.name + + +def child_skill_config( + config: RolloutConfig, + *, + skill_mode: str, + jobs_dir: Path, + job_name: str, + rollout_name: str, + source_provenance: dict[str, Any] | None = None, + config_override: dict[str, Any] | None = None, + environment_manifest: EnvironmentManifest | None = None, +) -> RolloutConfig: + """The parent's config with the delta's fields replaced — nothing else moved. + + Both ``skill_mode`` and ``artifact_skill_mode`` are written: the skill + policy resolves from ``recorded_skill_mode`` (``artifact_skill_mode or + skill_mode``), so setting only the former would leave a + ``from_legacy``-built parent's recorded mode in place and the delta would + silently do nothing. ``skills_dir`` is a with-skill-only field and is + dropped for a ``no-skill`` child; ``snapshot_stages`` is cleared because a + child forked *from* a stage does not re-checkpoint the parent's boundaries + (its environment plane belongs to the parent). + + ``config_override`` is the child's *delta* overlay (RFC §3.3): it is + deep-merged **over the parent's own overlay** so the child's effective + C-axis patch is "everything the parent ran under, plus the one recorded + change" — dropping the parent's overlay would silently vary two things. + The merged dict rides the existing seam: the child's ``setup()`` applies + it through :func:`~benchflow._utils.config_override.apply_config_override` + (same allowlist, same re-validation) and its ``config.json`` records the + merged keys + sha exactly as a run-level overlay is recorded (#790). + ``None`` inherits the parent's overlay unchanged — the zero delta. + + ``environment_manifest`` is the child's *resolved* ``environment_ref`` + delta (see :func:`resolve_environment_ref_delta`): when set it replaces + the parent's manifest wholesale, so the child's ``config.json`` records + the world it actually provisioned; ``None`` inherits the parent's — every + non-environment child re-provisions the parent's own manifest. + + The user is re-materialized from ``loop_strategy`` when the parent's was + strategy-built — ``RolloutConfig`` rejects carrying both, and the strategy + rebuilds an identical user — and left alone otherwise. + """ + merged_override = config.config_override + if config_override: + merged_override = deep_merge(merged_override or {}, config_override) + return dataclasses.replace( + config, + skill_mode=skill_mode, + artifact_skill_mode=skill_mode, + skills_dir=(config.skills_dir if skill_mode == SKILL_MODE_WITH_SKILL else None), + snapshot_stages=frozenset(), + jobs_dir=jobs_dir, + job_name=job_name, + rollout_name=rollout_name, + config_override=merged_override, + environment_manifest=( + environment_manifest + if environment_manifest is not None + else config.environment_manifest + ), + user=None if config.loop_strategy_spec is not None else config.user, + source_provenance=source_provenance, + ) + + +async def run_fresh_child( + rollout: Rollout, + config: RolloutConfig, + *, + prompts: list[str] | None = None, +) -> float: + """Run ``config`` as a fresh Rollout over ``rollout``'s restored sandbox. + + The branch engine has already rolled the container and environment back to + the ``env-ready`` snapshot, so the child adopts that sandbox + (``use_prebuilt_env``) and re-runs installation from it: this is the call + that installs the agent binary the restore removed and redeploys — or + refuses to deploy — the skill pack. ``start()`` is skipped (see the module + docstring) — except for its one piece the restore undoes: when the child's + config binds an environment manifest, the manifest's plane is provisioned + over the restored sandbox and gated on readiness before ``install_agent()`` + (:func:`provision_child_environment`), which is both how a service-level + ``environment_ref`` delta executes and how every other child gets the + parent's own services back. + + ``cleanup()`` always runs, including on failure: it releases the child's own + resources (its LiteLLM runtime, its staged task copy) and, because the + sandbox is caller-owned, leaves the parent's container up for the next + child. Exceptions propagate — a child that could not run is not a child + that scored zero. + + The child is a first-class rollout, so its terminal + :class:`~benchflow.models.RolloutResult` is materialized after cleanup — + that write is what leaves ``result.json`` / ``timing.json`` / ``prompts.json`` + beside its ``config.json`` (RFC §3.4). Artifact failure is isolated and + logged, exactly as the branch engine isolates lineage writes: the child's + reward stands either way. + + A child that *ran* but was never scored raises + :class:`~benchflow.branch.UnscoredChildError`. That is the same principle + one step further: a child whose reward does not exist is not a child that + scored zero either, and the engine records it as unscored. + """ + from benchflow.rollout import Rollout as _Rollout + + child_rollout = _Rollout(config) + child_rollout.use_prebuilt_env(rollout.env) + rewards: dict | None = None + try: + await child_rollout.setup() + if config.environment_manifest is not None: + await provision_child_environment( + child_rollout, config.environment_manifest + ) + await child_rollout.install_agent() + await child_rollout.connect() + await child_rollout.execute(prompts) + rewards = await child_rollout.verify() + finally: + await child_rollout.cleanup() + + result = None + try: + result = child_rollout.result + except Exception: + logger.warning( + "branch child result artifacts (%s) failed to build — the child's " + "reward is unaffected", + config.rollout_name, + exc_info=True, + ) + scored = (result.rewards if result is not None else None) or rewards + if not scored or scored.get("reward") is None: + # The child ran but was never scored — its verifier crashed, or its + # output never reached the host. Reporting 0.0 here is how a lost + # reward became a real-looking observation in a live ablation; the + # engine records this as an unscored child instead. + verifier_error = getattr(child_rollout, "_verifier_error", None) + raise UnscoredChildError( + f"branch child {config.rollout_name} produced no verifier reward " + f"(verify() returned {rewards!r})" + + (f" — {verifier_error}" if verifier_error else "") + ) + return float(scored["reward"]) + + +def make_fresh_child_runner( + rollout: Rollout, + *, + delta: BranchDelta | None, + parent: RolloutNode, + branch_stage: str, + run_dir: Path | None, +) -> Callable[[RolloutNode], Awaitable[float]]: + """Build the per-child runner for a child forked from ``env-ready``. + + Every engine-run child of that boundary uses this, whatever its delta — + the restored world has no agent installed, so there is nothing for an + in-place child to connect to (see the module docstring). ``delta`` may be + ``None`` (the zero delta): the child then re-installs the parent's own + recorded skill mode and runs the rollout's resolved prompts, which is what + makes it a genuine control arm rather than a differently-provisioned one. + + The child carries the branch provenance as its *own* ``source_provenance`` + — the seam a continued run already uses — so its ``config.json`` and + ``result.json`` record which rollout it forked from, at which stage, from + which snapshot refs, and that it ran as a fresh rollout. An + ``injected_prompt`` on the delta is delivered as the child's continuation + prompt, exactly as the default runner delivers it; a ``config_override`` + is deep-merged over the parent's own overlay into the child's config (see + :func:`child_skill_config`) and applied by the child's own ``setup()``. + """ + skill_mode = fresh_child_skill_mode( + rollout._config, delta.skill_mode if delta is not None else None + ) + injected_prompt = delta.injected_prompt if delta is not None else None + config_override = delta.config_override if delta is not None else None + environment_ref = delta.environment_ref if delta is not None else None + # Resolved through the same gates the engine validated the delta against + # (one source of truth), so a registry that changed between validation and + # run still cannot hand the child an unbranchable manifest. + environment_manifest = ( + resolve_environment_ref_delta( + rollout._config.environment_manifest, environment_ref + ) + if environment_ref is not None + else None + ) + + async def _runner(child: RolloutNode) -> float: + jobs_dir, job_name, rollout_name = child_run_location( + rollout._config, + run_dir=run_dir, + parent_id=parent.id, + child_id=child.id, + ) + config = child_skill_config( + rollout._config, + skill_mode=skill_mode, + jobs_dir=jobs_dir, + job_name=job_name, + rollout_name=rollout_name, + config_override=config_override, + environment_manifest=environment_manifest, + source_provenance=child_provenance( + str(run_dir) if run_dir is not None else str(rollout._config.task_path), + branch_stage=branch_stage, + snapshot=parent.state.get("snapshot"), + delta=child.state.get("delta"), + delta_execution=EXECUTION_FRESH_ROLLOUT, + ), + ) + logger.info( + "branch child %s runs as a fresh rollout with skill_mode=%s%s%s", + child.id, + skill_mode, + ( + f" config_override_keys={sorted(config_override)}" + if config_override + else "" + ), + ( + f" environment_ref={environment_ref!r}" + if environment_ref is not None + else "" + ), + ) + return await run_fresh_child( + rollout, + config, + prompts=([injected_prompt] if injected_prompt is not None else None), + ) + + return _runner diff --git a/src/benchflow/branch_stage.py b/src/benchflow/branch_stage.py new file mode 100644 index 000000000..de5f89066 --- /dev/null +++ b/src/benchflow/branch_stage.py @@ -0,0 +1,96 @@ +"""The branch stage taxonomy — the four cascade boundaries (RFC §3.2). + +The FrontierPhysics failure cascade names four stages a research-style agent +run can die in; each pins to an existing transition of the rollout lifecycle, +so this is a *naming* layer over ``rollout/__init__.py``, not a new phase +system: + +=================== ================================================== +``env-ready`` end of ``start()`` — sandbox up, environment plane + provisioned, readiness gate passed, **before** + ``install_agent()`` (so a child re-runs agent/skill + installation from a skill-free world) +``post-research`` a mid-``execute()`` cut point that cannot be + auto-detected — recorded by an explicit + ``Rollout.mark_stage()`` from the caller/harness that + knows when planning ended +``pre-verify`` agent quiesced, immediately **before** + ``planes.harden_before_verify`` +``post-verify`` after ``verify()`` completes +=================== ================================================== + +This module is a dependency-free leaf on purpose: ``RolloutConfig`` validates +``snapshot_stages`` at construction time and must not drag the branch engine's +snapshot/environment/sandbox import chain into config validation. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from typing import Any + +STAGE_ENV_READY = "env-ready" +STAGE_POST_RESEARCH = "post-research" +STAGE_PRE_VERIFY = "pre-verify" +STAGE_POST_VERIFY = "post-verify" + +#: Every stage, in lifecycle order — the order diagnostics and artifacts use. +BRANCH_STAGES: tuple[str, ...] = ( + STAGE_ENV_READY, + STAGE_POST_RESEARCH, + STAGE_PRE_VERIFY, + STAGE_POST_VERIFY, +) + +#: Stages the rollout captures by itself when they are requested. +AUTO_STAGES = frozenset({STAGE_ENV_READY, STAGE_PRE_VERIFY, STAGE_POST_VERIFY}) + +#: Stages only an explicit ``Rollout.mark_stage()`` can record. +MARKED_STAGES = frozenset(BRANCH_STAGES) - AUTO_STAGES + +_KNOWN_STAGES = frozenset(BRANCH_STAGES) + + +class BranchStageNotCaptured(LookupError): + """A branch was requested at a stage this rollout never snapshotted. + + Fail-closed counterpart of the capability errors: branching from a stage + that was never captured cannot silently degrade into a checkpoint at the + cursor — that would fork a *different* world state than the caller asked + for and mislabel it in provenance. + """ + + +def validate_stage(stage: str, *, field: str = "stage") -> str: + """Return ``stage`` if it names a taxonomy stage, else raise ``ValueError``.""" + if stage not in _KNOWN_STAGES: + raise ValueError( + f"unknown {field} {stage!r} — the branch stage taxonomy is " + f"{list(BRANCH_STAGES)!r}" + ) + return stage + + +def normalize_stages(stages: Iterable[str] | None) -> frozenset[str]: + """Coerce a requested stage set to a validated ``frozenset``. + + ``None`` and an empty iterable both mean *no stage snapshots* — today's + behavior exactly, with no snapshot call anywhere in the lifecycle. A + single string is rejected rather than silently iterated into characters. + """ + if stages is None: + return frozenset() + if isinstance(stages, str): + raise ValueError( + f"snapshot_stages must be a collection of stage names, got the " + f"string {stages!r} — pass e.g. {{'{STAGE_ENV_READY}'}}" + ) + resolved = frozenset(stages) + for stage in sorted(resolved): + validate_stage(stage, field="snapshot_stages entry") + return resolved + + +def captured_stages(registry: Mapping[str, Any]) -> list[str]: + """The recorded stages of ``registry``, in lifecycle order (deterministic).""" + return [stage for stage in BRANCH_STAGES if stage in registry] diff --git a/src/benchflow/branch_transaction.py b/src/benchflow/branch_transaction.py new file mode 100644 index 000000000..d0ea6b008 --- /dev/null +++ b/src/benchflow/branch_transaction.py @@ -0,0 +1,416 @@ +"""The branch transaction — checkpoint, per-child restore/run, state scoping. + +The execution heart of a fork (the Branch lifecycle's steps 2 and 3): once +:mod:`benchflow.branch_policy` has admitted the request, this module owns +everything that must happen *transactionally* so the parent comes out exactly +as it went in — the composed checkpoint at the branch point +(:func:`checkpoint_parent`), the scoped capture of the parent's linear and +result-bearing state (:class:`LinearState`), and the child loop +(:meth:`BranchTransaction.run_children`): restore the checkpointed layers, +reset the linear state, point the cursor at a pending child node, run the +child under the runner :mod:`benchflow.branch_children` selects, record its +reward / unscored reason / wall clock on the node, and hand its shared-mount +output into custody before the next child can inherit or destroy it. + +:class:`BranchTransaction` is one fork's execution context as a value — the +fourteen positional/keyword arguments the old ``_run_children`` took, named +once. The orchestrator (:func:`benchflow.rollout_branch.branch`) builds it +after quiesce + checkpoint and brackets it with artifact custody and the +final linear-state restore. +""" + +from __future__ import annotations + +import copy +import logging +import time +from dataclasses import dataclass, field +from datetime import datetime +from typing import TYPE_CHECKING, Any + +from benchflow.branch import UNSCORED_KEY, StageSnapshot, UnscoredChildError +from benchflow.branch import adopt_checkpoint as _adopt_checkpoint +from benchflow.branch import checkpoint as _checkpoint_branch +from benchflow.branch import checkpoint_composed as _checkpoint_composed +from benchflow.branch import restore as _restore_branch +from benchflow.branch import restore_composed as _restore_composed +from benchflow.branch_artifacts import MountedArtifacts, child_mount_dir +from benchflow.branch_children import ( + EXECUTION_FRESH_ROLLOUT, + ChildRunner, + make_fresh_child_runner, + select_child_runner, +) +from benchflow.branch_delta import BranchDelta +from benchflow.branch_result import ( + scope_child_result_state, + write_in_place_child_result, +) +from benchflow.models import TrajectorySource +from benchflow.trajectories.tree import RolloutNode + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from benchflow.rollout import Rollout + +logger = logging.getLogger(__name__) + +#: Node-state key holding how long a branch child took, in seconds. Recorded +#: on the child node (even when the child raised) so a caller comparing +#: children — ``bench eval ablate``'s per-arm cost column — can attribute time +#: without re-deriving it from per-child artifacts, which only a +#: fresh-rollout child leaves behind. Deliberately *not* serialized by +#: :mod:`benchflow.branch_lineage`: a measured duration would make +#: ``tree.json`` non-deterministic. +CHILD_WALL_CLOCK_KEY = "wall_clock_sec" + +#: The result-bearing Rollout attributes a child mutates that are not part of +#: the linear execution state above. Derived from an audit of everything +#: :meth:`Rollout._build_result` reads, against what a child's +#: ``connect() -> execute() -> verify() -> disconnect()`` writes: +#: +#: * ``_timing`` — ``execute()`` *accumulates* ``agent_execution`` and +#: ``verify()`` writes into the same dict, so without this the parent's +#: ``timing.json`` reported the parent's time plus every arm's (observed +#: live by @Galius5136 on a two-arm ``pre-verify`` ablation). +#: * ``_verifier_error`` — ``verify()`` assigns it unconditionally, so a +#: clean parent inherited the last child's verifier failure, and a failed +#: parent had its own diagnostic erased by a child that scored cleanly. +#: * ``_diagnostics`` — ``execute()`` records prompt-timeout diagnostics and +#: ``verify()`` records verifier-timeout ones; a child's timeout would +#: surface in the parent's ``result.json`` as the parent's own. +#: * ``_native_usage_metrics`` / ``_native_usage_checkpoint`` — ``execute()`` +#: accumulates native ACP token usage into the first and re-bases the +#: second; ``cleanup()``'s ``_finalize_usage_metrics`` then promotes the +#: first into ``_usage_metrics``, i.e. the children's tokens are billed to +#: the parent. The same accumulation bug as ``_timing``, one field further +#: from the result. +#: +#: The remaining four are result-bearing and reachable from a caller-supplied +#: ``run_child`` (which may drive any phase it likes), so they are scoped for +#: the same reason even though the engine's own runners do not write them: +#: ``_error``, ``_export_error``, ``_evolved_skills``, ``_usage_metrics``. +#: +#: Audited and deliberately *not* scoped: ``_rollout_dir`` / ``_rollout_name`` +#: / ``_resolved_prompts`` / ``_task_skill_policy`` (written by ``setup()``, +#: which no child re-runs on the shared instance), ``_agent_name`` (written by +#: ``connect()``, but from the parent's own config — a child cannot change what +#: it resolves to), ``_started_at`` (the parent's start is the parent's start, +#: and rolling it back per child would be wrong anyway), ``_terminal_timeout`` +#: (only ``_record_agent_timeout`` writes it, which lives in ``run()``), and +#: the ``_provider_*_cached`` trio (written by ``cleanup()`` only). A +#: fresh-rollout child touches none of these: it drives a Rollout of its own. +_RESULT_STATE_FIELDS: tuple[str, ...] = ( + "_timing", + "_verifier_error", + "_diagnostics", + "_native_usage_metrics", + "_native_usage_checkpoint", + "_error", + "_export_error", + "_evolved_skills", + "_usage_metrics", +) + + +@dataclass +class LinearState: + """A scoped snapshot of a Rollout's linear (non-tree) execution state. + + Captured before a branch child runs and restored after — this is what + makes a branch child an *isolated sub-rollout* rather than a re-entrant + mutation of the shared Rollout instance. The stage registry (RFC §3.2) is + scoped the same way: a child that runs through its own ``pre-verify`` + boundary records that snapshot on its own node, and it must not become the + *parent's* pre-verify — a later branch there would fork the child's world. + + Two kinds of state, scoped for two different reasons. The named fields are + the *execution* state — where the rollout is and what it has done. The + ``result_state`` bag is the **result-bearing** state + (:data:`_RESULT_STATE_FIELDS`): fields no child needs but every in-place + child writes, which end up in the parent's own ``result.json`` / + ``timing.json`` if they are left where the last child put them. An + in-place child continues the *shared* Rollout, so "the parent's linear + state is exactly what it was before" has to cover what the parent + *reports*, not only where its cursor sits. + + Everything mutable in that bag is deep-copied in both directions. + ``restore_onto`` runs once per child — not only at the end — so handing a + child the captured dict itself would let child *k* mutate the snapshot + child *k+1* and the parent are both restored from. + """ + + cursor: RolloutNode + trajectory: list[dict] + n_tool_calls: int + phase: str + rewards: dict | None + trajectory_source: TrajectorySource | None + partial_trajectory: bool + session_tool_count: int + session_traj_count: int + executed_prompts: list[str] + stage_snapshots: dict[str, StageSnapshot] + stage_nodes: dict[str, RolloutNode] + result_state: dict[str, Any] + + @classmethod + def capture(cls, rollout: Rollout) -> LinearState: + """Snapshot ``rollout``'s linear state — a shallow copy of the trajectory.""" + return cls( + cursor=rollout._cursor, + trajectory=list(rollout._trajectory), + n_tool_calls=rollout._n_tool_calls, + phase=rollout._phase, + rewards=copy.deepcopy(rollout._rewards), + trajectory_source=rollout._trajectory_source, + partial_trajectory=rollout._partial_trajectory, + session_tool_count=getattr(rollout, "_session_tool_count", 0), + session_traj_count=getattr(rollout, "_session_traj_count", 0), + executed_prompts=list(rollout._executed_prompts), + stage_snapshots=dict(getattr(rollout, "_stage_snapshots", {})), + stage_nodes=dict(getattr(rollout, "_stage_nodes", {})), + # hasattr, not a default: tests build a Rollout through + # ``__new__()`` (the established pattern in rollout.py), and + # inventing an attribute the instance never had would be a + # different kind of mutation. What is absent stays absent. + result_state={ + name: copy.deepcopy(getattr(rollout, name)) + for name in _RESULT_STATE_FIELDS + if hasattr(rollout, name) + }, + ) + + def restore_onto(self, rollout: Rollout) -> None: + """Write this snapshot back onto ``rollout`` — undoing a child's mutations.""" + rollout._cursor = self.cursor + rollout._trajectory = list(self.trajectory) + rollout._n_tool_calls = self.n_tool_calls + rollout._phase = self.phase + rollout._rewards = copy.deepcopy(self.rewards) + rollout._trajectory_source = self.trajectory_source + rollout._partial_trajectory = self.partial_trajectory + rollout._session_tool_count = self.session_tool_count + rollout._session_traj_count = self.session_traj_count + rollout._executed_prompts = list(self.executed_prompts) + rollout._stage_snapshots = dict(self.stage_snapshots) + rollout._stage_nodes = dict(self.stage_nodes) + for name, value in self.result_state.items(): + setattr(rollout, name, copy.deepcopy(value)) + + +async def checkpoint_parent( + rollout: Rollout, + parent: RolloutNode, + *, + stage_snapshot: StageSnapshot | None, + composed: bool, + snap_env: Any, + snap_sandbox: Any, + layers: frozenset[str], +) -> None: + """Record the fork's roll-back point on ``parent`` (Branch step 2). + + A stage branch adopts the snapshot the stage boundary already took — + re-snapshotting there would capture a world that has since moved on. A + cursor branch takes it now: the composed checkpoint when more than the + legacy environment layer is requested, else the legacy environment-only + path (same behavior, same bare ``StateSnapshot`` shape on the node). A + checkpoint failure fails the branch closed with a note naming the layers — + no partial checkpoint is recorded on the node. + """ + if stage_snapshot is not None: + _adopt_checkpoint(parent, stage_snapshot) + return + try: + if composed: + await _checkpoint_composed( + parent, environment=snap_env, sandbox=snap_sandbox + ) + else: + await _checkpoint_branch(parent, rollout._environment) + except Exception as exc: + exc.add_note( + f"branch(snapshot_layers={sorted(layers)!r}) failed during " + "checkpoint — the branch fails closed; no partial checkpoint " + "was recorded on the node." + ) + raise + + +@dataclass +class BranchTransaction: + """One fork's execution context — the parameters of the child loop, named. + + Built by :func:`benchflow.rollout_branch.branch` after quiesce and + checkpoint; :meth:`run_children` is the successor of the old + ``_run_children`` free function, whose fourteen arguments live here as + fields instead. ``children`` is the transaction's output: every child is + appended as soon as it is attached, so a caller that catches a propagating + child failure still sees the nodes that were forked — which is how ``bench + eval ablate`` reports the arm that raised and the arms that never ran. + """ + + rollout: Rollout + n: int + deltas: Sequence[BranchDelta | None] | None + parent: RolloutNode + saved: LinearState + runner: ChildRunner + run_child: ChildRunner | None + composed: bool + snap_env: Any + snap_sandbox: Any + fresh_children: bool + run_dir: Any + holder: MountedArtifacts | None + children: list[RolloutNode] = field(default_factory=list) + # Injected by the orchestrator from *its* module globals so + # ``benchflow.rollout_branch`` stays the call-time resolution point — the + # monkeypatch seam tests/test_branch_child_result.py pins (a fake + # fresh-child runner factory, a recording in-place result writer). The + # defaults are the canonical implementations for direct users. + fresh_runner_factory: Callable[..., ChildRunner] = make_fresh_child_runner + write_child_result: Callable[..., None] = write_in_place_child_result + + async def run_children(self) -> None: + """Run the fork's ``n`` children in order, appending to ``children``. + + Each iteration is the same bracket: attach a *pending* child node + carrying its delta provenance, restore the checkpointed layers and the + parent's linear state, run the child under the runner the boundary + selects, record its outcome on the node, and hand its shared-mount + output into custody. An **in-place** child of a run-dir-bearing + rollout additionally leaves its own ``result.json`` set + (:mod:`benchflow.branch_result`): its result-bearing state is scoped + to zero before it runs — so what the fields hold afterwards is the + child's own, never the parent's showing through — and the result is + synthesized from that state after the child completes, *before* the + next restore discards it. A fresh-rollout child writes its own result + already and is left alone; a child that raised ends the fork and is + evidenced by its ``mounted/`` archive instead. + """ + in_place_results = not self.fresh_children and self.run_dir is not None + for index in range(self.n): + delta = self.deltas[index] if self.deltas is not None else None + child = self._attach_pending_child(delta) + await self._restore_for(child, scope_result_state=in_place_results) + started_wall = datetime.now() + child_runner = select_child_runner( + self.rollout, + delta=delta, + default=self.runner, + run_child=self.run_child, + fresh_children=self.fresh_children, + parent=self.parent, + run_dir=self.run_dir, + fresh_runner_factory=self.fresh_runner_factory, + ) + await self._run_child(child, child_runner) + if in_place_results: + # The instance still carries the child's own state (the + # restore happens on the next iteration / at the end of the + # fork), so this is the one window where a first-class result + # can be built for an in-place child. Best-effort by contract: + # never raises. + self.write_child_result( + self.rollout, + parent=self.parent, + child=child, + run_dir=self.run_dir, + started_at=started_wall, + base_trajectory_len=len(self.saved.trajectory), + base_prompt_count=len(self.saved.executed_prompts), + base_tool_calls=self.saved.n_tool_calls, + ) + if self.holder is not None: + # Custody fails closed: a hand-off failure means the mounts + # may still carry this child's files, which the next child + # would both inherit and destroy. The child's own outcome + # (reward / unscored reason / result.json) is recorded above, + # so raising here loses no observation — the fork stops before + # evidence can cross arms. + self.holder.raise_pending() + + def _attach_pending_child(self, delta: BranchDelta | None) -> RolloutNode: + """Attach a *pending* branch-child node carrying its delta provenance. + + The child's real continuation Step is filled in place by its first + execute(), so the child's work lands on the child node, not a + descendant placeholder. The delta provenance is recorded on the node + itself at fork time (``None`` = the zero delta), so lineage + serialization reads it from the node — never by positional alignment — + and a second branch() at the same parent can never misattribute + deltas. How the delta is *executed* is recorded for the same reason: + an env-ready child re-runs installation as its own Rollout, and the + artifacts must say so rather than leaving a reader to infer it from + the delta — which they could not, since a zero-delta child of that + boundary is a fresh rollout too. + """ + child = self.rollout._tree.attach(self.parent) + child.state["delta"] = ( + delta if delta is not None else BranchDelta() + ).provenance_dict() + if self.fresh_children: + child.state["delta_execution"] = EXECUTION_FRESH_ROLLOUT + self.children.append(child) + return child + + async def _restore_for( + self, child: RolloutNode, *, scope_result_state: bool + ) -> None: + """Roll the world and the parent's linear state back for one child. + + Restore the checkpointed layers (sandbox first, then env — the reverse + of checkpoint order), reset the parent's linear state, and point the + cursor at the pending child for the sub-rollout. For an in-place child + the restore just put the *parent's* result-bearing values back; they + are scoped to zero so the child's result reports only what the child + itself produced — the next restore brings the parent's back. + """ + if self.composed: + await _restore_composed( + self.parent, environment=self.snap_env, sandbox=self.snap_sandbox + ) + else: + await _restore_branch(self.parent, self.rollout._environment) + self.saved.restore_onto(self.rollout) + self.rollout._cursor = child + if scope_result_state: + scope_child_result_state(self.rollout) + + async def _run_child(self, child: RolloutNode, child_runner: ChildRunner) -> None: + """Run one child and record its outcome on its node. + + A child that ran but was never scored records *why* + (:data:`~benchflow.branch.UNSCORED_KEY`), keeps its reward unset, and + does not end the fork: the next child restores the checkpoint for + itself, so one lost score does not cost the fork. The wall clock and + the mount hand-off happen in a ``finally`` so a child that raised + still reports what it cost — and still hands off whatever it wrote to + the shared mounts (a crashed child's verifier output is often the only + evidence of why it crashed). Handing off also empties the mounts, so + the next child cannot inherit this one's files. ``hand_off`` never + raises here — a raise inside this finally would replace the child's + own exception — it records the failure on the holder, and + ``raise_pending()`` in the loop surfaces it once the child's outcome + is safely on the node. + """ + started = time.monotonic() + unscored: str | None = None + try: + ret = await child_runner(child) + except UnscoredChildError as exc: + unscored = exc.reason + finally: + child.state[CHILD_WALL_CLOCK_KEY] = time.monotonic() - started + if self.holder is not None and self.run_dir is not None: + self.holder.hand_off( + child_mount_dir(self.run_dir, self.parent.id, child.id) + ) + if unscored is not None: + child.state[UNSCORED_KEY] = unscored + logger.error("branch child %s is unscored: %s", child.id, unscored) + else: + child.state["reward"] = float(ret) diff --git a/src/benchflow/cli/_failure_evidence.py b/src/benchflow/cli/_failure_evidence.py index cb1700371..4d1367e80 100644 --- a/src/benchflow/cli/_failure_evidence.py +++ b/src/benchflow/cli/_failure_evidence.py @@ -6,7 +6,10 @@ module so ``cli/_shared.py`` stays the side-effect-free display helpers it advertises. Also home of :func:`metric_breakdown`, the one canonical rewards-dict flattening — shared by ``_shared.py``'s in-memory metric tier and -the ``reward.json`` probe here, so the two render identically. +the ``reward.json`` probe here, so the two render identically. The CTRF report +this mines for a failure line is also the source of the *whole* per-test +outcome map (:func:`ctrf_test_outcomes`, :func:`artifact_test_outcomes`) that +``bench eval ablate`` attributes sub-outcomes with — one parser, two readings. Contract (the engine stays file-free — these reads are CLI-side, report-time only): @@ -167,6 +170,79 @@ def _bounded_json(path: Path) -> dict[str, Any] | None: return data if isinstance(data, dict) else None +def _ctrf_tests(ctrf_path: Path) -> list[dict[str, Any]]: + """The CTRF report's test entries — the one place ``ctrf.json`` is read. + + Both consumers of the report (the failure one-liner below and the per-test + outcome map :func:`ctrf_test_outcomes` mines for ablation attribution) go + through here, so there is exactly one notion of "the tests this verifier + reported": the ``results.tests`` array, non-object entries dropped. A file + that is missing, oversized or not a CTRF object yields no tests rather than + raising — the callers each decide what an empty report means for them. + """ + data = _bounded_json(ctrf_path) + if data is None: + return [] + raw_tests = (data.get("results") or {}).get("tests") or [] + return [test for test in raw_tests if isinstance(test, dict)] + + +def ctrf_test_outcomes(ctrf_path: Path) -> dict[str, str]: + """``{test name -> CTRF status}`` for every test the report names, sorted. + + The per-test half of the same report :func:`_ctrf_failure_line` mines for + its one-liner: every entry, not just the failures, because attribution + needs the passes too (a test that flips *to* passing is as much evidence as + one that flips away from it). Statuses are passed through verbatim — + ``passed`` / ``failed`` / ``skipped`` / ``pending`` / ``other`` are CTRF's + own vocabulary and this is a reporter, not a classifier. + + Names are shortened by :func:`_display_test_name` only when that stays + injective across the report: two files holding a same-named test would + otherwise collapse into one row and silently hide a difference, so such a + report keeps its raw node ids. Sorted keys, first entry wins on an exact + duplicate — the map has to be byte-stable across runs of the same input. + """ + named = [ + (str(test["name"]), str(test["status"])) + for test in _ctrf_tests(ctrf_path) + if isinstance(test.get("name"), str) + and test["name"].strip() + and isinstance(test.get("status"), str) + and test["status"].strip() + ] + display = {raw: _display_test_name(raw) for raw, _ in named} + injective = len(set(display.values())) == len(set(display)) + outcomes: dict[str, str] = {} + for raw, status in named: + outcomes.setdefault(display[raw] if injective else raw, status) + return dict(sorted(outcomes.items())) + + +def artifact_test_outcomes(rollout_dir: Path) -> dict[str, str] | None: + """One rollout's per-test outcomes, or ``None`` when it reported none. + + The rollout-dir twin of :func:`artifact_failure_evidence`, for callers that + need the whole outcome map rather than one failure line — ``bench eval + ablate`` reads it per branch child, whose child directory *is* its rollout + directory. CTRF is the only source: a verifier that emits no report yields + ``None`` (say so and attribute on the scalar), never an invented row. + Never raises. + """ + # Local import: RolloutPaths pulls in the task package (see + # artifact_failure_evidence). + from benchflow.task.paths import RolloutPaths + + try: + ctrf_path = RolloutPaths(rollout_dir=rollout_dir).verifier_dir / "ctrf.json" + if not ctrf_path.is_file(): + return None + outcomes = ctrf_test_outcomes(ctrf_path) + except Exception: + return None + return outcomes or None + + def _ctrf_failure_line(ctrf_path: Path) -> FailureLine | None: """`` failed[: ]`` from the first failed CTRF test. @@ -178,11 +254,7 @@ def _ctrf_failure_line(ctrf_path: Path) -> FailureLine | None: carries a count suffix (``(+N more failures; P/T checks passed)``) so the console never under-reports how much is broken. """ - data = _bounded_json(ctrf_path) - if data is None: - return None - raw_tests = (data.get("results") or {}).get("tests") or [] - tests = [test for test in raw_tests if isinstance(test, dict)] + tests = _ctrf_tests(ctrf_path) failed = [test for test in tests if test.get("status") == "failed"] if not failed: return None diff --git a/src/benchflow/cli/_options.py b/src/benchflow/cli/_options.py index d13fda604..a3b27a258 100644 --- a/src/benchflow/cli/_options.py +++ b/src/benchflow/cli/_options.py @@ -7,6 +7,7 @@ factored here; one-off variants stay inline in ``main.py``. """ +from pathlib import Path from typing import Annotated import typer @@ -14,6 +15,20 @@ from benchflow.sandbox.providers import providers_phrase AgentOption = Annotated[str, typer.Option("--agent", help="Agent name")] +EnvironmentManifestOption = Annotated[ + Path | None, + typer.Option( + "--environment-manifest", + help=( + "Environment-plane manifest applied to every rollout: a path to " + "an environment.toml, OR a 'name@version' registry spec (the S " + "axis) resolved via $BENCHFLOW_ENV_REGISTRY when set, else the " + "built-in registry shipped with benchflow (env0@prod, " + "env0@outage). The manifest-declared stateful environment is " + "provisioned, gated on readiness, and torn down." + ), + ), +] ModelOption = Annotated[str | None, typer.Option("--model", help="Model")] SandboxOption = Annotated[ str, typer.Option("--sandbox", help=f"Sandbox: {providers_phrase()}") diff --git a/src/benchflow/cli/ablate.py b/src/benchflow/cli/ablate.py new file mode 100644 index 000000000..4180ebcbb --- /dev/null +++ b/src/benchflow/cli/ablate.py @@ -0,0 +1,320 @@ +"""``bench eval ablate`` — stage-level ablation over branch children. + +Lives in its own module per the one-file-per-command-group convention; +:func:`register_eval_ablate` attaches it to the ``eval`` group. The command is +a thin caller: parsing arm specs, running the ablation, and turning the report +into a table or JSON all belong to :mod:`benchflow.ablate`, imported lazily +inside the command so ``bench --help`` never pays for the branch engine. + +Exit codes mirror ``bench review``: 0 when every arm produced a reward, 1 when +any arm errored or was skipped, and 1 for a request that cannot run — always +as a one-line error on stderr, never a traceback. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Annotated + +import typer +from rich.markup import escape +from rich.table import Table + +from benchflow.branch_stage import BRANCH_STAGES, STAGE_ENV_READY +from benchflow.cli._options import ( + AgentOption, + EnvironmentManifestOption, + ModelOption, + SandboxOption, +) +from benchflow.cli._shared import _apply_dotenv_to_process_env, console, print_error + +_STATUS_STYLES = { + "pass": "green", + "fail": "red", + "error": "red", + "skipped": "grey50", +} + +DEFAULT_ARMS = "with-skill,no-skill" + + +_TEST_STATUS_STYLES = {"passed": "green", "failed": "red", "skipped": "grey50"} + + +def _render_sub_test_attribution(report) -> None: + """Print the sub-test section: only the tests whose outcome differs. + + The second half of the attribution, and the half a scalar tie hides. Tying + tests are counted, never listed — they are noise for attribution — and the + section always ends on the summary line, so a reader is told *why* there is + no table (all tests tie / no per-test data was mined) rather than left to + read the absence. + """ + from benchflow.ablate import sub_test_attribution + + section = sub_test_attribution(report.arms) + differing = section["differing_tests"] + if differing: + arms = section["arms_with_tests"] + table = Table( + title=f"Sub-test outcomes that differ ({len(differing)})", + show_lines=True, + ) + table.add_column("Test") + for arm in arms: + table.add_column(escape(arm)) + for entry in differing: + cells = [] + for arm in arms: + outcome = entry["outcomes"].get(arm) + style = _TEST_STATUS_STYLES.get(outcome or "", "yellow") + cells.append( + f"[{style}]{escape(outcome) if outcome else 'not reported'}[/{style}]" + ) + table.add_row(escape(entry["test"]), *cells) + console.print(table) + if section["tying_tests"]: + console.print( + f"[dim]{len(section['tying_tests'])} test(s) tie across the " + "arms and are omitted.[/dim]" + ) + console.print(f"[bold]Sub-test attribution:[/bold] {escape(section['summary'])}") + + +def _render_ablation(report) -> None: + """Print the parent's own reward, then one row per arm.""" + parent = "n/a" if report.parent_reward is None else f"{report.parent_reward:.2f}" + console.print( + f"\n[bold]Task:[/bold] {escape(report.task_id)} " + f"[bold]stage:[/bold] {escape(report.stage)} " + f"[bold]parent reward:[/bold] {parent}" + ) + if report.environment: + env = report.environment + world = env.get("env_hash") or env.get("image") or env.get("base_image") + console.print( + f"[bold]Environment:[/bold] {escape(str(env.get('name')))} " + f"({escape(str(env.get('ref')))}, {escape(str(world))})" + ) + if report.stage_snapshot: + # The recorded roll-back handles of the branched stage — enough to + # restore this world and re-branch it by hand later. + snap = report.stage_snapshot + console.print( + f"[bold]Stage snapshot:[/bold] " + f"sandbox={escape(str(snap.get('sandbox_ref') or '-'))} " + f"environment={escape(str(snap.get('environment_ref') or '-'))}" + ) + if report.parent_error: + console.print( + f"[yellow]Parent run error:[/yellow] " + f"{escape(report.parent_error.splitlines()[0])}" + ) + table = Table( + title=f"Ablation at {escape(report.stage)}: {escape(report.task_id)}", + show_lines=True, + ) + for column in ("Arm", "Reward", "Result", "Wall clock", "Attribution"): + table.add_column(column) + for arm in report.arms: + table.add_row( + escape(arm.name), + "-" if arm.reward is None else f"{arm.reward:.2f}", + arm.status, + "-" if arm.wall_clock_sec is None else f"{arm.wall_clock_sec:.0f}s", + escape(arm.verdict), + style=_STATUS_STYLES.get(arm.status, "white"), + ) + console.print(table) + _render_sub_test_attribution(report) + if report.value is not None: + console.print(f"[bold]V(parent) over the arms:[/bold] {report.value:.2f}") + + +def _ablate_command( + tasks_dir: Annotated[ + Path, + typer.Option( + "--tasks-dir", + help=( + "Task directory to ablate (or a collection holding exactly one " + "task) — the arms are the axis, the task is fixed" + ), + ), + ], + agent: AgentOption = "claude-agent-acp", + model: ModelOption = None, + reasoning_effort: Annotated[ + str | None, + typer.Option( + "--reasoning-effort", + help=( + "Agent reasoning/thinking effort when the agent exposes one " + "(e.g. max) — the same control as bench eval run, recorded in " + "the parent's and every arm's config" + ), + ), + ] = None, + sandbox: SandboxOption = "docker", + environment_manifest: EnvironmentManifestOption = None, + at_stage: Annotated[ + str, + typer.Option( + "--at-stage", + help=( + f"Stage boundary to fork: {', '.join(BRANCH_STAGES)}. " + "'post-research' is a mid-execute() cut point that needs a " + "research-end trigger: pair it with --mark-research-end-on." + ), + ), + ] = STAGE_ENV_READY, + mark_research_end_on: Annotated[ + str | None, + typer.Option( + "--mark-research-end-on", + help=( + "Workspace file whose first appearance marks 'post-research' " + "(e.g. /app/PLAN.md, the FrontierPhysics convention): the " + "engine polls the sandbox during the agent's run (plus one " + "final check when it quiesces) and snapshots the stage the " + "first time the file exists. Required for --at-stage " + "post-research; rejected for any other stage." + ), + ), + ] = None, + arms: Annotated[ + str, + typer.Option( + "--arms", + help=( + "Comma-separated arms, one branch child each: 'with-skill', " + "'no-skill', 'inject:', " + "'config:', 'env:'" + ), + ), + ] = DEFAULT_ARMS, + out_dir: Annotated[ + Path | None, + typer.Option( + "--out-dir", + "-o", + help="Ablation output directory (default: jobs/ablate-)", + ), + ] = None, + keep_snapshots: Annotated[ + bool, + typer.Option( + "--keep-snapshots", + help=( + "Export the branched stage's sandbox snapshot image (docker " + "save) to /snapshots/.tar before cleanup " + "destroys it; ablation.json records the tar path and sha256. " + "Without it the snapshot dies with the run and its handle is " + "recorded as ephemeral" + ), + ), + ] = False, + output_json: Annotated[ + bool, + typer.Option("--json", help="Emit the ablation report as JSON on stdout"), + ] = False, +) -> None: + """Run a task once, branch a stage boundary, and compare the arms. + + The counterfactual entry point of the rollout-branching RFC: the task runs + once, the requested stage boundary is snapshotted as it passes, and that + one world is forked into a child per arm — ``with-skill`` / ``no-skill`` + switch the skill mode, ``inject:`` hands the child that file as its + continuation prompt, ``config:`` runs the child under the parent's + config with the allowlisted patch deep-merged on top (#790 machinery), + ``env:`` provisions a different registry manifest's service set over + the restored world (same image only — the tool-outage pattern). + Every arm therefore starts from a byte-identical world + and differs by exactly one recorded delta. At ``--at-stage env-ready`` + every arm re-runs agent installation as its own rollout (that boundary + precedes ``install_agent()``); at later boundaries the arms continue in + place. ``--environment-manifest`` binds the parent's (and therefore every + arm's) environment explicitly, beating the task-declared manifest — the + same precedence as ``bench eval run``. Results land in a table plus + ``ablation.json`` — which stamps the bound environment (name, ref, + content hash) and the branched stage's snapshot refs — followed by a + second table naming the sub-tests whose outcome differs across the arms: + the behavioral difference a tied scalar reward would otherwise hide. + """ + import asyncio + import json + from datetime import datetime + + from benchflow.ablate import ( + AblationError, + AblationRequest, + parse_arms, + resolve_ablation_task, + run_ablation, + validate_arms_for_stage, + write_ablation_report, + ) + + try: + parsed_arms = parse_arms(arms) + stage = validate_arms_for_stage( + parsed_arms, at_stage, research_end_marker=mark_research_end_on + ) + task_path = resolve_ablation_task(tasks_dir) + except (AblationError, ValueError) as exc: + print_error(str(exc)) + raise typer.Exit(1) from None + + if out_dir is None: + stamp = datetime.now().strftime("%Y-%m-%d__%H-%M-%S") + out_dir = Path("jobs") / f"ablate-{stamp}" + + _apply_dotenv_to_process_env() + if not output_json: + console.print( + f"\n[blue]Ablating {escape(task_path.name)} at {escape(stage)}: " + f"{escape(', '.join(arm.name for arm in parsed_arms))}[/blue]" + ) + try: + report = asyncio.run( + run_ablation( + AblationRequest( + task_path=task_path, + arms=parsed_arms, + agent=agent, + stage=stage, + model=model, + reasoning_effort=reasoning_effort, + sandbox=sandbox, + out_dir=out_dir, + environment_manifest=environment_manifest, + mark_research_end_on=mark_research_end_on, + keep_snapshots=keep_snapshots, + ) + ) + ) + except (AblationError, ValueError, FileNotFoundError) as exc: + print_error(str(exc)) + raise typer.Exit(1) from None + + report_path = write_ablation_report(report, out_dir) + if output_json: + payload = report.to_dict() + payload["report_path"] = str(report_path) + typer.echo(json.dumps(payload, sort_keys=True, indent=2)) + else: + _render_ablation(report) + console.print(f"\n[bold]Report:[/bold] {escape(str(report_path))}") + + if report.error: + print_error(report.error.splitlines()[0]) + for arm in report.arms: + if arm.error: + print_error(f"{arm.name}: {arm.error.splitlines()[0]}") + if report.has_errors: + raise typer.Exit(1) + + +def register_eval_ablate(eval_app: typer.Typer) -> None: + eval_app.command("ablate")(_ablate_command) diff --git a/src/benchflow/cli/continue_cmd.py b/src/benchflow/cli/continue_cmd.py index 051d74dc4..ee98bab96 100644 --- a/src/benchflow/cli/continue_cmd.py +++ b/src/benchflow/cli/continue_cmd.py @@ -95,6 +95,25 @@ def continue_cmd( "(no live model needed) — useful for testing.", ), ] = False, + max_exchanges: Annotated[ + int | None, + typer.Option( + "--max-exchanges", + help="Replay only the first K recorded LLM exchanges, then go " + "live (default: all recorded).", + ), + ] = None, + cut_stage: Annotated[ + str | None, + typer.Option( + "--cut-stage", + help="Cut the replay at a recorded stage boundary by name " + "(e.g. post-research): resolves the exchange index the run's " + "stage_snapshots.json recorded when that stage closed. " + "Mutually exclusive with --max-exchanges; an unrecorded stage " + "fails closed listing the stages the run did record.", + ), + ] = None, proxy_mode: Annotated[ str, typer.Option( @@ -108,6 +127,7 @@ def continue_cmd( ) -> None: """Resume a previous unfinished (timed-out) openhands run to completion.""" from benchflow.continue_run.orchestrator import continue_run + from benchflow.continue_run.replay_proxy import ReplayCutPointError from benchflow.continue_run.run_folder import RunFolderError _apply_dotenv_to_process_env() @@ -124,9 +144,11 @@ def continue_cmd( strict_divergence=strict_divergence, replay_only=replay_only, proxy_mode=proxy_mode, + max_exchanges=max_exchanges, + cut_stage=cut_stage, ) ) - except RunFolderError as exc: + except (RunFolderError, ReplayCutPointError) as exc: # Command-agnostic prefix: the same callback backs both the canonical # `bench eval continue` and the deprecated top-level `bench continue`. typer.secho(f"benchflow: {exc}", fg=typer.colors.RED, err=True) diff --git a/src/benchflow/cli/main.py b/src/benchflow/cli/main.py index 266e74ffa..d6d099af4 100644 --- a/src/benchflow/cli/main.py +++ b/src/benchflow/cli/main.py @@ -34,7 +34,12 @@ live_session, progress_enabled, ) -from benchflow.cli._options import AgentOption, ModelOption, SkillModeOption +from benchflow.cli._options import ( + AgentOption, + EnvironmentManifestOption, + ModelOption, + SkillModeOption, +) from benchflow.cli._shared import ( _apply_dotenv_to_process_env, _exit_if_evaluation_had_errors, @@ -44,6 +49,7 @@ err_console, print_error, ) +from benchflow.cli.ablate import register_eval_ablate from benchflow.cli.adopt import register_adopt_deprecated, register_eval_adopt from benchflow.cli.agent import register_agent from benchflow.cli.continue_cmd import register_continue @@ -189,6 +195,8 @@ def _show(sb, age_minutes, will_delete): # entry point; adopt makes a foreign benchmark runnable). register_eval_adopt(eval_app) register_eval_lift(eval_app) +# Stage-level ablation over branch children (rollout-branching RFC §5). +register_eval_ablate(eval_app) @eval_app.command("run") @@ -292,20 +300,7 @@ def eval_run( ), ), ] = None, - environment_manifest: Annotated[ - Path | None, - typer.Option( - "--environment-manifest", - help=( - "Environment-plane manifest applied to every rollout: a path to " - "an environment.toml, OR a 'name@version' registry spec (the S " - "axis) resolved via $BENCHFLOW_ENV_REGISTRY when set, else the " - "built-in registry shipped with benchflow (env0@prod, " - "env0@outage). The manifest-declared stateful environment is " - "provisioned, gated on readiness, and torn down." - ), - ), - ] = None, + environment_manifest: EnvironmentManifestOption = None, state: Annotated[ str | None, typer.Option( @@ -614,6 +609,22 @@ def eval_run( int, typer.Option("--trials", help="Number of trials for --matrix"), ] = 1, + keep_snapshots: Annotated[ + bool, + typer.Option( + "--keep-snapshots", + help=( + "Export each captured stage-snapshot image (docker save) to " + "/snapshots/.tar before cleanup destroys it; " + "stage_snapshots.json records the tar's path, sha256 and " + "image id. Without it the images die with the run and each " + "recorded ref is marked ephemeral. Load them back later with " + "`bench eval import-snapshots`. Only meaningful when the run " + "captures stage snapshots (a task-declared branch_execution: " + "forked-snapshot, or an SDK snapshot_stages request)" + ), + ), + ] = False, ) -> None: # The supported --sandbox values are rendered from the provider registry # into that option's own help text. This docstring used to hand-copy them @@ -684,6 +695,7 @@ def eval_run( eval_results_task=eval_results_task, matrix=matrix, trials=trials, + keep_snapshots=keep_snapshots, ) # --source-path/--source-ref only apply to --source-repo; otherwise they're # silently ignored (e.g. `--dataset X --source-ref abc` drops the ref). @@ -1102,6 +1114,51 @@ def _run_source_env_eval( eval_create = eval_run +@eval_app.command("import-snapshots") +def eval_import_snapshots( + run_dir: Annotated[ + Path, + typer.Argument( + help=( + "A completed run directory holding stage_snapshots.json " + "(and, from a --keep-snapshots run, snapshots/.tar)" + ), + ), + ], + stage: Annotated[ + list[str] | None, + typer.Option( + "--stage", + help=( + "Import only this stage's snapshot; repeatable " + "(default: every exported stage)" + ), + ), + ] = None, +) -> None: + """Load a run's exported stage-snapshot images back into Docker. + + The import half of ``bench eval run --keep-snapshots`` (rollout-branching + RFC §3.6): verifies each exported tar's recorded sha256, ``docker load``s + it, and confirms the recorded ``bf-snap-…`` ref resolves to the recorded + image id — after which the recorded stage boundary can be restored and + branched even though the run that captured it is gone. Entries recorded + ``ephemeral: true`` fail closed with the flag to re-run with. + """ + from benchflow.snapshot_import import SnapshotImportError, import_stage_snapshots + + try: + imported = import_stage_snapshots(run_dir, stages=stage) + except SnapshotImportError as exc: + print_error(str(exc)) + raise typer.Exit(1) from None + for snap in imported: + console.print( + f"[green]✓[/green] {escape(snap.stage)}: {escape(snap.sandbox_ref)} " + f"({escape(snap.image_id)})" + ) + + @eval_app.command("list") # Evaluation inspection commands diff --git a/src/benchflow/continue_run/orchestrator.py b/src/benchflow/continue_run/orchestrator.py index 3b9a92310..fc6c1d141 100644 --- a/src/benchflow/continue_run/orchestrator.py +++ b/src/benchflow/continue_run/orchestrator.py @@ -21,6 +21,7 @@ from __future__ import annotations +import asyncio import json import logging import os @@ -31,7 +32,14 @@ from typing import Any, cast from benchflow._utils.text import describe_exception -from benchflow.continue_run.replay_proxy import ReplayProxy, ReplayRouter +from benchflow.continue_run.replay_proxy import ( + REQUEST_DIGEST_BASIS, + ReplayCutPointError, + ReplayProxy, + ReplayRouter, + comparable_request_digest, + validate_max_exchanges, +) from benchflow.continue_run.run_folder import RunFolder, RunFolderError, load_run_folder from benchflow.continue_run.sandbox_proxy import ( SandboxReplayProxy, @@ -89,6 +97,188 @@ def continued_rollout_name(run: RunFolder) -> str: return f"{cleaned}__continued" +def stage_tags_from_run(run: RunFolder, cut_stage: str) -> dict[str, int]: + """The recorded stage→exchange-index tags backing a ``--cut-stage`` request. + + The run-folder half of the RFC §3.5 stage-tagged cut: stage captures + record their completed-exchange index into the run's + ``stage_snapshots.json`` (see ``benchflow.branch_policy.capture_stage``), + and this resolves that registry for ``cut_stage``, failing closed with a + typed :class:`ReplayCutPointError` that tells the caller what *was* + recorded: + + * no registry at all — the run never captured a stage; + * ``cut_stage`` absent — names the stages the run did record; + * recorded without an index (``exchanges_completed: null``) — the usage + gateway could not count at capture time, so the stage cannot name a cut; + * recorded at index 0 — the stage closed before the first exchange, so + there is no recorded prefix to replay. + + An explicit SDK ``stage_tags`` mapping bypasses this entirely (it stays + the override seam); the CLI resolves through here. + """ + registry = run.stage_registry + if not registry: + raise ReplayCutPointError( + f"--cut-stage {cut_stage!r}: the run folder recorded no stage " + "snapshots (no stage_snapshots.json) — re-run the task with stage " + "capture (RolloutConfig.snapshot_stages / Rollout.mark_stage()) " + "or cut by number with --max-exchanges." + ) + if cut_stage not in registry: + raise ReplayCutPointError( + f"unknown cut stage {cut_stage!r}; this run recorded: " + f"{', '.join(sorted(registry))}." + ) + tags = run.stage_exchange_tags + if cut_stage not in tags: + raise ReplayCutPointError( + f"stage {cut_stage!r} was recorded without an exchange index " + "(exchanges_completed is null — the run's usage gateway could not " + "count exchanges at capture time), so it cannot name a cut; use " + "--max-exchanges instead." + ) + if tags[cut_stage] == 0: + raise ReplayCutPointError( + f"stage {cut_stage!r} closed before the first LLM exchange " + "(exchanges_completed is 0) — there is no recorded prefix to " + "replay; run the task fresh instead of continuing it." + ) + return tags + + +def resolve_cut_point( + max_exchanges: int | None, + *, + stage_tags: dict[str, int] | None = None, + cut_stage: str | None = None, +) -> tuple[int | None, str | None]: + """Resolve an optional stage-named cut into a ``max_exchanges`` value. + + ``stage_tags[stage]`` is the 1-based count of exchanges that had completed + when the stage closed — a cut at that stage replays exactly that many + exchanges (rollout-branching RFC §3.5). The mapping comes from the run + folder's recorded registry (:func:`stage_tags_from_run`) unless the SDK + caller supplies its own override dict. With ``cut_stage`` given, the cut + resolves to ``stage_tags[cut_stage]``; every misuse — combining + ``cut_stage`` with ``max_exchanges``, a missing ``stage_tags`` mapping, an + unknown stage, or a tag below 1 — fails closed with + :class:`ReplayCutPointError`. Returns ``(max_exchanges, branch_stage)``. + """ + if cut_stage is None: + return max_exchanges, None + if max_exchanges is not None: + raise ReplayCutPointError( + "pass either max_exchanges or cut_stage, not both — a stage-named " + "cut resolves max_exchanges from stage_tags." + ) + if not stage_tags: + raise ReplayCutPointError( + f"cut_stage={cut_stage!r} given but no stage_tags were provided; a " + "stage-named cut needs the stage-name → completed-exchange-count " + "mapping." + ) + if cut_stage not in stage_tags: + raise ReplayCutPointError( + f"unknown cut stage {cut_stage!r}; available stages: " + f"{', '.join(sorted(stage_tags))}." + ) + resolved = stage_tags[cut_stage] + if resolved < 1: + raise ReplayCutPointError( + f"stage_tags[{cut_stage!r}] must be >= 1 — a stage tag is the " + f"1-based count of exchanges completed when the stage closed, got " + f"{resolved}." + ) + return resolved, cut_stage + + +def cut_point_provenance( + exchanges: list[LLMExchange], + *, + max_exchanges: int | None = None, + branch_stage: str | None = None, +) -> dict[str, Any]: + """The configured ``cut_point`` block recorded in source_provenance (RFC §3.5). + + ``n_replayed_exchanges`` is the replay prefix the run is *configured* to + serve (the natural end when ``max_exchanges`` is ``None``) and + ``recorded_request_digest`` the comparable digest of the last replayed + exchange's *recorded* request — both computed from the recording at + config time, so the block carries ``accounting: "configured"`` and every + field only the live run could observe is recorded honestly as null: + ``served_request_digest`` (no incoming request exists yet), + ``divergences`` (not observed on this basis) and ``workspace_digest`` + (no sandbox is reachable at the cut point on this basis — the in-sandbox + replay proxy crosses the cut without notifying the host). In sandbox + proxy mode the uploaded recording is truncated to exactly this prefix, + so configured accounting is the honest basis there; in host proxy mode + the orchestrator reconciles the block after the run with what the live + router actually served (:func:`served_cut_point`). ``branch_stage`` is + recorded only for a stage-named cut. + """ + n_replayed = validate_max_exchanges(max_exchanges, len(exchanges)) + block: dict[str, Any] = { + "n_replayed_exchanges": n_replayed, + "recorded_request_digest": ( + comparable_request_digest(exchanges[n_replayed - 1].request.body) + if n_replayed > 0 + else None + ), + "served_request_digest": None, + "request_digest_basis": REQUEST_DIGEST_BASIS, + "accounting": "configured", + "divergences": None, + "workspace_digest": None, + "workspace_digest_reason": ( + "configured-basis block computed from the recording alone; no " + "live sandbox is reachable at the cut point on this basis" + ), + } + if branch_stage is not None: + block["branch_stage"] = branch_stage + return block + + +def served_cut_point( + router: ReplayRouter, + *, + configured_max_exchanges: int | None = None, + branch_stage: str | None = None, +) -> dict[str, Any]: + """The post-run ``cut_point`` block reconciled with what was served. + + Host proxy mode keeps the live :class:`ReplayRouter` in-process, so after + the run the block records the replay prefix the proxy *actually* served + (``accounting: "served"``): ``served_request_digest`` is the comparable + digest of the ACTUAL request the agent sent at the cut, + ``recorded_request_digest`` the recorded request it answered for — named + separately so a reader can compare them — plus every divergence event the + router detected and the workspace digest taken as the run crossed the cut + (null with a reason when none could be, never fabricated). When an + explicit cut was requested, ``configured_max_exchanges`` preserves the + requested K so a served/configured divergence (e.g. the agent stopped + asking before the cut) stays visible; ``branch_stage`` is kept for a + stage-named cut. + """ + block: dict[str, Any] = { + "n_replayed_exchanges": router.n_replayed_exchanges, + "served_request_digest": router.served_request_digest, + "recorded_request_digest": router.recorded_request_digest, + "request_digest_basis": REQUEST_DIGEST_BASIS, + "accounting": "served", + "divergences": list(router.divergence_events), + "workspace_digest": router.workspace_digest, + } + if router.workspace_digest is None: + block["workspace_digest_reason"] = router.workspace_digest_reason + if configured_max_exchanges is not None: + block["configured_max_exchanges"] = configured_max_exchanges + if branch_stage is not None: + block["branch_stage"] = branch_stage + return block + + @dataclass(frozen=True) class ContinuedUsageSummary: """Token usage recovered from the stitched LLM trajectory.""" @@ -161,6 +351,8 @@ def build_rollout_config( timeout: int | None, output_dir: str | Path, rollout_name: str, + max_exchanges: int | None = None, + branch_stage: str | None = None, ) -> Any: """Assemble the RolloutConfig that re-runs the task through the proxy. @@ -198,6 +390,11 @@ def build_rollout_config( "original_model": run.model, "live_model": live_model, "n_recorded_exchanges": run.n_recorded_exchanges, + "cut_point": cut_point_provenance( + run.exchanges, + max_exchanges=max_exchanges, + branch_stage=branch_stage, + ), }, ) @@ -254,19 +451,30 @@ def __call__(self, request_body: dict[str, Any]) -> dict[str, Any]: def stitched_trajectory_lines( - original_llm_trajectory: Path, live_exchanges: list[LLMExchange] + recorded_lines: list[str], + live_exchanges: list[LLMExchange], + *, + max_recorded: int | None = None, ) -> list[str]: """Build the continuous llm_trajectory: recorded prefix + live suffix. - The recorded prefix is taken verbatim from the source file (already redacted - and byte-identical to what the agent replayed); the live suffix is the - exchanges the proxy captured after the cut-point, redacted on the way out. + ``recorded_lines`` are the *parsed* exchanges' verbatim source lines + (:func:`~benchflow.continue_run.run_folder.load_llm_exchanges_with_lines` + — already redacted and byte-identical to what the agent replayed). The + replay counts parsed exchanges, so ``max_recorded`` selects exactly the + first K replayed exchanges' lines — a malformed (never-replayed) line in + the source file can neither appear in the stitched prefix nor shift the + cut. The live suffix is the exchanges the proxy captured after the + cut-point, redacted on the way out. + + ``max_recorded`` is the prefix that was *served*, which is the configured + cut only when the agent consumed all of it — callers pass the router's + served count in host proxy mode, so a run the agent abandoned early does + not get recorded responses it never received appended to its trajectory. """ - lines: list[str] = [] - if original_llm_trajectory.is_file(): - for raw in original_llm_trajectory.read_text().splitlines(): - if raw.strip(): - lines.append(raw) + lines = list(recorded_lines) + if max_recorded is not None: + lines = lines[:max_recorded] for exchange in live_exchanges: raw = json.dumps(exchange.model_dump(mode="json"), default=str) lines.append(redact_trajectory_text(raw)) @@ -274,12 +482,18 @@ def stitched_trajectory_lines( def write_stitched_trajectory( - rollout_dir: Path, original_llm_trajectory: Path, live_exchanges: list[LLMExchange] + rollout_dir: Path, + recorded_lines: list[str], + live_exchanges: list[LLMExchange], + *, + max_recorded: int | None = None, ) -> Path: """Write the stitched continuous trajectory into the new rollout folder.""" out = rollout_dir / "trajectory" / "llm_trajectory.jsonl" out.parent.mkdir(parents=True, exist_ok=True) - lines = stitched_trajectory_lines(original_llm_trajectory, live_exchanges) + lines = stitched_trajectory_lines( + recorded_lines, live_exchanges, max_recorded=max_recorded + ) out.write_text("\n".join(lines) + ("\n" if lines else "")) return out @@ -295,7 +509,13 @@ def _usage_int(usage: dict[str, Any], *keys: str) -> int: def summarize_llm_trajectory_usage( trajectory_path: Path, *, n_recorded: int ) -> ContinuedUsageSummary: - """Recover aggregate token usage from provider responses in a trajectory.""" + """Recover aggregate token usage from provider responses in a trajectory. + + ``n_recorded`` is the recorded/live boundary inside the stitched file and + must be the same prefix length that produced it (the served count in host + proxy mode, the configured one in sandbox mode) — otherwise live tokens + are billed as replayed, or the reverse. + """ input_tokens = 0 output_tokens = 0 cache_read = 0 @@ -351,6 +571,7 @@ def update_continued_metadata( live_model: str | None, usage: ContinuedUsageSummary, environment: str, + cut_point: dict[str, Any] | None = None, ) -> None: """Patch output metadata that Rollout could not know while model=None. @@ -358,12 +579,19 @@ def update_continued_metadata( provider resolution does not clobber the replay proxy. After the run, the HF-compatible artifacts still need the actual live model and token usage. The stitched LLM trajectory is authoritative for provider usage. + + ``cut_point`` (host proxy mode) replaces the configured ``cut_point`` + block in both files' ``source`` with the served-accounting block + (:func:`served_cut_point`) — reconciling provenance with what the live + router actually replayed. """ config_path = rollout_dir / "config.json" if config_path.is_file(): config = json.loads(config_path.read_text()) config["model"] = live_model config.setdefault("source", {})["usage_source"] = "stitched_llm_trajectory" + if cut_point is not None: + config.setdefault("source", {})["cut_point"] = cut_point config["usage_tracking"] = { "requested": "required", "status": ( @@ -384,6 +612,12 @@ def update_continued_metadata( return result = json.loads(result_path.read_text()) result["model"] = live_model + if cut_point is not None: + source = result.get("source") + if not isinstance(source, dict): + source = {} + result["source"] = source + source["cut_point"] = cut_point agent_result = result.setdefault("agent_result", {}) if isinstance(agent_result, dict): agent_result.update(usage.as_agent_result_patch()) @@ -411,6 +645,88 @@ def update_continued_metadata( result_path.write_text(json.dumps(result, indent=2) + "\n") +def write_continuation_artifacts( + rollout_dir: Path, + run: RunFolder, + live_exchanges: list[LLMExchange], + *, + n_recorded: int, + cut_point: dict[str, Any], + live_model: str | None, +) -> Path: + """Write the stitched trajectory and its metadata on one replay basis. + + The three artifacts a continuation leaves behind only make sense together: + the stitched ``llm_trajectory.jsonl`` (recorded prefix + live suffix), the + recorded-vs-live token split derived from it, and the ``cut_point`` + provenance block. Producing them here, from a single ``n_recorded``, + is what keeps them from disagreeing — a prefix length used for two of the + three and a different one for the last leaves artifacts no experiment can + be run on. + + The caller chooses the basis and passes the matching block: host proxy + mode reads the prefix the live router actually *served* + (:func:`served_cut_point`), sandbox proxy mode uploads a recording already + truncated to the *configured* prefix and keeps that + (:func:`cut_point_provenance`). The block's ``accounting`` field is what + tells a reader which of the two produced these numbers. + """ + stitched_path = write_stitched_trajectory( + rollout_dir, + run.exchange_lines, + live_exchanges, + max_recorded=n_recorded, + ) + update_continued_metadata( + rollout_dir, + live_model=live_model, + usage=summarize_llm_trajectory_usage(stitched_path, n_recorded=n_recorded), + environment=run.environment, + cut_point=cut_point, + ) + return stitched_path + + +def _host_workspace_digest_fn( + rollout: Any, loop: asyncio.AbstractEventLoop +) -> Callable[[], dict[str, Any]]: + """The router's cut-point workspace hook, bound to a live host-mode run. + + Called by :class:`ReplayRouter` on the proxy's HTTP handler thread the + moment the first live-leg request arrives — the workspace the replay just + rebuilt is still live in the rollout's sandbox at that moment, so the + digest is scheduled onto the run's event loop + (:func:`~benchflow.sandbox.workspace_digest.compute_workspace_digest`) + and awaited from the handler thread. Raises — with the reason — when no + sandbox is attached or the hook is somehow invoked on the loop thread + itself (waiting there would deadlock the run); the router records the + reason and never fabricates a digest. + """ + from benchflow.sandbox.workspace_digest import compute_workspace_digest + + def _digest() -> dict[str, Any]: + sandbox = getattr(rollout, "env", None) + if sandbox is None: + raise RuntimeError( + "no live sandbox is attached to the continuation rollout" + ) + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + if running is loop: + raise RuntimeError( + "cut point crossed on the run's own event-loop thread; " + "waiting for the digest here would deadlock the run" + ) + future = asyncio.run_coroutine_threadsafe( + compute_workspace_digest(sandbox), loop + ) + return future.result(timeout=180) + + return _digest + + def _host_proxy_binding(environment: str) -> tuple[str, str]: """(bind_host, advertise_host) so a Docker agent can reach the host proxy. @@ -486,6 +802,9 @@ async def continue_run( strict_divergence: bool = False, replay_only: bool = False, proxy_mode: str = "auto", + max_exchanges: int | None = None, + stage_tags: dict[str, int] | None = None, + cut_stage: str | None = None, ) -> ContinueResult: """Resume ``folder`` to completion via record-replay + live continuation. @@ -493,8 +812,25 @@ async def continue_run( ``gemini-3.1-flash-lite-preview``); defaults to the original run's model. ``replay_only`` skips the live leg (rebuild-and-stop) — useful for testing the replay without provider credentials. + ``max_exchanges`` replays only the first K recorded exchanges, then goes + live (the RFC §3.5 cut-point). ``cut_stage`` names the cut by recorded + stage instead: the run folder's ``stage_snapshots.json`` carries the + 1-based count of exchanges that had completed when each recorded stage + closed (:func:`stage_tags_from_run`), and a cut at that stage replays + exactly that many exchanges. ``stage_tags`` overrides the recorded + registry when the SDK caller supplies its own stage→count mapping. """ run = load_run_folder(folder, require_timeout=require_timeout) + if cut_stage is not None and stage_tags is None: + stage_tags = stage_tags_from_run(run, cut_stage) + max_exchanges, branch_stage = resolve_cut_point( + max_exchanges, stage_tags=stage_tags, cut_stage=cut_stage + ) + # Fail closed on an out-of-range cut here, before anything runs. Each proxy + # mode then resolves its own replay basis: host mode reads the served count + # off the live router after the run, sandbox mode uploads a recording + # truncated to the configured prefix and keeps that basis. + validate_max_exchanges(max_exchanges, run.n_recorded_exchanges) task_path = resolve_task_path(run, tasks_dir) live_model = model or run.model @@ -528,12 +864,15 @@ async def continue_run( output_dir=out_root, rollout_name=rollout_name, strict_divergence=strict_divergence, + max_exchanges=max_exchanges, + branch_stage=branch_stage, ) router = ReplayRouter( run.exchanges, live_forwarder=live_forwarder, strict_divergence=strict_divergence, + max_exchanges=max_exchanges, ) bind_host, advertise_host = _host_proxy_binding(run.environment) @@ -550,33 +889,51 @@ async def continue_run( timeout=timeout if timeout is not None else run.timeout_sec, output_dir=out_root, rollout_name=rollout_name, + max_exchanges=max_exchanges, + branch_stage=branch_stage, ) rollout = await Rollout.create(config) + # RFC §3.5 workspace accounting: the router digests the continuation + # workspace the moment the run crosses the cut into the live leg — + # the one moment the replay-rebuilt workspace is both complete and + # still live in the sandbox. If the run never crosses (or the digest + # fails), the cut_point block records null with the reason. + router.workspace_digest_fn = _host_workspace_digest_fn( + rollout, asyncio.get_running_loop() + ) result = await rollout.run() rollout_dir = Path(rollout._rollout_dir or (out_root / rollout_name)) finally: proxy.stop() - stitched_path = write_stitched_trajectory( + # The served prefix, not the configured one, is what the agent actually + # consumed: an agent that exits or errors before reaching the cut leaves + # ``n_replayed_exchanges < n_replay``. Stitching the configured prefix + # would append recorded responses the agent never received, and the usage + # split would bill those tokens as "replayed" — while the cut_point block + # below reported the smaller served count, contradicting both. One basis + # for all three (``accounting: "served"`` names it in the artifact). + n_served = router.n_replayed_exchanges + write_continuation_artifacts( rollout_dir, - run.path / "trajectory" / "llm_trajectory.jsonl", + run, router.live_exchanges, - ) - update_continued_metadata( - rollout_dir, - live_model=live_model, - usage=summarize_llm_trajectory_usage( - stitched_path, - n_recorded=run.n_recorded_exchanges, + n_recorded=n_served, + # Reconcile the configured cut_point block with what the live router + # actually served — host mode's post-run accounting. + cut_point=served_cut_point( + router, + configured_max_exchanges=max_exchanges, + branch_stage=branch_stage, ), - environment=run.environment, + live_model=live_model, ) return ContinueResult( rollout_dir=rollout_dir, rewards=getattr(result, "rewards", None), error=getattr(result, "error", None), - n_recorded=run.n_recorded_exchanges, + n_recorded=n_served, n_live=len(router.live_exchanges), divergences=router.divergences, ) @@ -591,6 +948,8 @@ async def _continue_run_with_sandbox_proxy( output_dir: Path, rollout_name: str, strict_divergence: bool, + max_exchanges: int | None = None, + branch_stage: str | None = None, ) -> ContinueResult: """Run continuation with replay and provider proxies inside the sandbox.""" from benchflow.providers.runtime import ( @@ -599,6 +958,10 @@ async def _continue_run_with_sandbox_proxy( ) from benchflow.rollout import Rollout + # The in-sandbox proxy replays exactly the recording it is given, so a cut + # is threaded by truncating the uploaded prefix — the switch to live then + # happens exactly as if the recording had ended there. + n_replay = validate_max_exchanges(max_exchanges, run.n_recorded_exchanges) config = build_rollout_config( run, task_path=task_path, @@ -607,6 +970,8 @@ async def _continue_run_with_sandbox_proxy( timeout=timeout if timeout is not None else run.timeout_sec, output_dir=output_dir, rollout_name=rollout_name, + max_exchanges=max_exchanges, + branch_stage=branch_stage, ) rollout = await Rollout.create(config) replay_proxy: SandboxReplayProxy | None = None @@ -627,19 +992,22 @@ async def _write_artifacts_before_cleanup() -> None: return rollout_dir = Path(rollout._rollout_dir or (output_dir / rollout_name)) live_exchanges = replay_proxy.live_exchanges if replay_proxy is not None else [] - stitched_path = write_stitched_trajectory( + write_continuation_artifacts( rollout_dir, - run.path / "trajectory" / "llm_trajectory.jsonl", + run, live_exchanges, - ) - update_continued_metadata( - rollout_dir, - live_model=live_model, - usage=summarize_llm_trajectory_usage( - stitched_path, - n_recorded=run.n_recorded_exchanges, + n_recorded=n_replay, + # No live router on this side to read a served count off: the + # in-sandbox proxy is handed a recording already truncated to the + # configured prefix, so configured accounting is the honest basis + # here — and the artifact says so, rather than leaving the reader + # to infer which basis produced the numbers beside it. + cut_point=cut_point_provenance( + run.exchanges, + max_exchanges=max_exchanges, + branch_stage=branch_stage, ), - environment=run.environment, + live_model=live_model, ) artifacts_written = True @@ -661,7 +1029,7 @@ async def _write_artifacts_before_cleanup() -> None: ) replay_proxy = await SandboxReplayProxy.start( sandbox=rollout.env, - recorded=run.exchanges, + recorded=run.exchanges[:n_replay], upstream_url=provider_env["LLM_BASE_URL"], upstream_api_key=provider_env["LLM_API_KEY"], upstream_model=provider_env["LLM_MODEL"], @@ -741,7 +1109,7 @@ async def _write_artifacts_before_cleanup() -> None: rollout_dir=rollout_dir, rewards=getattr(result, "rewards", None), error=getattr(result, "error", None), - n_recorded=run.n_recorded_exchanges, + n_recorded=n_replay, n_live=len(live_exchanges), divergences=0, ) diff --git a/src/benchflow/continue_run/replay_proxy.py b/src/benchflow/continue_run/replay_proxy.py index 0d730503b..8ba903e30 100644 --- a/src/benchflow/continue_run/replay_proxy.py +++ b/src/benchflow/continue_run/replay_proxy.py @@ -29,6 +29,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any, cast +from benchflow._utils.config_override import overlay_hash from benchflow.trajectories.types import LLMExchange, LLMRequest, LLMResponse logger = logging.getLogger(__name__) @@ -53,9 +54,75 @@ def _n_messages(body: dict[str, Any]) -> int: return len(msgs) if isinstance(msgs, list) else 0 +class ReplayCutPointError(ValueError): + """Raised when a requested replay cut-point lies outside the recording.""" + + +def validate_max_exchanges(max_exchanges: int | None, n_recorded: int) -> int: + """Resolve the replay prefix length; fail closed on an out-of-range cut. + + ``None`` means the full recorded prefix (today's behavior). An explicit cut + must satisfy ``0 < K <= n_recorded`` — validated here, at configuration + time, never mid-flight. + """ + if max_exchanges is None: + return n_recorded + if not 0 < max_exchanges <= n_recorded: + raise ReplayCutPointError( + f"max_exchanges must be in 1..{n_recorded} (the recorded exchange " + f"count), got {max_exchanges}" + ) + return max_exchanges + + +def request_body_digest(body: dict[str, Any]) -> str: + """Content address of a request body — canonical JSON, ``sha256:`` form. + + A thin wrapper over the config-overlay hash + (:func:`benchflow._utils.config_override.overlay_hash`, #790) so the two + canonicalizations can never drift: sorted keys, compact separators, + ``sha256:``-prefixed digest. + """ + return overlay_hash(body) + + +#: The basis every request digest in cut-point accounting is computed on — +#: recorded next to the digests so a reader knows what "equal" means. +REQUEST_DIGEST_BASIS = ( + "sha256 over canonical JSON of the comparable projection {messages, tools}" +) + + +def comparable_request_projection(body: dict[str, Any]) -> dict[str, Any]: + """The content-bearing fields a recorded and a live request share. + + The recorded request is a *normalized projection* written by the LiteLLM + logger (model/messages/tools/stream…), not the verbatim HTTP body — and + the live replay request carries the proxy's placeholder model + (``openai/replay``) plus transport fields (``stream``, ``temperature``, + …) the recording may lack. Digesting either side whole would therefore + flag every exchange as divergent. ``messages`` and ``tools`` are the + fields that exist on both sides and actually carry the conversation, so + they are the honest comparison basis (:data:`REQUEST_DIGEST_BASIS`). + """ + return { + key: body[key] for key in ("messages", "tools") if body.get(key) is not None + } + + +def comparable_request_digest(body: dict[str, Any]) -> str: + """Digest of :func:`comparable_request_projection` — the comparison unit.""" + return request_body_digest(comparable_request_projection(body)) + + class ReplayRouter: """Serve recorded responses in order, then live — the core replay logic. + ``max_exchanges`` cuts the replay short (rollout-branching RFC §3.5): at + most the first K recorded exchanges are served, then the router switches to + live passthrough exactly as if the recording had ended there. ``None`` + replays the full recorded prefix. + Thread-safe: a single agent is typically serial, but ``ThreadingHTTPServer`` may overlap requests, so the cursor is advanced under a lock. """ @@ -66,47 +133,161 @@ def __init__( *, live_forwarder: LiveForwarder | None = None, strict_divergence: bool = False, + max_exchanges: int | None = None, + workspace_digest_fn: Callable[[], dict[str, Any]] | None = None, ) -> None: self._recorded = recorded + self._n_replay = validate_max_exchanges(max_exchanges, len(recorded)) self._live_forwarder = live_forwarder self._strict = strict_divergence self._lock = threading.Lock() self._cursor = 0 self.divergences = 0 + #: One event per detected divergence: the exchange index and both + #: sides' comparable digests, so the artifact can say *where* the + #: replay left the rails, not only that it did. + self.divergence_events: list[dict[str, Any]] = [] + # Comparable digests of the last replayed exchange — the actually + # served request and the recorded one it answered for. + self._last_served_digest: str | None = None + self._last_recorded_digest: str | None = None + # RFC §3.5 workspace accounting: called exactly once, at the first + # request past the cut, while the continuation's sandbox still holds + # the workspace the replay rebuilt. ``None`` = no sandbox reachable + # in this mode; the provenance then records null with a reason. + self.workspace_digest_fn = workspace_digest_fn + self.workspace_digest: dict[str, Any] | None = None + self._workspace_digest_error: str | None = None + self._cut_crossed = False # Live-leg exchanges, in order, for stitching onto the recorded prefix. self.live_exchanges: list[LLMExchange] = [] @property def exhausted(self) -> bool: - return self._cursor >= len(self._recorded) + return self._cursor >= self._n_replay + + @property + def n_replayed_exchanges(self) -> int: + """Recorded exchanges actually served so far (== the cut once live).""" + return min(self._cursor, self._n_replay) + + @property + def served_request_digest(self) -> str | None: + """Comparable digest of the ACTUAL incoming request last replayed. + + Cut-point accounting (RFC §3.5): this is the request the agent really + sent at the cut, digested on :data:`REQUEST_DIGEST_BASIS` — compare it + to :attr:`recorded_request_digest` to see whether the replayed + conversation was still the recorded one. ``None`` while nothing has + been replayed. + """ + return self._last_served_digest + + @property + def recorded_request_digest(self) -> str | None: + """Comparable digest of the recorded request at the same position.""" + return self._last_recorded_digest + + @property + def cut_crossed(self) -> bool: + """Whether any request arrived past the replay prefix.""" + return self._cut_crossed + + @property + def workspace_digest_reason(self) -> str | None: + """Why :attr:`workspace_digest` is ``None`` — never fabricated.""" + if self.workspace_digest is not None: + return None + if self._workspace_digest_error is not None: + return self._workspace_digest_error + if not self._cut_crossed: + return ( + "the run never crossed the cut point into the live leg, so " + "there was no moment to digest the workspace at" + ) + return ( + "no live sandbox was reachable at the cut point in this mode; " + "no workspace digest hook was configured" + ) def _check_divergence( - self, incoming: dict[str, Any], recorded_req: dict[str, Any] + self, + incoming: dict[str, Any], + recorded_req: dict[str, Any], + *, + index: int, + served_digest: str, + recorded_digest: str, ) -> None: - """Soft-validate that replay is still on the original rails. - - The recorded request is a *normalized projection* (model/messages/tools), - not the verbatim HTTP body, so we only compare message counts — a cheap - signal that the agent's conversation has the expected shape at this turn. + """Validate that replay is still on the original rails — by content. + + Compares the comparable projections' canonical-JSON digests, so a + same-message-count change to prompt content, tool definitions or tool + arguments is detected — the case the old message-count heuristic was + blind to. A recorded request with no ``messages`` (a differently + shaped recording) offers nothing to compare against and is skipped, + exactly as the count heuristic skipped it. + + A divergence *annotates* rather than aborts by default: replay + fidelity is best-effort by design (RFC §3.5 — "recorded, not + hidden"), and legitimately shifting content (timestamps in prompts, + nondeterministic tool output) would otherwise kill continuations the + user asked for. Every event is recorded in + :attr:`divergence_events`, so the artifact is truthful either way; + ``strict_divergence`` remains the opt-in abort. """ - want = _n_messages(recorded_req) - got = _n_messages(incoming) - if want and got and got != want: - self.divergences += 1 - msg = ( - f"replay divergence at turn {self._cursor}: agent sent {got} " - f"messages, recorded turn had {want}" + if not comparable_request_projection(recorded_req).get("messages"): + return + if served_digest == recorded_digest: + return + event = { + "exchange_index": index, + "served_request_digest": served_digest, + "recorded_request_digest": recorded_digest, + "n_messages_served": _n_messages(incoming), + "n_messages_recorded": _n_messages(recorded_req), + } + self.divergences += 1 + self.divergence_events.append(event) + msg = ( + f"replay divergence at exchange {index}: the agent's request " + f"({event['n_messages_served']} messages, {served_digest}) does " + f"not match the recorded one ({event['n_messages_recorded']} " + f"messages, {recorded_digest})" + ) + if self._strict: + raise ReplayDivergenceError(msg) + logger.warning(msg) + + def _capture_workspace_digest(self) -> None: + """Run the cut-point workspace hook once; record failure, never raise.""" + if self.workspace_digest_fn is None: + return + try: + self.workspace_digest = self.workspace_digest_fn() + except Exception as exc: + self._workspace_digest_error = ( + f"workspace digest failed at the cut point: {exc}" ) - if self._strict: - raise ReplayDivergenceError(msg) - logger.warning(msg) + logger.warning(self._workspace_digest_error, exc_info=True) def next_response(self, request_body: dict[str, Any]) -> ReplayResult: """Route one request to its recorded response, or to the live model.""" with self._lock: - if self._cursor < len(self._recorded): + if self._cursor < self._n_replay: exchange = self._recorded[self._cursor] - self._check_divergence(request_body, exchange.request.body) + index = self._cursor + served_digest = comparable_request_digest(request_body) + recorded_digest = comparable_request_digest(exchange.request.body) + self._last_served_digest = served_digest + self._last_recorded_digest = recorded_digest + self._check_divergence( + request_body, + exchange.request.body, + index=index, + served_digest=served_digest, + recorded_digest=recorded_digest, + ) self._cursor += 1 return ReplayResult( source="replay", @@ -115,8 +296,16 @@ def next_response(self, request_body: dict[str, Any]) -> ReplayResult: ) # Past the cut-point: live continuation. self._cursor += 1 + crossing = not self._cut_crossed + self._cut_crossed = True forwarder = self._live_forwarder + if crossing: + # Outside the lock: the hook may exec into the sandbox. The + # workspace it digests is the one the replay just rebuilt — + # this request is the first the recording no longer covers. + self._capture_workspace_digest() + if forwarder is None: logger.error( "recorded responses exhausted at turn %d and no live forwarder " diff --git a/src/benchflow/continue_run/run_folder.py b/src/benchflow/continue_run/run_folder.py index 7064bb62b..57dd20fba 100644 --- a/src/benchflow/continue_run/run_folder.py +++ b/src/benchflow/continue_run/run_folder.py @@ -19,7 +19,7 @@ import json import logging -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -43,6 +43,10 @@ class RunFolder: Only fields ``benchflow continue`` needs are surfaced; the raw ``config`` and ``result`` dicts are kept for anything else the orchestrator wants. + ``exchange_lines[i]`` is the verbatim source line of ``exchanges[i]`` — + malformed lines are skipped at load time and therefore never appear here, + so a stitched trajectory built from these lines contains exactly what the + replay serves. """ path: Path @@ -50,6 +54,13 @@ class RunFolder: result: dict[str, Any] prompts: list[str] exchanges: list[LLMExchange] + exchange_lines: list[str] + # ``stage_snapshots.json``'s ``stages`` mapping when the original run + # recorded stage boundaries (rollout-branching RFC §3.2/§3.5) — empty for + # runs captured before stage snapshots existed or without them. Each entry + # may carry ``exchanges_completed``, the completed-exchange index recorded + # at capture time, which is what a ``--cut-stage`` request resolves. + stage_registry: dict[str, Any] = field(default_factory=dict) # ── derived task identity (from config.json) ────────────────────────── @property @@ -109,6 +120,30 @@ def is_timeout(self) -> bool: def n_recorded_exchanges(self) -> int: return len(self.exchanges) + # ── recorded stage boundaries (stage_snapshots.json) ────────────────── + @property + def recorded_stages(self) -> list[str]: + """The stage boundaries the original run recorded, in file order.""" + return list(self.stage_registry) + + @property + def stage_exchange_tags(self) -> dict[str, int]: + """``stage -> completed-exchange count`` for stages with a usable index. + + A stage recorded with ``exchanges_completed: null`` (the run's usage + gateway could not count at capture time) is omitted — callers that + need to distinguish "never recorded" from "recorded without an index" + read :attr:`stage_registry` directly. + """ + tags: dict[str, int] = {} + for stage, entry in self.stage_registry.items(): + value = ( + entry.get("exchanges_completed") if isinstance(entry, dict) else None + ) + if isinstance(value, int) and not isinstance(value, bool): + tags[stage] = value + return tags + def _read_json(path: Path, *, required: bool) -> dict[str, Any]: if not path.is_file(): @@ -126,6 +161,34 @@ def _read_json(path: Path, *, required: bool) -> dict[str, Any]: return data +def _load_stage_registry(path: Path) -> dict[str, Any]: + """Read ``stage_snapshots.json``'s ``stages`` mapping, tolerantly. + + The file is optional (runs recorded before stage snapshots existed, or + without any stage capture, do not have it) and advisory: a malformed + registry degrades to empty with a warning rather than stranding a + continuation that never asked for a stage-named cut. The strict, typed + errors live where a ``--cut-stage`` request actually resolves + (:func:`benchflow.continue_run.orchestrator.stage_tags_from_run`). + """ + if not path.is_file(): + return {} + try: + data = _read_json(path, required=False) + except RunFolderError as exc: + logger.warning("ignoring unreadable stage registry: %s", exc) + return {} + stages = data.get("stages") + if not isinstance(stages, dict): + logger.warning( + "ignoring %s: expected a 'stages' mapping, got %s", + path, + type(stages).__name__, + ) + return {} + return stages + + def _load_prompts(path: Path) -> list[str]: """Read ``prompts.json`` — a JSON list of strings (or ``{"prompts": [...]}``).""" if not path.is_file(): @@ -141,12 +204,19 @@ def _load_prompts(path: Path) -> list[str]: return [str(p) for p in data if p is not None] -def load_llm_exchanges(path: Path) -> list[LLMExchange]: - """Parse ``llm_trajectory.jsonl`` into ordered :class:`LLMExchange` records. +def load_llm_exchanges_with_lines( + path: Path, +) -> tuple[list[LLMExchange], list[str]]: + """Parse ``llm_trajectory.jsonl`` into exchanges plus their raw lines. One exchange per line (``Trajectory.to_jsonl``). Blank lines are skipped; a malformed line is skipped with a warning rather than aborting the whole - resume (a single bad record should not strand a recoverable run). + resume (a single bad record should not strand a recoverable run). The + second list carries, per *parsed* exchange, its verbatim source line — + the replay counts parsed exchanges, so consumers that reconstruct the + replayed prefix (trajectory stitching) must select these lines rather + than truncating by raw file-line index, or a skipped malformed line would + shift the cut. """ if not path.is_file(): raise RunFolderError( @@ -154,6 +224,7 @@ def load_llm_exchanges(path: Path) -> list[LLMExchange]: "trajectory. Was this run captured with usage tracking enabled?" ) exchanges: list[LLMExchange] = [] + lines: list[str] = [] for lineno, raw in enumerate(path.read_text().splitlines(), start=1): if not raw.strip(): continue @@ -161,11 +232,22 @@ def load_llm_exchanges(path: Path) -> list[LLMExchange]: exchanges.append(LLMExchange.model_validate_json(raw)) except Exception as exc: logger.warning("skipping malformed llm_trajectory line %d: %s", lineno, exc) + else: + lines.append(raw) if not exchanges: raise RunFolderError( f"{path} contained no usable LLM exchanges — nothing to replay." ) - return exchanges + return exchanges, lines + + +def load_llm_exchanges(path: Path) -> list[LLMExchange]: + """Parse ``llm_trajectory.jsonl`` into ordered :class:`LLMExchange` records. + + The original public contract, kept backward compatible; see + :func:`load_llm_exchanges_with_lines` for the raw-line-aware variant. + """ + return load_llm_exchanges_with_lines(path)[0] def load_run_folder(folder: str | Path, *, require_timeout: bool = False) -> RunFolder: @@ -183,7 +265,10 @@ def load_run_folder(folder: str | Path, *, require_timeout: bool = False) -> Run config = _read_json(path / "config.json", required=True) result = _read_json(path / "result.json", required=False) prompts = _load_prompts(path / "prompts.json") - exchanges = load_llm_exchanges(path / "trajectory" / "llm_trajectory.jsonl") + exchanges, exchange_lines = load_llm_exchanges_with_lines( + path / "trajectory" / "llm_trajectory.jsonl" + ) + stage_registry = _load_stage_registry(path / "stage_snapshots.json") run = RunFolder( path=path, @@ -191,6 +276,8 @@ def load_run_folder(folder: str | Path, *, require_timeout: bool = False) -> Run result=result, prompts=prompts, exchanges=exchanges, + exchange_lines=exchange_lines, + stage_registry=stage_registry, ) if run.agent != "openhands": diff --git a/src/benchflow/contracts/planes.py b/src/benchflow/contracts/planes.py index 839be0937..abfdafa57 100644 --- a/src/benchflow/contracts/planes.py +++ b/src/benchflow/contracts/planes.py @@ -24,10 +24,21 @@ class LiveUsageGateway(Protocol): ``benchflow.providers.litellm_runtime``; renaming the accessor on either side fails the type checker there instead of silently blanking the dashboard signal. + + ``live_exchange_count`` is the stage-marker seam (rollout-branching RFC + §3.5), read the same duck-typed way by + ``benchflow.branch_policy.capture_stage``: the count of provider LLM + exchanges that had *completed* when a stage boundary was recorded, + drained to the gateway log's end so a stage-named replay cut cannot + silently undercount. ``None`` means the count is genuinely unknown + (capture never started, or the tail could not catch up) — never a stale + lower bound. """ def live_usage_tokens(self) -> int | None: ... + async def live_exchange_count(self) -> int | None: ... + @runtime_checkable class RolloutPlanes(Protocol): diff --git a/src/benchflow/environment/manifest.py b/src/benchflow/environment/manifest.py index 6b8cc2a56..00fb5dfab 100644 --- a/src/benchflow/environment/manifest.py +++ b/src/benchflow/environment/manifest.py @@ -22,11 +22,14 @@ import logging import os import tomllib +from dataclasses import dataclass, replace from pathlib import Path from typing import Literal from pydantic import BaseModel, Field, model_validator +from benchflow._utils.content_address import sha256_prefixed + logger = logging.getLogger(__name__) @@ -201,17 +204,32 @@ def model_validate_path(cls, path: str | Path) -> EnvironmentManifest: return cls.model_validate_toml(text) -def load_manifest(path: str | Path) -> EnvironmentManifest: - """Load and validate an environment manifest. +@dataclass(frozen=True) +class ManifestBinding: + """A loaded manifest plus the provenance of how it was bound. - ``path`` is either a TOML file path (the historical behavior) or a registry - spec ``name@version`` resolved via ``$BENCHFLOW_ENV_REGISTRY`` when set, - else the built-in registry shipped inside the wheel - (``benchflow/environment/_registry/``). The spec form - lets a run bind its environment (the ``S`` axis) by name at the command line — - decoupled from the task and swappable per run, like ``--agent`` / ``--model`` - / ``--sandbox``. Resolution is content-addressed so the bound world is - recorded for replay. + ``ref`` is the caller's reference verbatim — a manifest path, a registry + ``name@version`` spec, or the string a ``task.md`` declared — which stays + deterministic across machines in a way a resolved absolute path is not. + ``env_hash`` is the registry's content address (``sha256:`` of the + manifest file bytes), computed for file refs too, so the bound world is + identifiable even when it never went through the registry. ``bench eval + ablate`` stamps both into ``ablation.json``, answering "which environment + did this run compare arms in" from the report alone. + """ + + manifest: EnvironmentManifest + ref: str + env_hash: str + + +def load_manifest_binding(path: str | Path) -> ManifestBinding: + """Load a manifest and keep the provenance of its resolution. + + Same resolution as :func:`load_manifest` — a manifest file path or a + registry ``name@version`` spec — returning the manifest together with the + verbatim ref and its content address, for callers that record which world + they bound (:class:`ManifestBinding`). """ p = Path(path) if not p.is_file(): @@ -228,8 +246,85 @@ def load_manifest(path: str | Path) -> EnvironmentManifest: resolved.manifest_path, resolved.env_hash, ) - return EnvironmentManifest.model_validate_path(resolved.manifest_path) - return EnvironmentManifest.model_validate_path(p) + return ManifestBinding( + manifest=EnvironmentManifest.model_validate_path( + resolved.manifest_path + ), + ref=str(path), + env_hash=resolved.env_hash, + ) + manifest = EnvironmentManifest.model_validate_path(p) + return ManifestBinding( + manifest=manifest, ref=str(path), env_hash=sha256_prefixed(p.read_bytes()) + ) + + +def load_manifest(path: str | Path) -> EnvironmentManifest: + """Load and validate an environment manifest. + + ``path`` is either a TOML file path (the historical behavior) or a registry + spec ``name@version`` resolved via ``$BENCHFLOW_ENV_REGISTRY`` when set, + else the built-in registry shipped inside the wheel + (``benchflow/environment/_registry/``). The spec form + lets a run bind its environment (the ``S`` axis) by name at the command line — + decoupled from the task and swappable per run, like ``--agent`` / ``--model`` + / ``--sandbox``. Resolution is content-addressed so the bound world is + recorded for replay (:func:`load_manifest_binding` returns that record). + """ + return load_manifest_binding(path).manifest + + +def manifest_binding_from_task_document(task_dir: Path) -> ManifestBinding | None: + """The bound manifest ``task_dir``'s ``task.md`` declares, or ``None``. + + The binding-returning form of :func:`manifest_from_task_document` — same + extraction, same resolution, same fail-closed behavior — with the + declared string kept verbatim as the binding's ``ref`` (never the resolved + absolute path, which would vary by machine), for callers that stamp which + world the task bound. + """ + task_md = task_dir / "task.md" + if not task_md.is_file(): + return None + + from benchflow.task.document import TaskDocument + + document = TaskDocument.from_path(task_md) + environment = document.benchflow.get("environment") + if environment is None: + return None + if not isinstance(environment, dict): + raise ValueError("task.md benchflow.environment must be a mapping") + manifest = environment.get("manifest") + if manifest is None: + return None + if not isinstance(manifest, str) or not manifest.strip(): + raise ValueError("task.md benchflow.environment.manifest must be a path") + + manifest_path = Path(manifest) + if not manifest_path.is_absolute(): + manifest_path = task_dir / manifest_path + return replace(load_manifest_binding(manifest_path), ref=manifest) + + +def manifest_from_task_document(task_dir: Path) -> EnvironmentManifest | None: + """The manifest ``task_dir``'s ``task.md`` declares, or ``None``. + + A task binds its own world by declaring ``benchflow.environment.manifest`` + in its task document; the value is a manifest path relative to the task + directory (absolute paths and registry ``name@version`` specs also resolve, + via :func:`load_manifest`). Every command that runs a task resolves it + through *this* function — ``bench eval`` for a normal evaluation + (:mod:`benchflow.evaluation`) and ``bench eval ablate`` for the parent its + arms fork from (:mod:`benchflow.ablate`) — so the two cannot drift into + running the same task against different environments. + + A declaration that is present but unusable raises rather than degrading to + ``None``: a task that says it needs an image, services and readiness gates + must not silently run without them. + """ + binding = manifest_binding_from_task_document(task_dir) + return None if binding is None else binding.manifest def resolve_manifest_runtime_env( diff --git a/src/benchflow/eval_plan.py b/src/benchflow/eval_plan.py index af5fd9f44..7cd85e2b6 100644 --- a/src/benchflow/eval_plan.py +++ b/src/benchflow/eval_plan.py @@ -132,6 +132,7 @@ class EvalCreateRequest: eval_results_task: str | None = None matrix: Path | None = None trials: int = 1 + keep_snapshots: bool = False @dataclass @@ -216,6 +217,7 @@ def make_eval_config( environment_manifest=self.eval_env_manifest, config_override=self.eval_config_override, loop_strategy=self.eval_loop_strategy, + keep_snapshots=req.keep_snapshots, ) @@ -442,6 +444,22 @@ def build_eval_plan(request: EvalCreateRequest) -> EvalPlan: raise EvalPlanError( f"Invalid --reasoning-effort {request.reasoning_effort!r}: {exc}" ) from None + if eval_reasoning_effort and (request.tasks_dir or request.source_repo): + # Pre-flight the agent/effort pairing while it is knowable statically + # (PR #1046 second review, P2-A): an agent with no ACP effort channel + # otherwise rejects the run only after the sandbox was built and the + # agent installed — and ablation then retried the same doomed config + # in a branch child. Gated to the CLI-authoritative agent paths, the + # same rule as the effective_model() validation above: --config / + # --source-env may resolve their agent from YAML / the hosted source + # later, so pre-validating those would falsely reject. + from benchflow.acp.runtime import reasoning_effort_preflight_error + + effort_error = reasoning_effort_preflight_error( + eval_agent, eval_reasoning_effort + ) + if effort_error is not None: + raise EvalPlanError(effort_error) output_jobs_dir = request.jobs_dir or "jobs" # Resolve the optional Environment-plane manifest once and reuse across diff --git a/src/benchflow/eval_sharding.py b/src/benchflow/eval_sharding.py index ea462a510..62d114dd2 100644 --- a/src/benchflow/eval_sharding.py +++ b/src/benchflow/eval_sharding.py @@ -156,6 +156,7 @@ def _config_payload( "loop_strategy": ( config.loop_strategy.to_mapping() if config.loop_strategy else None ), + "keep_snapshots": config.keep_snapshots, } payload.update(config.usage_tracking.to_mapping()) return payload diff --git a/src/benchflow/eval_worker.py b/src/benchflow/eval_worker.py index 0df90656e..f5ff287d4 100644 --- a/src/benchflow/eval_worker.py +++ b/src/benchflow/eval_worker.py @@ -82,6 +82,7 @@ def _evaluation_config(raw: dict[str, Any]) -> EvaluationConfig: if raw.get("loop_strategy") else None ), + keep_snapshots=bool(raw.get("keep_snapshots", False)), ) diff --git a/src/benchflow/evaluation.py b/src/benchflow/evaluation.py index 1c96c1f41..372baecff 100644 --- a/src/benchflow/evaluation.py +++ b/src/benchflow/evaluation.py @@ -22,7 +22,7 @@ from dataclasses import dataclass, field from datetime import datetime from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import yaml @@ -53,6 +53,7 @@ PROVIDER_AUTH, PROVIDER_RATE_LIMIT, PROVIDER_REJECTED, + REQUEST_GLOBAL, SANDBOX_SETUP, SUSPECTED_API_ERROR, VERIFIER_DEP_INSTALL, @@ -71,7 +72,10 @@ from benchflow._utils.source_provenance import summary_source_fields from benchflow._utils.text import truncate_end from benchflow.diagnostics import DIAGNOSTIC_REGISTRY, summary_warning -from benchflow.environment.manifest import EnvironmentManifest +from benchflow.environment.manifest import ( + EnvironmentManifest, + manifest_from_task_document, +) from benchflow.learner_store import LearnerState, LearnerStore from benchflow.loop_strategies import ( LoopStrategySpec, @@ -94,6 +98,9 @@ from benchflow.trajectories.tree import RolloutNode from benchflow.usage_tracking import UsageTrackingConfig +if TYPE_CHECKING: + from benchflow.rollout import RolloutConfig + # Backward-compat alias RunResult = RolloutResult @@ -111,35 +118,6 @@ # flight, skip — there's nothing new to clean since the in-flight one started. _PRUNE_LOCK = threading.Lock() - -def _environment_manifest_from_task_document( - task_dir: Path, -) -> EnvironmentManifest | None: - task_md = task_dir / "task.md" - if not task_md.is_file(): - return None - - from benchflow.environment.manifest import load_manifest - from benchflow.task.document import TaskDocument - - document = TaskDocument.from_path(task_md) - environment = document.benchflow.get("environment") - if environment is None: - return None - if not isinstance(environment, dict): - raise ValueError("task.md benchflow.environment must be a mapping") - manifest = environment.get("manifest") - if manifest is None: - return None - if not isinstance(manifest, str) or not manifest.strip(): - raise ValueError("task.md benchflow.environment.manifest must be a path") - - manifest_path = Path(manifest) - if not manifest_path.is_absolute(): - manifest_path = task_dir / manifest_path - return load_manifest(manifest_path) - - _SENTINEL: Any = object() # default value for _sdk; tests replace with AsyncMock @@ -210,6 +188,11 @@ class RetryConfig: PROVIDER_AUTH, PROVIDER_RATE_LIMIT, PROVIDER_REJECTED, + # A rejected request-global setting (unsupported reasoning + # effort / model for the agent) re-fails identically on every + # attempt (PR #1046 second review, P2-A) — same family as the + # #917 provider_auth exclusion. + REQUEST_GLOBAL, } ) @@ -526,6 +509,10 @@ class EvaluationConfig: # and stamped in summary.json; None = single-shot. A dict (the to_mapping() # shape) is also accepted at runtime — __post_init__ materializes it. loop_strategy: LoopStrategySpec | str | None = None + # Durable stage-snapshot retention (`bench eval run --keep-snapshots`, + # RFC §3.6): threaded to RolloutConfig.keep_snapshots so each rollout + # exports its captured stage images before cleanup destroys them. + keep_snapshots: bool = False def __post_init__(self): from benchflow._utils.config import ( @@ -576,6 +563,93 @@ def __post_init__(self): ) +def task_rollout_config( + cfg: EvaluationConfig, + task_dir: Path, + *, + job_name: str | None, + jobs_dir: str | Path, + **overrides: Any, +) -> RolloutConfig: + """The canonical per-task :class:`RolloutConfig` an evaluation executes. + + This is the one place the evaluation's controls and provenance become a + rollout's config: the dataset identity and per-task content digest (registry + digest for a pinned ``--dataset`` run, live-computed for a dev run — every + trajectory stays attributable to the exact task content it ran), the + task-declared environment manifest fallback, the task-level source + provenance, and every normalized control the plan resolved (reasoning + effort, prompts, agent env, usage tracking, idle timeout, sandbox settings). + + ``overrides`` are the *caller-owned* fields layered on top after the + canonical assembly: the continual-learning path's per-generation + ``skills_dir`` / ``skill_mode`` / ``export_generated_skills_to``, and + ``bench eval ablate``'s stage-capture request plus its pinned ``no-skill`` + parent. Callers overlay only the axis they own — everything else must flow + from here, so an ablation's parent/child ``config.json`` carries the same + ``task_digest`` / ``reasoning_effort`` / provenance a plain ``bench eval + run`` of the task would (the PR #1046 review finding: a hand-rolled reduced + config published ``task_digest: null`` from a real E2E run). + """ + from benchflow._utils.benchmark_repos import task_source_provenance + from benchflow.rollout import RolloutConfig + + dataset = None + if cfg.dataset_name: + dataset = {"name": cfg.dataset_name, "version": cfg.dataset_version} + task_digest_value = ( + cfg.dataset_task_digests.get(task_dir.name) if cfg.dataset_name else None + ) + if task_digest_value is None: + # Dev runs (--tasks-dir / --source-repo) stamp a live-computed + # digest so every trajectory stays attributable to the exact + # task content it ran, not just a directory name. + from benchflow._utils.task_authoring import task_digest + + try: + task_digest_value = task_digest(task_dir) + except (OSError, ValueError, UnicodeError) as e: + logger.debug("Could not compute task digest for %s: %s", task_dir, e) + environment_manifest = overrides.pop("environment_manifest", None) + if environment_manifest is None: + environment_manifest = cfg.environment_manifest + if environment_manifest is None: + environment_manifest = manifest_from_task_document(task_dir) + kwargs: dict[str, Any] = dict( + task_path=task_dir, + agent=cfg.agent, + model=cfg.model, + reasoning_effort=cfg.reasoning_effort, + prompts=cfg.prompts, + agent_env=cfg.agent_env, + job_name=job_name, + jobs_dir=str(jobs_dir), + concurrency=cfg.concurrency, + environment=cfg.environment, + environment_manifest=environment_manifest, + config_override=cfg.config_override, + skills_dir=cfg.skills_dir, + sandbox_user=cfg.sandbox_user, + sandbox_locked_paths=cfg.sandbox_locked_paths, + sandbox_setup_timeout=cfg.sandbox_setup_timeout, + skip_agent_install=cfg.skip_agent_install, + agent_idle_timeout=cfg.agent_idle_timeout, + context_root=cfg.context_root, + base_image_override=cfg.base_image_override, + skill_mode=cfg.skill_mode, + skill_creator_dir=cfg.skill_creator_dir, + self_gen_no_internet=cfg.self_gen_no_internet, + source_provenance=task_source_provenance(cfg.source_provenance, task_dir), + dataset=dataset, + task_digest=task_digest_value, + usage_tracking=cfg.usage_tracking, + loop_strategy=cfg.loop_strategy, + keep_snapshots=cfg.keep_snapshots, + ) + kwargs.update(overrides) + return RolloutConfig.from_legacy(**kwargs) + + @dataclass(frozen=True) class TaskFailure: """Cheap failure evidence for one FAILED (scored, reward != 1) task. @@ -920,6 +994,7 @@ def _from_native_yaml(cls, raw: dict, **kwargs) -> Evaluation: environment_manifest=env_manifest, config_override=raw.get("config_override"), loop_strategy=raw.get("loop_strategy"), + keep_snapshots=bool(raw.get("keep_snapshots", False)), ) return cls(tasks_dir=tasks_dir, jobs_dir=jobs_dir, config=config, **kwargs) @@ -1232,25 +1307,8 @@ async def _run_single_task( skill set (``_learner_skills_dir``) and its agent-evolved skills are captured back through ``export_generated_skills_to``. """ - from benchflow._utils.benchmark_repos import task_source_provenance - from benchflow.rollout import Rollout, RolloutConfig + from benchflow.rollout import Rollout - dataset = None - if cfg.dataset_name: - dataset = {"name": cfg.dataset_name, "version": cfg.dataset_version} - task_digest_value = ( - cfg.dataset_task_digests.get(task_dir.name) if cfg.dataset_name else None - ) - if task_digest_value is None: - # Dev runs (--tasks-dir / --source-repo) stamp a live-computed - # digest so every trajectory stays attributable to the exact - # task content it ran, not just a directory name. - from benchflow._utils.task_authoring import task_digest - - try: - task_digest_value = task_digest(task_dir) - except (OSError, ValueError, UnicodeError) as e: - logger.debug("Could not compute task digest for %s: %s", task_dir, e) skills_dir = ( str(self._learner_skills_dir) if self._learner_skills_dir is not None @@ -1266,39 +1324,14 @@ async def _run_single_task( if self._learner_export_dir is not None else None ) - environment_manifest = cfg.environment_manifest - if environment_manifest is None: - environment_manifest = _environment_manifest_from_task_document(task_dir) - rollout_config = RolloutConfig.from_legacy( - task_path=task_dir, - agent=cfg.agent, - model=cfg.model, - reasoning_effort=cfg.reasoning_effort, - prompts=cfg.prompts, - agent_env=cfg.agent_env, + rollout_config = task_rollout_config( + cfg, + task_dir, job_name=self._job_name, jobs_dir=str(self._jobs_dir), - concurrency=cfg.concurrency, - environment=cfg.environment, - environment_manifest=environment_manifest, - config_override=cfg.config_override, skills_dir=skills_dir, - sandbox_user=cfg.sandbox_user, - sandbox_locked_paths=cfg.sandbox_locked_paths, - sandbox_setup_timeout=cfg.sandbox_setup_timeout, - skip_agent_install=cfg.skip_agent_install, - agent_idle_timeout=cfg.agent_idle_timeout, - context_root=cfg.context_root, - base_image_override=cfg.base_image_override, skill_mode=skill_mode, - skill_creator_dir=cfg.skill_creator_dir, - self_gen_no_internet=cfg.self_gen_no_internet, export_generated_skills_to=export_to, - source_provenance=task_source_provenance(cfg.source_provenance, task_dir), - dataset=dataset, - task_digest=task_digest_value, - usage_tracking=cfg.usage_tracking, - loop_strategy=cfg.loop_strategy, ) if skill_mode == SKILL_MODE_SELF_GEN: from benchflow.self_gen import run_self_gen diff --git a/src/benchflow/providers/litellm_runtime.py b/src/benchflow/providers/litellm_runtime.py index fc53f1186..29efb08e4 100644 --- a/src/benchflow/providers/litellm_runtime.py +++ b/src/benchflow/providers/litellm_runtime.py @@ -103,6 +103,10 @@ # normal backlog; one that has stopped advancing is the state that renders a # frozen-but-plausible token count. _LIVE_CAPTURE_STALL_WARN_TICKS = 30 +# Bounded drain for live_exchange_count(): how many back-to-back polls a +# stage-marker read may spend catching the tail up to the gateway log's end +# before answering "unknown" instead of a stale lower bound. +_LIVE_EXCHANGE_DRAIN_ATTEMPTS = 8 # Agents that cannot make model calls through LiteLLM. ``oracle`` has no model # at all. Gemini is routable through LiteLLM's native Google GenerateContent @@ -214,10 +218,55 @@ def live_usage_tokens(self) -> int | None: return None return self._live_usage_total_tokens + async def live_exchange_count(self) -> int | None: + """Completed provider exchanges at this instant, drained to the log's end. + + The stage-marker seam (rollout-branching RFC §3.5): + ``benchflow.branch_policy.capture_stage`` records this count into + ``stage_snapshots.json`` so a stage-named replay cut + (``bench eval continue --cut-stage``) can replay exactly the prefix + that had completed when the stage closed. Unlike + :meth:`live_usage_tokens` (display-only, allowed to lag), a cut index + must not undercount, so this read *drains* the callback tail to EOF + first — bounded at ``_LIVE_EXCHANGE_DRAIN_ATTEMPTS`` polls — and + returns ``None`` rather than a stale lower bound when the tail still + is not caught up (or the live capture never started): an honest + "unknown" beats a plausible wrong cut. Counts every completed + exchange, failure records included — the same per-record basis + ``trajectory_from_litellm_callback_log`` gives both the live mirror + and the end-of-run import, so the index names the same + ``llm_trajectory.jsonl`` line the replay prefix counts. + """ + trajectory = getattr(self, "_live_trajectory", None) + if trajectory is None: + return None + for _ in range(_LIVE_EXCHANGE_DRAIN_ATTEMPTS): + await self._capture_live_records() + if self._live_capture_lag_bytes == 0: + return len(trajectory.exchanges) + return None + async def _read_callback_chunk(self, offset: int, limit: int) -> _CallbackChunk: raise NotImplementedError async def _capture_live_records(self) -> None: + """Serialized entry to one tail poll. + + Historically only the capture loop (and ``_stop_live_capture``, after + cancelling it) called the poll, so it could assume exclusive access to + the offset/remainder cursor. ``live_exchange_count`` now drains + on demand from the event loop *while the loop task may be mid-poll at + an await point*; two interleaved polls would read the same offset and + append the same records twice. The lock makes every poll atomic with + respect to other polls without changing single-caller behavior. + """ + lock = getattr(self, "_live_capture_lock", None) + if lock is None: + lock = self._live_capture_lock = asyncio.Lock() + async with lock: + await self._capture_live_records_locked() + + async def _capture_live_records_locked(self) -> None: trajectory = getattr(self, "_live_trajectory", None) writer = getattr(self, "_live_writer", None) if trajectory is None or writer is None: diff --git a/src/benchflow/rollout/__init__.py b/src/benchflow/rollout/__init__.py index 8e73b0f96..d54411cee 100644 --- a/src/benchflow/rollout/__init__.py +++ b/src/benchflow/rollout/__init__.py @@ -52,6 +52,7 @@ import shutil import tarfile import tempfile +from collections.abc import Sequence from dataclasses import replace from datetime import datetime from pathlib import Path @@ -78,6 +79,12 @@ from benchflow.acp.types import McpServerSpec from benchflow.agents.credentials import upload_credential from benchflow.agents.registry import AGENTS +from benchflow.branch import StageSnapshot +from benchflow.branch_delta import BranchDelta +from benchflow.branch_stage import STAGE_ENV_READY as STAGE_ENV_READY +from benchflow.branch_stage import STAGE_POST_RESEARCH as STAGE_POST_RESEARCH +from benchflow.branch_stage import STAGE_POST_VERIFY as STAGE_POST_VERIFY +from benchflow.branch_stage import STAGE_PRE_VERIFY as STAGE_PRE_VERIFY from benchflow.contracts import ( AgentProtocolError, AskUserRequest, @@ -197,6 +204,10 @@ from benchflow.rollout.task_runtime import TaskRuntimeResult as TaskRuntimeResult from benchflow.rollout_branch import ChildRunner from benchflow.rollout_branch import branch as _branch_engine +from benchflow.rollout_branch import capture_stage as _capture_stage_engine +from benchflow.rollout_branch import ( + finalize_stage_snapshots as _finalize_stage_snapshots_engine, +) from benchflow.sandbox.metadata import persist_sandbox_info from benchflow.scenes import compile_scenes_to_steps from benchflow.scenes import scene_step_prompt as scene_step_prompt @@ -229,11 +240,12 @@ logger = logging.getLogger(__name__) _SETUP_COMMAND_LOCK_SAFE_RE = re.compile(r"[^A-Za-z0-9._-]+") -# Lifecycle phases from verify() onward. The agent will not run again in this -# rollout once one of these is set, so nothing may rewind ``_phase`` out of -# them — the live dashboard renders the phase as a label and a backwards step -# reads as the run having restarted (see disconnect()). -_TERMINAL_PHASES = frozenset({"verifying", "verified", "cleaned"}) +# Lifecycle phases from verify() onward, plus the branch-first terminal +# ("branched"). The agent will not run again in this rollout once one of these +# is set, so nothing may rewind ``_phase`` out of them — the live dashboard +# renders the phase as a label and a backwards step reads as the run having +# restarted (see disconnect()). +_TERMINAL_PHASES = frozenset({"verifying", "verified", "cleaned", "branched"}) _MCP_TRANSPORT_TO_ACP_TYPE = { @@ -746,6 +758,15 @@ def __init__(self, config: RolloutConfig) -> None: self._tree: RolloutTree = RolloutTree() self._cursor: RolloutNode = self._tree.root + # Stage-boundary snapshot registry (rollout-branching RFC §3.2): + # stage name -> the composed StageSnapshot taken at that boundary, + # plus the tree node it was taken on (the node a branch_at_stage() + # forks, since that node's checkpoint *is* the stage snapshot). Both + # stay empty unless config.snapshot_stages requests a stage or a + # caller calls mark_stage() — the default rollout takes no snapshot. + self._stage_snapshots: dict[str, StageSnapshot] = {} + self._stage_nodes: dict[str, RolloutNode] = {} + # Populated by verify() self._rewards: dict | None = None # Canonical plan-review status/provenance, populated after verify(). @@ -895,7 +916,25 @@ def timing(self) -> dict[str, float]: @property def result(self) -> RolloutResult | None: - if self._phase not in ("verified", "cleaned"): + """The terminal :class:`RolloutResult`, or ``None`` before one exists. + + "branched" is terminal for a branch-first workflow (rollout-branching + RFC §3.4) — the children carried the verification, so the result is + reachable without the linear verify()/cleanup() path. Two branch + specifics: a branch-first rollout that never ran setup() has no run + directory to build artifacts in, so ``result`` stays ``None`` + (the pre-branch contract — never a RuntimeError); and because + branch() restores the parent's linear state (``_rewards`` rolls back + to its pre-branch value), a built branched result surfaces the branch + aggregate V(cursor) recorded by the engine as + ``rewards={"reward": , "source": "branch_aggregate"}`` when no + linear rewards exist (see :meth:`_build_result`). + """ + if self._phase not in ("verified", "cleaned", "branched"): + return None + if self._phase == "branched" and self._rollout_dir is None: + # Branch-first rollout without setup(): no run directory to build + # result artifacts in — stay graceful instead of raising. return None return self._build_result() @@ -1006,8 +1045,17 @@ async def setup(self) -> None: and effective_task_path != cfg.task_path ): effective_skills_dir = task_bundled_skills_dir(effective_task_path) - if effective_skills_dir is not None and not _environment_uses_prebuilt_image( - env_config, cfg.environment_manifest + # A caller-owned sandbox (use_prebuilt_env — a branch child forking a + # stage snapshot) is never built from this Dockerfile, so injecting the + # COPY line there deploys nothing and makes deploy_skills believe the + # pack is already baked in, skipping the runtime upload the child + # actually needs. Leave the staged Dockerfile alone so the upload runs. + if ( + effective_skills_dir is not None + and not self._env_externally_owned + and not _environment_uses_prebuilt_image( + env_config, cfg.environment_manifest + ) ): self._planes.inject_skills_into_dockerfile( effective_task_path, @@ -1143,6 +1191,12 @@ def _capture_and_persist_sandbox() -> None: self._phase = "started" + # env-ready (RFC §3.2): the sandbox is up, the environment plane is + # provisioned and past its readiness gate, and install_agent() has NOT + # run — so a child forked here re-runs agent/skill installation from a + # skill-free world, which is what makes the skills ablation honest. + await self._capture_stage(STAGE_ENV_READY) + # Phase 3: INSTALL AGENT async def install_agent(self) -> None: @@ -1789,7 +1843,9 @@ async def branch( run_child: ChildRunner | None = None, *, require_sandbox_snapshot: bool = False, - ) -> float: + snapshot_layers: frozenset[str] | set[str] = frozenset({"environment"}), + deltas: Sequence[BranchDelta | None] | None = None, + ) -> float | None: """Branch the rollout at the cursor into ``n`` child continuations. Thin entry point — the Branch engine lives in @@ -1809,9 +1865,146 @@ async def branch( without that capability (Modal, Daytona DinD) fail closed with a clear diagnostic rather than running with a half-consistent checkpoint (#384, Branch lifecycle in docs/architecture.md). + + ``snapshot_layers`` selects which checkpoint layers compose the + roll-back point (rollout-branching RFC §3.1). The default + ``{"environment"}`` is the legacy environment-state-only checkpoint; + adding ``"sandbox"`` composes a container-level snapshot with it, and + ``{"sandbox"}`` alone branches a stateless environment on the + container layer only. Missing capability on a requested layer fails + closed before anything is snapshotted. + + ``deltas`` — one :class:`~benchflow.branch_delta.BranchDelta` (or + ``None``) per child, the recorded exactly-one-controlled-change each + child runs under (RFC §3.3). At the cursor only ``injected_prompt`` + executes (the child's user-visible first message, hash-recorded in + provenance); ``skill_mode`` needs the ``env-ready`` boundary and so + runs through :meth:`branch_at_stage`, and every remaining field fails + closed with ``BranchDeltaNotSupported`` before any child runs. """ return await _branch_engine( - self, n, run_child, require_sandbox_snapshot=require_sandbox_snapshot + self, + n, + run_child, + require_sandbox_snapshot=require_sandbox_snapshot, + snapshot_layers=snapshot_layers, + deltas=deltas, + ) + + # Phase 3e: STAGE-BOUNDARY SNAPSHOTS (rollout-branching RFC §3.2) + + @property + def stage_snapshots(self) -> dict[str, StageSnapshot]: + """The lifecycle boundaries this rollout has snapshotted, by stage name. + + Empty unless ``RolloutConfig.snapshot_stages`` requested a stage or a + caller marked one — a default rollout takes no stage snapshot at all. + Each value is the composed :class:`~benchflow.branch.StageSnapshot` + recorded at that boundary, with ``.stage`` set; pass the stage name to + :meth:`branch_at_stage` to fork from it. + """ + return dict(self._stage_snapshots) + + async def _capture_stage(self, stage: str) -> None: + """Snapshot ``stage`` iff this run's config asked for it. + + The lifecycle's own boundaries call this; a stage nobody requested + costs one set membership test and no snapshot call anywhere. A + capability gap (no Environment plane, a sandbox without + ``supports_snapshot``, a stateless environment) fails closed here with + a diagnostic naming the stage — the run does not continue past a + boundary it was told to checkpoint and could not. + """ + if stage not in self._config.snapshot_stages: + return + await _capture_stage_engine( + self, stage, snapshot_layers=self._config.snapshot_layers + ) + + async def mark_stage( + self, + name: str, + *, + snapshot_layers: frozenset[str] | set[str] | None = None, + ) -> StageSnapshot: + """Snapshot the current point and register it under stage ``name``. + + The explicit half of the stage policy: ``post-research`` is a + mid-``execute()`` boundary no lifecycle transition can detect, so the + caller/harness that knows when planning ended marks it. ``name`` must + be a taxonomy stage (``ValueError`` otherwise) but need not appear in + ``RolloutConfig.snapshot_stages`` — marking *is* the request. Quiescing + the agent first is the caller's job, exactly as it is for + :meth:`branch`. + + Marking also records which LLM exchange had completed at that moment + (read from the usage gateway's drained live capture; ``None`` when + unavailable) into the snapshot's meta and ``stage_snapshots.json`` — + the stage→exchange-index data ``bench eval continue --cut-stage`` + resolves (RFC §3.5). + + ``snapshot_layers`` defaults to the run's configured layers. + """ + return await _capture_stage_engine( + self, + name, + snapshot_layers=( + self._config.snapshot_layers + if snapshot_layers is None + else snapshot_layers + ), + ) + + async def branch_at_stage( + self, + stage: str, + n: int, + *, + deltas: Sequence[BranchDelta | None] | None = None, + run_child: ChildRunner | None = None, + snapshot_layers: frozenset[str] | set[str] | None = None, + ) -> float | None: + """Branch from a recorded stage boundary instead of from the cursor. + + The counterfactual entry point (RFC §3.2): the children restore the + world ``stage`` was snapshotted in — not the world the cursor is in — + and fork at the node that stage was captured on, so V aggregates the + value of *that* state. Everything else is the ordinary branch path: + the same delta validation, the same lineage artifacts, and + ``branch_stage`` recorded in each child's provenance as the stage name + rather than the ``cursor:`` fallback. + + A stage that was never captured raises + :class:`~benchflow.branch_stage.BranchStageNotCaptured` listing the + stages that were. ``snapshot_layers`` defaults to the layers the stage + actually captured; passing a different set is rejected rather than + silently re-snapshotting. + + This is also the entry point for the skills ablation (RFC §3.3): at + ``env-ready`` — captured before ``install_agent()``, with the + ``sandbox`` layer so the container rolls back too — a + ``BranchDelta(skill_mode=...)`` child runs as a fresh rollout over the + restored sandbox and re-runs skill deployment under the switched mode. + **The rollout being branched must itself be a ``no-skill`` run.** A + ``with-skill`` parent bakes its pack into the image its ``setup()`` + builds, and the ``env-ready`` snapshot is a commit of that image — so + a ``no-skill`` child would restore the pack, deploy nothing on top of + it, and be reported as ``no-skill`` while running with the parent's + skills. The engine refuses that fork + (:class:`~benchflow.rollout_branch.BranchParentSkillModeConflict`); + run the parent ``no-skill`` and let each arm's delta deploy its own + skills. Every engine-run child of ``env-ready`` is a fresh rollout + whatever its delta, so a ``None`` delta is a genuine control arm — it + re-installs the parent's own recorded mode rather than running in + place. + """ + return await _branch_engine( + self, + n, + run_child, + snapshot_layers=snapshot_layers, + deltas=deltas, + at_stage=stage, ) # Phase 4: VERIFY @@ -1841,6 +2034,12 @@ async def verify(self) -> dict | None: self._env, self._trajectory, self._rollout_paths.agent_dir ) + # pre-verify (RFC §3.2): the agent is quiesced and the workspace is + # still exactly as it left it — this is the last point before + # planes.harden_before_verify (inside _verify_rollout) kills processes + # and restores the verifier-owned workspace. + await self._capture_stage(STAGE_PRE_VERIFY) + ( self._rewards, self._verifier_error, @@ -1858,6 +2057,11 @@ async def verify(self) -> dict | None: self._diagnostics.set(verifier_timeout_diag) self._phase = "verified" + + # post-verify (RFC §3.2): the self-judgment stage of the cascade — + # the verifier has run, so a child forked here re-judges the same + # scored world. + await self._capture_stage(STAGE_POST_VERIFY) return self._rewards async def soft_verify(self) -> tuple[dict | None, str | None, str | None]: @@ -2022,6 +2226,14 @@ async def cleanup(self) -> None: # caller — leave it running so they can reuse it or stop it # themselves. #388. getattr() keeps tests that bypass __init__ # via Rollout.__new__() working. + # + # The stop below destroys every committed ``bf-snap-*`` stage + # image (``compose down --rmi all``), so first make + # ``stage_snapshots.json`` truthful about each ref's lifetime — + # exporting the images to /snapshots/ when + # ``RolloutConfig.keep_snapshots`` asks for it, marking them + # ephemeral otherwise (never raises; PR #1046 second review). + await _finalize_stage_snapshots_engine(self) try: await self._env.stop(delete=True) except Exception as e: @@ -2656,6 +2868,19 @@ def _build_result(self) -> RolloutResult: # to the resolved base prompts when no execute() ran (e.g. setup # failure paths). prompts = self._executed_prompts or self._resolved_prompts + rewards = self._rewards + # getattr() keeps tests that bypass __init__ via Rollout.__new__() + # working (the established pattern in this module). + if rewards is None and getattr(self, "_phase", None) == "branched": + # branch() restores the parent's linear state, so ``_rewards`` + # rolls back to its pre-branch value even though the children + # verified. The honest terminal signal for a branched rollout is + # the branch aggregate V(cursor) the engine recorded on the + # branch-point node (rollout_branch.branch -> aggregate()). + cursor = getattr(self, "_cursor", None) + value = cursor.state.get("value") if cursor is not None else None + if value is not None: + rewards = {"reward": float(value), "source": "branch_aggregate"} return _build_rollout_result( rollout_dir, task_name=self._config.task_path.name, @@ -2671,7 +2896,7 @@ def _build_result(self) -> RolloutResult: trajectory=self._trajectory, partial_trajectory=self._partial_trajectory, trajectory_source=self._trajectory_source, - rewards=self._rewards, + rewards=rewards, started_at=self._require_started_at(), timing=self._timing, scenes=self._config.effective_scenes, diff --git a/src/benchflow/rollout/_config.py b/src/benchflow/rollout/_config.py index d88a47c76..0d5cf8f9c 100644 --- a/src/benchflow/rollout/_config.py +++ b/src/benchflow/rollout/_config.py @@ -24,6 +24,7 @@ normalize_reasoning_effort, normalize_sandbox_user, ) +from benchflow.branch_stage import normalize_stages from benchflow.contracts import BaseUser, RolloutPlanes from benchflow.environment.manifest import EnvironmentManifest from benchflow.loop_strategies import ( @@ -69,8 +70,15 @@ def _task_document_user_runtime( *, prompts: list[str | None] | None, skill_mode: str, -) -> tuple[BaseUser, int | None] | None: - """Compile a supported document-declared user into runtime configuration.""" +) -> Any | None: + """Compile a supported document-declared user into runtime configuration. + + Returns the full :class:`~benchflow.task.prompts.CompiledUserRuntime` + (``None`` when the task declares no adoptable user): beside the user loop + itself it carries the task's stage-capture request — a + ``branch_execution: forked-snapshot`` declaration compiles to the + ``snapshot_stages`` the launch policy should honor (RFC §3.2). + """ if prompts is not None or skill_mode == SKILL_MODE_SELF_GEN: return None @@ -84,7 +92,7 @@ def _task_document_user_runtime( runtime = compile_document_user_runtime(TaskDocument.from_path(document_path)) if runtime.user is None: return None - return runtime.user, runtime.max_rounds + return runtime @dataclass @@ -117,6 +125,29 @@ class RolloutConfig: # C-axis overlay (parsed dict) deep-merged into the task's resolved config # at rollout setup. None => no overlay (default). config_override: dict | None = None + # Stage-boundary snapshot policy (rollout-branching RFC §3.2). The + # lifecycle boundaries named here get a composed checkpoint as the rollout + # passes them, registered on the Rollout for a later branch_at_stage(). + # Empty (default) => today's behavior exactly: no stage is captured and no + # snapshot call is made anywhere in the lifecycle. Unknown names fail + # closed here, at construction, not at the boundary. ``post-research`` is + # not auto-detectable — declaring it says the harness will call + # ``Rollout.mark_stage()`` when planning ends. + snapshot_stages: frozenset[str] = frozenset() + # The checkpoint layers each stage snapshot composes (RFC §3.1), the same + # notion Rollout.branch() takes. Layer names are validated by the branch + # engine's capability gate, the single source of truth for which layers + # exist and which the active planes can actually take. + snapshot_layers: frozenset[str] = frozenset({"environment"}) + # Durable stage-snapshot retention (RFC §3.6, `bench eval run + # --keep-snapshots`): before cleanup destroys the committed ``bf-snap-*`` + # images, export each captured stage's sandbox image (``docker save``) to + # ``/snapshots/.tar`` and record the tar's path, sha256 and + # image id in ``stage_snapshots.json``. False (default) keeps today's + # lifecycle — the images die with the run — and cleanup marks each + # recorded ref ``ephemeral: true`` so the artifact never shows a bare ref + # that no longer resolves. + keep_snapshots: bool = False # Abort the prompt if no tool call arrives for this many seconds. # Catches agents that hung silently while the local process is alive # (e.g. gemini-cli not responding). None disables idle detection. @@ -194,6 +225,8 @@ def __post_init__(self) -> None: if self.skills_dir is not None and not isinstance(self.skills_dir, Path): self.skills_dir = Path(self.skills_dir) self.skill_mode = normalize_skill_mode(self.skill_mode) + self.snapshot_stages = normalize_stages(self.snapshot_stages) + self.snapshot_layers = frozenset(self.snapshot_layers) if self.artifact_skill_mode is not None: self.artifact_skill_mode = normalize_skill_mode(self.artifact_skill_mode) explicit_scenes = bool(self.scenes) @@ -302,9 +335,41 @@ def _resolve_user(self) -> None: ) return if document_user is not None: - self.user, max_rounds = document_user - if max_rounds is not None: - self.max_user_rounds = max_rounds + self.user = document_user.user + if document_user.max_rounds is not None: + self.max_user_rounds = document_user.max_rounds + self._adopt_document_snapshot_request(document_user) + + def _adopt_document_snapshot_request(self, runtime: Any) -> None: + """Honor a task-declared forked-snapshot stage-capture request. + + ``branch_execution: forked-snapshot`` compiles to the stages the task + asks the rollout to snapshot (RFC §3.2 — snapshot policy is opt-in + per run *or per task*); adopting them here is what makes the + declaration real: a plain evaluation of the task leaves + ``stage_snapshots.json`` (with exchange indices) behind, so an + ablation/branch — or ``bench eval continue --cut-stage`` — can fork + the recorded boundaries later. The run-level request wins outright: a + caller that set ``snapshot_stages`` keeps exactly what it asked for. + When the task's request is adopted and the layers are still the + config default, they resolve to the container layer (plus the + environment layer iff an Environment plane is bound): the sandbox + snapshot is what a forked-snapshot branch restores, and the + environment layer without a plane would fail every capture closed. A + sandbox that cannot snapshot still fails the first capture closed at + run time — the capability gate is the branch engine's, not this + adoption's — and ``bench tasks check`` flags that combination before + launch. + """ + stages = getattr(runtime, "snapshot_stages", frozenset()) + if not stages or self.snapshot_stages: + return + self.snapshot_stages = normalize_stages(stages) + if self.snapshot_layers == frozenset({"environment"}): + layers = {"sandbox"} + if self.environment_manifest is not None: + layers.add("environment") + self.snapshot_layers = frozenset(layers) @classmethod def from_legacy( diff --git a/src/benchflow/rollout_branch.py b/src/benchflow/rollout_branch.py index 598998507..644ddea4b 100644 --- a/src/benchflow/rollout_branch.py +++ b/src/benchflow/rollout_branch.py @@ -1,100 +1,115 @@ -"""The Branch -> Rollout engine wiring. +"""The Branch -> Rollout engine wiring — the fork orchestrator. The pure Branch primitives live in :mod:`benchflow.branch` — ``checkpoint``, ``restore``, ``aggregate`` operate on a ``RolloutTree`` node and an ``Environment`` with no I/O beyond the env contract. This module is the -*engine*: it drives those primitives against a live -:class:`~benchflow.rollout.Rollout` — quiescing the agent, running each forked -child as an **isolated sub-rollout**, and restoring the parent's linear state -afterward. - -Why a separate module: ``rollout.py`` is the 5-phase lifecycle; the Branch -path is a distinct, optional capability. Keeping it here holds ``rollout.py`` -under the size threshold and keeps the branch logic independently testable. - -The engine functions are free functions taking a ``Rollout`` as their first -argument — ``Rollout.branch`` is a thin one-line entry point that delegates -here. +*orchestrator*: :func:`branch` drives one fork end to end against a live +:class:`~benchflow.rollout.Rollout` — admit the request, quiesce the agent, +checkpoint, run the children, restore the parent, aggregate, leave lineage — +delegating each concern to its focused module: + +* :mod:`benchflow.branch_policy` — what may be captured and forked, and when: + layer resolution and capability gating, the stage-snapshot capture path + (:func:`~benchflow.branch_policy.capture_stage`), recorded-stage resolution, + and the whole delta-vector gate. Everything fails closed before anything is + quiesced. +* :mod:`benchflow.branch_transaction` — the transactional heart: the composed + checkpoint at the branch point, the scoped + :class:`~benchflow.branch_transaction.LinearState` capture that makes each + child an isolated sub-rollout, and the per-child restore/run/record loop + (:class:`~benchflow.branch_transaction.BranchTransaction`). +* :mod:`benchflow.branch_children` — how a child executes its delta: the + in-place default runner, and (via :mod:`benchflow.branch_skill`) the + fresh-rollout path every ``env-ready`` child takes. +* :mod:`benchflow.branch_report` — lineage and ablation reporting; the + fork's failure-isolated lineage write stays here (:func:`_write_lineage`) + because its writers must resolve through *this* module's namespace — the + patch seam ``tests/test_rollout_branch.py`` pins. + +``Rollout.branch`` / ``Rollout.mark_stage`` / ``Rollout.branch_at_stage`` are +thin entry points that delegate here, and the public names the engine has +always exported (``ChildRunner``, ``CHILD_WALL_CLOCK_KEY``, the gate +exceptions, ``capture_stage``, ``make_default_runner``, the fresh-child API) +are re-exported so existing import paths keep working. Isolation invariant (the architecture's "tree is additive / no-regression"): after :func:`branch` returns, the parent Rollout's linear state — ``_cursor``, ``_trajectory``, ``_rewards``, ``_phase``, ``_n_tool_calls`` (and the session -bookkeeping) — is *exactly* what it was before. A branch child never -re-entrantly mutates the shared instance: it runs against a scoped snapshot of -that state, captured before and restored after each child, and its real -continuation Steps attach to a *pending* branch-child node so the reward and -value land on the right node. +bookkeeping) — is *exactly* what it was before, and so is the result-bearing +state an in-place child writes on its way through (``_timing``, +``_verifier_error``, ``_diagnostics``, the usage counters — see +:data:`~benchflow.branch_transaction._RESULT_STATE_FIELDS`): the parent's own +``result.json`` describes the parent's run, never the last arm's. A branch +child never re-entrantly mutates the shared instance: it runs against a scoped +snapshot of that state, captured before and restored after each child, and its +real continuation Steps attach to a *pending* branch-child node so the reward +and value land on the right node. """ from __future__ import annotations -from collections.abc import Awaitable, Callable -from dataclasses import dataclass -from typing import TYPE_CHECKING +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any +from benchflow.branch import ( # noqa: F401 (StageSnapshot/UnscoredChildError re-exported) + UNSCORED_KEY, + StageSnapshot, + UnscoredChildError, +) from benchflow.branch import aggregate as _aggregate_branch -from benchflow.branch import checkpoint as _checkpoint_branch -from benchflow.branch import restore as _restore_branch -from benchflow.models import TrajectorySource -from benchflow.trajectories.tree import RolloutNode +from benchflow.branch_artifacts import MountedArtifacts, child_mount_dir # noqa: F401 +from benchflow.branch_children import ( # noqa: F401 (re-exports) + EXECUTION_FRESH_ROLLOUT, + FRESH_CHILD_LAYER, + FRESH_CHILD_STAGE, + SKILL_DELTA_LAYER, + SKILL_DELTA_STAGE, + ChildRunner, + make_default_runner, + make_fresh_child_runner, + resolve_child_skill_policy, + resolve_environment_ref_delta, +) +from benchflow.branch_delta import BranchDelta, BranchDeltaNotSupported # noqa: F401 +from benchflow.branch_lineage import ( + write_branch_artifacts, + write_stage_snapshots, +) +from benchflow.branch_policy import ( # noqa: F401 (re-exports) + _EXECUTABLE_DELTA_FIELDS, + _UNSUPPORTED_DELTA_FIELDS, + BranchChildExecutionNotSupported, + BranchParentSkillModeConflict, + capture_stage, + finalize_stage_snapshots, + gate_layers, + recorded_stage_checkpoint, + resolve_layers, + runs_fresh_children, + validate_deltas, + validate_fresh_children, +) +from benchflow.branch_result import ( # noqa: F401 (scope_… re-exported) + scope_child_result_state, + write_in_place_child_result, +) +from benchflow.branch_transaction import ( # noqa: F401 (wall-clock key re-exported) + CHILD_WALL_CLOCK_KEY, + BranchTransaction, + LinearState, + checkpoint_parent, +) if TYPE_CHECKING: from benchflow.rollout import Rollout + from benchflow.trajectories.tree import RolloutNode -# The per-child runner: given the child's branch node, run its continuation and -# return the scalar return. No ``int`` index — a caller that needs per-child -# prompts binds them into a closure (see ``run_child`` in :func:`branch`). -ChildRunner = Callable[[RolloutNode], Awaitable[float]] - - -@dataclass -class _LinearState: - """A scoped snapshot of a Rollout's linear (non-tree) execution state. - - Captured before a branch child runs and restored after — this is what - makes a branch child an *isolated sub-rollout* rather than a re-entrant - mutation of the shared Rollout instance. - """ - - cursor: RolloutNode - trajectory: list[dict] - n_tool_calls: int - phase: str - rewards: dict | None - trajectory_source: TrajectorySource | None - partial_trajectory: bool - session_tool_count: int - session_traj_count: int - executed_prompts: list[str] - - @classmethod - def capture(cls, rollout: Rollout) -> _LinearState: - """Snapshot ``rollout``'s linear state — a shallow copy of the trajectory.""" - return cls( - cursor=rollout._cursor, - trajectory=list(rollout._trajectory), - n_tool_calls=rollout._n_tool_calls, - phase=rollout._phase, - rewards=rollout._rewards, - trajectory_source=rollout._trajectory_source, - partial_trajectory=rollout._partial_trajectory, - session_tool_count=getattr(rollout, "_session_tool_count", 0), - session_traj_count=getattr(rollout, "_session_traj_count", 0), - executed_prompts=list(rollout._executed_prompts), - ) +logger = logging.getLogger(__name__) - def restore_onto(self, rollout: Rollout) -> None: - """Write this snapshot back onto ``rollout`` — undoing a child's mutations.""" - rollout._cursor = self.cursor - rollout._trajectory = list(self.trajectory) - rollout._n_tool_calls = self.n_tool_calls - rollout._phase = self.phase - rollout._rewards = self.rewards - rollout._trajectory_source = self.trajectory_source - rollout._partial_trajectory = self.partial_trajectory - rollout._session_tool_count = self.session_tool_count - rollout._session_traj_count = self.session_traj_count - rollout._executed_prompts = list(self.executed_prompts) +#: Back-compat alias — tests and callers imported the scoped-state class under +#: its original private name. +_LinearState = LinearState async def branch( @@ -103,7 +118,10 @@ async def branch( run_child: ChildRunner | None = None, *, require_sandbox_snapshot: bool = False, -) -> float: + snapshot_layers: frozenset[str] | set[str] | None = None, + deltas: Sequence[BranchDelta | None] | None = None, + at_stage: str | None = None, +) -> float | None: """Branch ``rollout`` at its cursor into ``n`` child continuations. The Branch lifecycle (``docs/architecture.md``, "Lifecycles"): @@ -117,8 +135,13 @@ async def branch( continuation Steps attach directly to its branch node (a *pending* node, no content-free placeholder Step), so the reward lands on the real leaf. 4. ``score / aggregate`` — each child's return is recorded on - ``child.state["reward"]``; their mean is V(parent), recorded on - ``parent.state["value"]`` and returned. + ``child.state["reward"]`` (and its duration on + :data:`CHILD_WALL_CLOCK_KEY`, in memory only); their mean is V(parent), + recorded on ``parent.state["value"]`` and returned. A child that ran + but produced no reward is recorded as *unscored* + (:data:`~benchflow.branch.UNSCORED_KEY`) with no ``reward`` of its own, + and makes the fork's value ``None`` — an unobserved score is never + averaged in as a zero. The branch point is always the current cursor. ``run_child`` is the per-child runner — injected for unit tests; the default @@ -126,97 +149,294 @@ async def branch( runs the continuation, scores it, and disconnects. A caller that needs per-child prompts binds them into the ``run_child`` closure. + ``snapshot_layers`` selects which checkpoint layers compose the roll-back + point (RFC §3.1). The default ``{"environment"}`` keeps the legacy + environment-state-only checkpoint — same behavior, same bare + ``StateSnapshot`` shape on the node. Adding ``"sandbox"`` composes a + container-level snapshot with it (environment first on checkpoint, sandbox + first on restore); ``{"sandbox"}`` alone branches a stateless environment + on the container layer only, never touching environment snapshot/restore. + Missing capability on any requested layer fails closed before anything is + snapshotted. + + ``deltas`` records the exactly-one-controlled-change each child runs + under (RFC §3.3) — one :class:`BranchDelta` (or ``None`` = zero delta) + per child, validated before anything runs (see + :func:`~benchflow.branch_policy.validate_deltas` for the per-field gates + and :mod:`benchflow.branch_children` for how each executable field runs). + + **How a child is executed depends on the boundary, not on the delta.** + ``at_stage="env-ready"`` precedes ``install_agent()``, so every child the + engine runs from there is a *fresh rollout* over the restored sandbox + (``use_prebuilt_env``) that installs the agent for itself — including a + child with no delta at all, which re-installs the parent's own recorded + skill mode. An in-place child there would connect to an agent the restore + deleted, or score a world missing everything installation deploys and + report it as an ordinary single-delta child; that fork therefore requires + the container layer and raises + :class:`~benchflow.branch_policy.BranchChildExecutionNotSupported` without + it. At every other boundary the agent is already installed and children + run in place. + After this returns, ``rollout``'s linear state is exactly what it was - before — the tree gained ``n`` children at the cursor, nothing else moved. + before — the tree gained ``n`` children at the cursor, nothing else + moved; when the rollout knows its run directory, the branch also leaves + lineage artifacts (``tree.json``, + ``branches//children//``), with any + artifact-write failure logged and isolated from the branch result. + + A child that raises anything other than + :class:`~benchflow.branch.UnscoredChildError` ends the fork and propagates + — but the parent's state is restored and the *partial* lineage is written + on the way out, because the caller may report on it: ``bench eval ablate`` + catches that exception and publishes the completed, failed and skipped + arms, and an experiment that reports arms it cannot evidence is worse than + one that reports none. + + ``at_stage`` branches from a recorded stage boundary instead of + checkpointing at the cursor (RFC §3.2): the stage's :class:`StageSnapshot` + becomes the roll-back point every child restores to, the fork happens at + the node that stage was captured on (so the children continue that stage's + world, and V aggregates over this fork only), and the stage name — not the + ``cursor:`` fallback — is the ``branch_stage`` recorded in provenance. """ - if rollout._environment is None: - raise RuntimeError( - "branch() needs the Environment plane — there is no world to " - "snapshot. Pass RolloutConfig(environment_manifest=...)." + stage_snapshot: StageSnapshot | None = None + if at_stage is not None: + stage_snapshot, stage_node, layers = recorded_stage_checkpoint( + rollout, at_stage, snapshot_layers + ) + subject = f"branch_at_stage({at_stage!r})" + else: + subject = "branch" + layers = resolve_layers( + frozenset({"environment"}) if snapshot_layers is None else snapshot_layers, + subject=subject, ) if n < 2: raise ValueError(f"a branch forks into >= 2 children, got n={n}") - # Fail closed when the run requires a three-layer checkpoint but the - # sandbox cannot snapshot the container layer (#384). The Branch - # lifecycle composes container ⊃ environment-state ⊃ agent-session; - # without the container layer, restoring only environment state can - # produce inconsistent state for runs that mutate process/service state - # the Environment manifest does not capture. - if require_sandbox_snapshot: - sandbox = getattr(rollout, "_env", None) - supports = getattr(sandbox, "supports_snapshot", False) - if not supports: - sandbox_name = type(sandbox).__name__ if sandbox else "" - raise RuntimeError( - f"branch(require_sandbox_snapshot=True) cannot run: the active " - f"sandbox {sandbox_name!r} does not implement container-level " - "snapshot/restore. Use a provider whose Sandbox satisfies the " - "checkpoint contract (DockerSandbox or DaytonaSandbox in direct " - "mode), or drop require_sandbox_snapshot if Environment-state " - "checkpoint is sufficient for this run." - ) - - parent = rollout._cursor + # Admit the request — the whole delta vector and the boundary's own gate + # fail closed with nothing quiesced, checkpointed, or run. + validate_deltas( + rollout, deltas, n=n, at_stage=at_stage, layers=layers, run_child=run_child + ) + fresh_children = runs_fresh_children(at_stage, run_child) + if fresh_children: + validate_fresh_children(layers, at_stage=str(at_stage)) + + # Fail closed when a requested layer's plane or capability is missing + # (#384). ``require_sandbox_snapshot`` keeps its original check-only + # semantics; requesting the layer via ``snapshot_layers`` gates the same + # way and then actually composes the sandbox into the checkpoint. A stage + # branch gates too — the planes must still be live to restore what the + # stage captured. + snap_env, snap_sandbox = gate_layers( + rollout, + layers, + subject=subject, + requested=( + "require_sandbox_snapshot=True" + if require_sandbox_snapshot + else f"snapshot_layers={sorted(layers)!r}" + ), + gate_sandbox=require_sandbox_snapshot, + ) + composed = stage_snapshot is not None or layers != frozenset({"environment"}) + parent = stage_node if stage_snapshot is not None else rollout._cursor runner = run_child if run_child is not None else make_default_runner(rollout) + run_dir = getattr(rollout, "_rollout_dir", None) # quiesce — pause the agent before snapshotting so the checkpoint is - # consistent (the Branch lifecycle quiesces first). + # consistent (the Branch lifecycle quiesces first); then checkpoint — the + # roll-back point (a stage branch adopts the snapshot the boundary took). await rollout.disconnect() - - # checkpoint — snapshot the env at the parent; the roll-back point. - await _checkpoint_branch(parent, rollout._environment) + await checkpoint_parent( + rollout, + parent, + stage_snapshot=stage_snapshot, + composed=composed, + snap_env=snap_env, + snap_sandbox=snap_sandbox, + layers=layers, + ) # The parent's linear state, captured once. Each child runs against a fresh # restore of this; the parent is restored to it at the end. - saved = _LinearState.capture(rollout) - - for _ in range(n): - # Attach a *pending* branch-child node — its real continuation Step is - # filled in place by the child's first execute(), so the child's work - # lands on the child node, not a descendant placeholder. - child = rollout._tree.attach(parent) - - # restore the env to the parent's checkpoint, reset the parent's linear - # state, and point the cursor at the pending child for the sub-rollout. - await _restore_branch(parent, rollout._environment) - saved.restore_onto(rollout) - rollout._cursor = child - - ret = await runner(child) - child.state["reward"] = float(ret) + saved = LinearState.capture(rollout) + + # The parent's *on-disk* state needs the same treatment, for the same + # reason. Children share the parent's sandbox, and the sandbox bind-mounts + # the parent rollout's agent/artifacts/verifier directories into the + # container — so without this every child's verifier writes (and clears) + # the parent's own evidence, and the run ends with the last child's + # reward.txt sitting where the parent's belongs. Held aside before the + # first child, handed to each child after it ran, put back at the end. + # :mod:`benchflow.branch_artifacts` fails closed: a custody failure + # preserves the hold directory and raises ``ArtifactCustodyError`` — the + # fork must not report a world whose evidence it lost — except while a + # child's own exception is propagating, where the failure is recorded and + # logged instead so the child's failure stays the diagnosis. + holder = ( + MountedArtifacts.hold(run_dir=run_dir, parent_id=parent.id) + if run_dir is not None + else None + ) + transaction = BranchTransaction( + rollout=rollout, + n=n, + deltas=deltas, + parent=parent, + saved=saved, + runner=runner, + run_child=run_child, + composed=composed, + snap_env=snap_env, + snap_sandbox=snap_sandbox, + fresh_children=fresh_children, + run_dir=run_dir, + holder=holder, + # Resolved through *this* module's namespace at call time — the patch + # seam tests/test_branch_child_result.py pins (a fake fresh-child + # runner, a recording result writer). + fresh_runner_factory=make_fresh_child_runner, + write_child_result=write_in_place_child_result, + ) + try: + try: + await transaction.run_children() + except BaseException: + # The exception on its way out — a child's crash, or a custody + # failure the transaction itself surfaced — is the caller's + # diagnosis. release() must not replace it: it records and logs a + # failure of its own (preserving the hold directory) instead of + # raising over the real one. + if holder is not None: + holder.release(raising=False) + raise + # Success path: a release failure here is the most important thing + # that happened — it raises ArtifactCustodyError (the hold directory + # is preserved), which the except below treats like any other fork + # failure: restore the parent, persist the partial lineage, propagate. + if holder is not None: + holder.release() + except BaseException: + # A child raised something other than UnscoredChildError (an agent + # connection failure, a verifier crash) and the fork ends here. The + # caller may well survive it — ``run_ablation`` catches exactly this + # and reports the completed, failed and skipped arms — so the evidence + # for the arms that *did* run has to be on disk before the exception + # leaves: without this, a partially completed experiment reported arms + # whose tree.json and per-child provenance were never written. + # + # Same two steps as the success path, minus the aggregate: a fork that + # did not finish has no V, and the failing child carries neither a + # reward nor an unscored reason — which is what marks the tree partial. + # The linear-state restore comes first so the parent's own state (and + # its stage registry, re-serialized below) is the parent's again + # before anything reads it. + # + # Nothing on this path may raise: the exception on its way out is the + # caller's diagnosis, and neither best-effort step is worth replacing + # it with. (On the success path a failed restore is loud, as it should + # be — there is no more important failure to preserve there.) + try: + saved.restore_onto(rollout) + except Exception: + logger.warning( + "branch could not restore the parent's linear state after a " + "child failed — the parent's own result may carry the child's", + exc_info=True, + ) + _write_lineage( + rollout, run_dir=run_dir, parent=parent, children=transaction.children + ) + raise # restore the parent's linear state — the tree grew, nothing else moved. saved.restore_onto(rollout) - - # aggregate — per-child return -> V(parent). - value = _aggregate_branch(parent) - parent.state["value"] = value + value = _aggregate_fork(parent, transaction.children) rollout._phase = "branched" + _write_lineage( + rollout, run_dir=run_dir, parent=parent, children=transaction.children + ) return value -def make_default_runner(rollout: Rollout) -> ChildRunner: - """Build the default per-child runner bound to ``rollout``. +def _aggregate_fork( + parent: RolloutNode, children: Sequence[RolloutNode] +) -> float | None: + """Per-child return -> V(parent), over *this* fork's children only. + + A stage branch point can already carry the linear continuation, whose node + has no reward to average. An unscored child makes V *undefined*: averaging + it in as a zero would report a number computed from a score that was never + observed, so the value is left unrecorded and returned as ``None``. - The default runner re-runs the child from the parent's env checkpoint with - a *fresh agent session* — agent-session snapshot is the unsolved hard part - (``docs/architecture.md``, "The hard part"), so the agent restarts per - child. Each child connects a fresh agent and disconnects it at the end, so - no two children's agents overlap (the next child connects only after the - previous one disconnected). ``verify()`` returning ``None`` or an empty - dict falls back to a ``0.0`` return. + "Unrecorded" has to mean *removed*, not merely "not written": a node can + be branched more than once, and skipping the write would leave the + previous fork's V sitting on the node — ``branch()`` returning None while + ``tree.json`` publishes a number for a fork that has none. """ + unscored_children = [child for child in children if UNSCORED_KEY in child.state] + if unscored_children: + parent.state.pop("value", None) + logger.error( + "branch at %s left %d of %d children unscored — V(parent) is " + "undefined and no value is recorded", + parent.id, + len(unscored_children), + len(children), + ) + return None + value = _aggregate_branch(parent, over=children) + parent.state["value"] = value + return value + - async def _runner(child: RolloutNode) -> float: - await rollout.connect() - # Fill the pending branch-child node in place — the continuation Step - # lands on `child` itself, no content-free placeholder. - await rollout.execute(node=child) - rewards = await rollout.verify() - await rollout.disconnect() - if not rewards: - return 0.0 - return float(rewards.get("reward", 0.0)) - - return _runner +def _write_lineage( + rollout: Rollout, + *, + run_dir: Any, + parent: RolloutNode, + children: Sequence[RolloutNode], +) -> None: + """Write the fork's lineage artifacts (RFC §3.4), failure-isolated. + + Written only when the rollout knows its run directory, and never allowed + to propagate: an artifact-write error is logged and neither corrupts the + branch result nor — on the partial path, where this is called from an + ``except`` block — masks the child failure that is on its way out. That + isolation is the reason this is a plain ``except Exception`` rather than a + chained re-raise: the exception being carried is the one the caller has to + see, and it is the more important of the two. + + Called on both paths, so a fork that died mid-way leaves the same shape of + evidence as one that finished: ``tree.json`` for the nodes that exist and + per-child ``provenance.json`` / ``reward.json`` for the children that were + attached, plus the parent's own stage registry (a child that ran through a + stage boundary of its own already overwrote the file with its version). + + Kept in this module (not :mod:`benchflow.branch_report`) so the writers + resolve through *this* namespace — ``tests/test_rollout_branch.py`` + patches ``benchflow.rollout_branch.write_branch_artifacts`` to prove the + isolation. + """ + if run_dir is None: + return + try: + write_branch_artifacts( + run_dir=run_dir, + tree=rollout._tree, + parent=parent, + children=children, + ) + stages = getattr(rollout, "_stage_snapshots", {}) + if stages: + write_stage_snapshots(run_dir=run_dir, snapshots=stages) + except Exception: + logger.warning( + "branch lineage artifact write under %s failed — the branch " + "result is unaffected", + run_dir, + exc_info=True, + ) diff --git a/src/benchflow/sandbox/docker.py b/src/benchflow/sandbox/docker.py index 4ae9f362d..35617ecca 100644 --- a/src/benchflow/sandbox/docker.py +++ b/src/benchflow/sandbox/docker.py @@ -85,6 +85,66 @@ def _is_compose_up_network_race_error(message: str) -> bool: return is_compose_up_network_race_error(message) +def _replayed_run_args(container: dict[str, Any], *, default_network: str) -> list[str]: + """``docker run`` flags that re-create ``container``'s *host* configuration. + + Input is one object from ``docker inspect``. The split is what ``docker + commit`` can and cannot carry into a snapshot image: + + * **Not replayed, because commit already carries it** — ``Env``, + ``WorkingDir``, ``User``, ``Entrypoint``, ``Labels`` and the rest of + ``Config`` are recorded *in the committed image*, so a container created + from it inherits them. Replaying them would be redundant and would let + the two sources disagree. + * **Replayed, because commit cannot carry it** — the host configuration: + bind mounts and volumes (the reason this function exists: a restored + container without them writes into itself, and everything the host + expected to see, above all the verifier's reward file, is lost), + ``none`` networking (a task that opted out of the network must not get + one back through a restore), and the CPU/memory limits. + + Deliberately *not* replayed, and therefore not equivalent: published + ports, capabilities / ``privileged`` / ``security-opt`` / ``sysctls``, + devices, ulimits, extra hosts, restart policy, and membership of more than + one network (``docker run`` attaches one at creation; the compose project + network is used unless the container had none). + """ + args: list[str] = [] + host_config = container.get("HostConfig") or {} + + if str(host_config.get("NetworkMode") or "") == "none": + args += ["--network", "none"] + else: + args += ["--network", default_network] + + for mount in container.get("Mounts") or []: + destination = mount.get("Destination") + if not destination: + continue + kind = mount.get("Type") + if kind == "tmpfs": + args += ["--tmpfs", str(destination)] + continue + # A volume's source is its name; a bind's is the host path. + source = mount.get("Name") if kind == "volume" else mount.get("Source") + if not source: + continue + spec = f"type={kind},src={source},dst={destination}" + if mount.get("RW") is False: + spec += ",readonly" + args += ["--mount", spec] + + # Resource limits: a restored container that ignored the task's caps could + # starve the host the rest of the run shares. + nano_cpus = host_config.get("NanoCpus") or 0 + if nano_cpus: + args += ["--cpus", f"{nano_cpus / 1_000_000_000:g}"] + memory = host_config.get("Memory") or 0 + if memory: + args += ["--memory", str(memory)] + return args + + class DockerSandboxEnvVars(BaseModel): main_image_name: str context_dir: str @@ -265,6 +325,34 @@ async def _probe_verifier_log_mount(self) -> None: "container; verifier outputs will be copied back explicitly." ) + async def has_host_mount(self, *, host_dir: Path, container_dir: str) -> bool: + """Whether the *live* ``main`` container really mounts ``host_dir``. + + :attr:`is_mounted` is a static declaration about how this backend + starts containers; it stays ``True`` for the lifetime of the object. + The container underneath can be replaced — :meth:`restore` re-creates + it — so a caller about to skip a download because "it is mounted + anyway" asks the daemon instead of the declaration. Any failure + answers ``False``: a redundant download is cheap, a skipped one loses + the file. + """ + container_id = await self._main_container_id() + if not container_id: + return False + try: + container = await self._inspect_container(container_id) + except RuntimeError as exc: + self.logger.warning("Could not inspect %s: %s", container_id, exc) + return False + expected = Path(host_dir).resolve() + for mount in container.get("Mounts") or []: + if str(mount.get("Destination") or "") != str(container_dir): + continue + source = mount.get("Source") + if source and Path(source).resolve() == expected: + return True + return False + @property def _dockerfile_path(self) -> Path: return self.environment_dir / "Dockerfile" @@ -751,12 +839,29 @@ async def snapshot(self, name: str | None = None) -> SandboxImage: async def restore(self, image: SandboxImage) -> None: """Restore the ``main`` container from a previously committed image. - Stops and removes the current ``main`` container, then ``docker - run``s a replacement from ``image.ref``. Sibling compose services - are untouched — they keep running as before, matching the - documented container-only scope of the Sandbox layer. + Inspects the live ``main`` container, stops and removes it, then + ``docker run``s a replacement from ``image.ref`` **carrying the host + configuration the inspection recorded** — above all the bind mounts. + Sibling compose services are untouched — they keep running as before, + matching the documented container-only scope of the Sandbox layer. + + The mounts are the whole point: compose bind-mounts the rollout's + ``verifier`` / ``agent`` / ``artifacts`` directories into the + container, and a replacement created without them looks healthy while + silently dropping every file the sandbox writes to those paths — a + branch child's verifier wrote its ``reward.txt`` into a container-local + directory nobody read, and the run reported a fabricated 0.0. + + Fails closed: if the live container cannot be resolved or inspected, + restore raises :class:`~benchflow.sandbox.protocol.SandboxRestoreHostConfigUnavailable` + rather than creating a container whose host config is unknown. Both + halves matter — there is no host config to read once ``main`` is gone, + and a mountless replacement is the exact regression above. """ - from benchflow.sandbox.protocol import SandboxSnapshotNotSupported + from benchflow.sandbox.protocol import ( + SandboxRestoreHostConfigUnavailable, + SandboxSnapshotNotSupported, + ) if image.provider != "docker": raise SandboxSnapshotNotSupported( @@ -765,12 +870,36 @@ async def restore(self, image: SandboxImage) -> None: "across providers." ) + project_name = _sanitize_docker_compose_project_name(self.session_id) + default_network = f"{project_name}_default" + container_id = await self._main_container_id() - if container_id: - await self._docker_cli(["stop", container_id]) - await self._docker_cli(["rm", "-f", container_id]) + if not container_id: + raise SandboxRestoreHostConfigUnavailable( + f"DockerSandbox.restore({image.ref!r}) cannot resolve the " + f"'main' container of compose project {project_name!r}: " + "`docker compose ps -q main` named none, so its bind mounts, " + "network and resource limits cannot be read and the " + "replacement cannot be made equivalent to it. Restoring " + "anyway would create a mountless container that looks healthy " + "while every file the sandbox writes to /logs stays inside " + "it. The container was live when snapshot() committed this " + "image; something removed it since." + ) + # Inspected *before* removal — the host config only exists while + # the container does. + try: + inspected = await self._inspect_container(container_id) + except RuntimeError as exc: + raise SandboxRestoreHostConfigUnavailable( + f"DockerSandbox.restore({image.ref!r}) cannot read the host " + f"config of the 'main' container {container_id!r}, so the " + f"replacement cannot be made equivalent to it: {exc}" + ) from exc + replayed = _replayed_run_args(inspected, default_network=default_network) + await self._docker_cli(["stop", container_id]) + await self._docker_cli(["rm", "-f", container_id]) - project_name = _sanitize_docker_compose_project_name(self.session_id) new_name = f"{project_name}-main-restored-{uuid.uuid4().hex[:8]}" run_cmd = [ @@ -778,12 +907,11 @@ async def restore(self, image: SandboxImage) -> None: "--detach", "--name", new_name, - "--network", - f"{project_name}_default", "--label", f"com.docker.compose.project={project_name}", "--label", "com.docker.compose.service=main", + *replayed, image.ref, "sleep", "infinity", @@ -794,7 +922,54 @@ async def restore(self, image: SandboxImage) -> None: f"docker run from snapshot {image.ref!r} failed: " f"{result.stderr or result.stdout}" ) - self.logger.info(f"Snapshot restored: {image.ref} -> {new_name}") + self.logger.info( + "Snapshot restored: %s -> %s (replayed host config: %s)", + image.ref, + new_name, + " ".join(replayed) or "none", + ) + + async def export_image(self, ref: str, target_path: Path | str) -> None: + """``docker save`` the image ``ref`` to ``target_path`` (a tar). + + The durable half of the snapshot lifecycle (rollout-branching RFC + §3.6): committed ``bf-snap-*`` images die with the rollout's + ``compose down --rmi all``, so a caller that wants the branched + stage's world to outlive the run exports it *before* cleanup. + Fails closed — a missing image raises rather than leaving a + zero-length tar behind. + """ + result = await self._docker_cli( + ["save", "-o", str(target_path), ref], check=False + ) + if result.return_code != 0: + # docker save may leave a partial/empty file on failure; a tar + # that does not restore must not look like one that does. + Path(target_path).unlink(missing_ok=True) + raise RuntimeError( + f"docker save {ref!r} failed: {result.stderr or result.stdout}" + ) + self.logger.info("Snapshot exported: %s -> %s", ref, target_path) + + async def _inspect_container(self, container_id: str) -> dict[str, Any]: + """``docker inspect`` one container as a dict — raises if it cannot.""" + result = await self._docker_cli(["inspect", container_id], check=False) + if result.return_code != 0: + raise RuntimeError( + f"docker inspect {container_id!r} failed: " + f"{result.stderr or result.stdout}" + ) + try: + payload = json.loads(result.stdout) + except json.JSONDecodeError as exc: + raise RuntimeError( + f"docker inspect {container_id!r} returned unparseable JSON: {exc}" + ) from exc + if not isinstance(payload, list) or not payload: + raise RuntimeError( + f"docker inspect {container_id!r} returned no container object" + ) + return payload[0] async def _main_container_id(self) -> str | None: """Return the container id of the ``main`` compose service, or None.""" diff --git a/src/benchflow/sandbox/protocol.py b/src/benchflow/sandbox/protocol.py index ad4a044c7..20c934564 100644 --- a/src/benchflow/sandbox/protocol.py +++ b/src/benchflow/sandbox/protocol.py @@ -46,6 +46,24 @@ class SandboxSnapshotNotSupported(NotImplementedError): """ +class SandboxRestoreHostConfigUnavailable(RuntimeError): + """Restore could not resolve the live container's host configuration. + + A container-level restore replaces the running container with one built + from the snapshot, and the replacement is only equivalent if it carries + the host config the original had — above all the bind mounts, without + which everything the sandbox writes to a mounted path stays inside the + container and the host reads an empty verifier directory. That config + exists only while the original container does, so a restore that cannot + read it has no way to reproduce it and must not create the replacement + anyway: raising loses the branch child, guessing loses the evidence. + + Lives in the core protocol module, next to :class:`SandboxStartupError` + and for the same reason: a base install can catch it without pulling in + any provider SDK. + """ + + class SandboxStartupError(RuntimeError): """Raised when sandbox creation fails or times out. diff --git a/src/benchflow/sandbox/providers.py b/src/benchflow/sandbox/providers.py index f26fbcaa3..ecb006bf8 100644 --- a/src/benchflow/sandbox/providers.py +++ b/src/benchflow/sandbox/providers.py @@ -44,6 +44,15 @@ class SandboxProvider: #: Whether the backend can run a task's docker-compose side services. #: ``False`` means a multi-service task must be refused, not run partially. supports_compose: bool = False + #: Whether the backend can take container-level snapshots — the sandbox + #: layer of the rollout-branching RFC's composed checkpoint (§3.1/§4: + #: docker via ``docker commit``, daytona direct via provider snapshots). + #: Declared registry-level so pre-launch gates (a task declaring + #: ``branch_execution: forked-snapshot``) read a fact instead of growing + #: per-backend special cases; the runtime capability gate + #: (``Sandbox.supports_snapshot``) stays the final authority — it also + #: catches the per-strategy holes a name cannot express (Daytona DinD). + supports_container_snapshot: bool = False @property def off_box_model(self) -> bool: @@ -59,6 +68,7 @@ def off_box_model(self) -> bool: extra=None, model_proxy=ModelProxyLocation.HOST, supports_compose=True, + supports_container_snapshot=True, ), SandboxProvider( "daytona", @@ -66,6 +76,9 @@ def off_box_model(self) -> bool: model_proxy=ModelProxyLocation.SANDBOX, # The DinD strategy runs compose inside the sandbox VM. supports_compose=True, + # Direct sandboxes snapshot via provider images; the DinD strategy + # cannot — the runtime gate fails that combination closed. + supports_container_snapshot=True, ), SandboxProvider( "modal", @@ -115,6 +128,12 @@ def off_box_model(self) -> bool: NO_NETWORK_UNSUPPORTED_PROVIDERS: frozenset[str] = frozenset( p.name for p in _PROVIDERS if not p.enforces_no_network ) +#: Providers whose sandboxes can take container-level snapshots (the branch +#: engine's sandbox layer). Registry-level fact for pre-launch gates; the +#: runtime ``Sandbox.supports_snapshot`` gate remains the final authority. +CONTAINER_SNAPSHOT_PROVIDERS: frozenset[str] = frozenset( + p.name for p in _PROVIDERS if p.supports_container_snapshot +) def is_known_provider(name: str) -> bool: diff --git a/src/benchflow/sandbox/workspace_digest.py b/src/benchflow/sandbox/workspace_digest.py new file mode 100644 index 000000000..fca4e041b --- /dev/null +++ b/src/benchflow/sandbox/workspace_digest.py @@ -0,0 +1,124 @@ +"""Deterministic digest of a sandbox workspace directory. + +The reusable form of the state-digest pipeline the branch docker proofs run +(``tests/test_branch_composed_docker.py``): file contents plus a ``stat`` +listing of names and permission bits, all piped through a final +``sha256sum``. One line out, stable for identical trees — the oracle those +tests use to prove snapshot→restore losslessness, and the workspace half of +the replay cut-point accounting the rollout-branching RFC §3.5 promises +("record a content digest of the last replayed request *and the workspace*"). + +The value is echoed behind a marker prefix and extracted by prefix match: +sandbox ``exec`` implementations merge stderr into stdout, so compose +warnings ("Found orphan containers…") would otherwise corrupt the digest. +""" + +from __future__ import annotations + +import shlex +from typing import Any + +#: The agent's conventional working directory inside benchflow sandboxes. +WORKSPACE_DIGEST_ROOT = "/app" + +#: What the digest is computed over — recorded next to every digest so a +#: reader knows two digests are comparable before comparing them. The +#: null-safe/fail-closed wording is deliberate: the original ``find|sort| +#: xargs`` form word-split legal filenames and recorded a *successful wrong* +#: digest (PR #1046 second review, P1-A), so a digest recorded under that +#: basis must never read as comparable to one recorded under this one. +WORKSPACE_DIGEST_BASIS = ( + "sha256 over sorted per-file sha256sums + a sorted stat listing of names " + "and permission bits (null-safe find -print0|sort -z|xargs -0, " + "fail-closed staged pipeline)" +) + +_MARKER = "BFWSDIGEST:" + + +def workspace_digest_command( + root: str = WORKSPACE_DIGEST_ROOT, *, exclude_basename: str | None = None +) -> str: + """The in-sandbox pipeline: one deterministic line summarizing ``root``. + + Null-safe and fail-closed (PR #1046 second review, P1-A): + + * File names travel NUL-terminated end to end + (``find -print0 | sort -z | xargs -0``), so a legal name containing + spaces — or even a newline — stays one argument. The original + newline-separated form word-split ``file with spaces.txt`` into three + arguments; the inner ``sha256sum`` failed on all of them and two + materially different workspaces recorded the *same* digest. + * Every stage writes to its own file under a private temp dir with its + exit status checked by ``&&`` — an inner failure fails the whole + command, so a digest is either correct or absent-with-reason, never + wrong. This is intrinsic failure propagation: the sandboxes run + ``sh -c`` (busybox ash on minimal images), where ``pipefail`` is not + reliably available, and a trailing ``| sha256sum`` would otherwise + succeed over a broken producer (a pipeline's status is its last + command's). ``LC_ALL=C`` pins the sort to byte order so the digest + does not depend on the image's locale. + + For workspaces whose filenames the old pipeline handled correctly, the + byte stream reaching the final ``sha256sum`` is unchanged, so recorded + digest values are stable; trees with pathological names now digest + correctly (a genuinely different value) instead of colliding. + + ``exclude_basename`` drops files whose basename matches the glob from + both listings — the docker proofs exclude a ``.backup``-restored sqlite + DB that is logically, not byte-, identical. + """ + quoted = shlex.quote(root) + skip = f" ! -name {shlex.quote(exclude_basename)}" if exclude_basename else "" + d = '"$__bf_td"' + return ( + f"cd {quoted} && __bf_td=$(mktemp -d) && " + f"find . -type f{skip} -print0 >{d}/f0 && " + f"LC_ALL=C sort -z <{d}/f0 >{d}/f1 && " + f"xargs -0 -r sha256sum <{d}/f1 >{d}/sums && " + f"find .{skip} -print0 >{d}/a0 && " + f"LC_ALL=C sort -z <{d}/a0 >{d}/a1 && " + f"xargs -0 -r stat -c '%n %a' <{d}/a1 >{d}/stats && " + f"cat {d}/sums {d}/stats >{d}/all && " + f"sha256sum <{d}/all && " + f"rm -rf {d}" + ) + + +async def compute_workspace_digest( + sandbox: Any, + *, + root: str = WORKSPACE_DIGEST_ROOT, + timeout_sec: int = 120, +) -> dict[str, Any]: + """Compute the workspace digest of ``root`` inside a live sandbox. + + Returns ``{"digest": "sha256:", "basis": WORKSPACE_DIGEST_BASIS, + "root": root}`` or raises — the caller records the failure reason instead + of the digest; a digest is never fabricated. + """ + command = workspace_digest_command(root) + wrapped = f'__bf_out="$({command})" && echo "{_MARKER}${{__bf_out}}"' + result = await sandbox.exec(wrapped, timeout_sec=timeout_sec) + if result.return_code != 0: + raise RuntimeError( + f"workspace digest command failed (rc={result.return_code}): " + f"{(result.stderr or result.stdout or '').strip()}" + ) + values = [ + line[len(_MARKER) :].strip() + for line in (result.stdout or "").splitlines() + if line.startswith(_MARKER) + ] + if len(values) != 1: + raise RuntimeError( + f"workspace digest expected exactly one {_MARKER} line, got {values!r}" + ) + hex_digest = values[0].split()[0] if values[0] else "" + if len(hex_digest) != 64 or any(c not in "0123456789abcdef" for c in hex_digest): + raise RuntimeError(f"workspace digest output is not a sha256: {values[0]!r}") + return { + "digest": f"sha256:{hex_digest}", + "basis": WORKSPACE_DIGEST_BASIS, + "root": root, + } diff --git a/src/benchflow/snapshot_import.py b/src/benchflow/snapshot_import.py new file mode 100644 index 000000000..6b0159bfa --- /dev/null +++ b/src/benchflow/snapshot_import.py @@ -0,0 +1,217 @@ +"""Load a completed run's exported stage snapshots back into Docker. + +The import half of ``--keep-snapshots`` (rollout-branching RFC §3.6): a run +that captured stage boundaries leaves ``stage_snapshots.json`` behind, and +with the flag each captured ``bf-snap-*`` image was ``docker save``d to +``/snapshots/.tar`` before cleanup destroyed it — with the +tar's path, content sha256 and image id recorded per stage. This module makes +"branch later from a completed run" real: it verifies the recorded sha256, +``docker load``s the tar, and confirms the recorded ref resolves to the +recorded image id — a snapshot is only reported restored when +``docker image inspect`` would agree. + +Fail-closed by design: an entry recorded ``ephemeral: true`` (the image died +with the run), a tar whose digest does not match its record, or a loaded +image whose id differs from the recorded one all raise +:class:`SnapshotImportError` naming exactly what disagreed. CLI surface: +``bench eval import-snapshots ``. +""" + +from __future__ import annotations + +import json +import logging +import subprocess +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from benchflow.branch_policy import _file_sha256 + +logger = logging.getLogger(__name__) + +#: Seam for the two docker CLI calls the import makes (``load`` and +#: ``image inspect``); tests substitute a recorder, the default shells out. +DockerRunner = Callable[[list[str]], "subprocess.CompletedProcess[str]"] + + +class SnapshotImportError(RuntimeError): + """A stage-snapshot import that cannot restore what the run recorded.""" + + +@dataclass(frozen=True) +class ImportedSnapshot: + """One stage snapshot restored into the local Docker image store.""" + + stage: str + #: The recorded (and now once again resolvable) image ref, ``bf-snap-…``. + sandbox_ref: str + #: The id ``docker image inspect`` reports for the loaded ref — equal to + #: the recorded id whenever the run recorded one. + image_id: str + tar_path: Path + + +def _run_docker(args: list[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["docker", *args], capture_output=True, text=True, check=False, timeout=600 + ) + + +def read_stage_snapshots(run_dir: Path | str) -> dict[str, dict[str, Any]]: + """The ``stages`` mapping of ``/stage_snapshots.json``. + + Raises :class:`SnapshotImportError` when the file is absent or does not + carry a ``stages`` mapping — there is nothing recorded to import. + """ + path = Path(run_dir) / "stage_snapshots.json" + if not path.exists(): + raise SnapshotImportError( + f"no stage_snapshots.json under {Path(run_dir)} — this is not a " + "run directory of a rollout that captured stage snapshots" + ) + try: + payload = json.loads(path.read_text()) + except (json.JSONDecodeError, OSError) as exc: + raise SnapshotImportError(f"unreadable {path}: {exc}") from exc + stages = payload.get("stages") if isinstance(payload, dict) else None + if not isinstance(stages, dict): + raise SnapshotImportError(f"{path} carries no 'stages' mapping") + return stages + + +def _resolve_tar_path(run_dir: Path, exported: dict[str, Any]) -> Path: + """The exported tar's location, tolerating a relocated run directory. + + The export records an absolute path on the machine that ran the eval; a + run folder copied elsewhere (scp, artifact download) keeps the tar at + ``/snapshots/``, so that is the fallback when the + recorded path does not exist here. + """ + recorded = Path(str(exported.get("path"))) + if recorded.exists(): + return recorded + relocated = Path(run_dir) / "snapshots" / recorded.name + if relocated.exists(): + return relocated + raise SnapshotImportError( + f"exported snapshot tar not found: neither the recorded path " + f"{recorded} nor {relocated} exists" + ) + + +def _load_one( + run_dir: Path, + stage: str, + entry: dict[str, Any], + *, + run_docker: DockerRunner, +) -> ImportedSnapshot: + """Verify, ``docker load``, and identity-check one exported stage entry.""" + exported = entry.get("exported") + if not isinstance(exported, dict): + raise SnapshotImportError( + f"stage {stage!r} has no exported snapshot — its recorded ref " + f"({entry.get('sandbox_ref')!r}) is ephemeral: the image died " + "with the run's cleanup. Re-run the evaluation with " + "--keep-snapshots to retain an importable tar." + ) + sandbox_ref = entry.get("sandbox_ref") + if not sandbox_ref: + raise SnapshotImportError( + f"stage {stage!r} records no sandbox-layer image ref to restore" + ) + tar_path = _resolve_tar_path(run_dir, exported) + recorded_sha = exported.get("sha256") + if not recorded_sha: + raise SnapshotImportError( + f"stage {stage!r}'s export record carries no sha256 — the tar at " + f"{tar_path} cannot be verified against what the run exported" + ) + actual_sha = _file_sha256(tar_path) + if actual_sha != recorded_sha: + raise SnapshotImportError( + f"stage {stage!r}'s tar {tar_path} does not match its record: " + f"recorded {recorded_sha}, found {actual_sha} — refusing to load " + "a tar that is not the one the run exported" + ) + load = run_docker(["load", "-i", str(tar_path)]) + if load.returncode != 0: + raise SnapshotImportError( + f"docker load of {tar_path} failed: " + f"{(load.stderr or load.stdout or '').strip()}" + ) + inspect = run_docker(["image", "inspect", "--format", "{{.Id}}", str(sandbox_ref)]) + if inspect.returncode != 0: + raise SnapshotImportError( + f"docker load of {tar_path} succeeded but the recorded ref " + f"{sandbox_ref!r} does not resolve: " + f"{(inspect.stderr or inspect.stdout or '').strip()}" + ) + loaded_id = (inspect.stdout or "").strip() + recorded_id = exported.get("image_id") + if recorded_id and loaded_id != recorded_id: + raise SnapshotImportError( + f"stage {stage!r} loaded, but {sandbox_ref!r} now names image " + f"{loaded_id}, not the recorded {recorded_id} — the ref resolves " + "to a different world than the run snapshotted" + ) + logger.info("stage %r snapshot restored: %s (%s)", stage, sandbox_ref, loaded_id) + return ImportedSnapshot( + stage=stage, + sandbox_ref=str(sandbox_ref), + image_id=loaded_id, + tar_path=tar_path, + ) + + +def import_stage_snapshots( + run_dir: Path | str, + *, + stages: Iterable[str] | None = None, + run_docker: DockerRunner | None = None, +) -> list[ImportedSnapshot]: + """Restore a completed run's exported stage snapshots into Docker. + + ``stages=None`` (default) imports every stage whose entry carries an + ``exported`` record; naming stages explicitly fails closed on one that + was not recorded or not exported. After a successful import each returned + :class:`ImportedSnapshot`'s ``sandbox_ref`` resolves locally again — + verified sha256 on the tar, verified image id on the loaded ref — so the + recorded boundary can be branched (``DockerSandbox.restore`` / + ``Rollout.branch_at_stage`` machinery, or a plain ``docker run``). + + Raises :class:`SnapshotImportError`; never partially lies — a stage is + only present in the result when its image verifiably resolves. + """ + # Resolved at call time (not bound as a parameter default) so tests can + # substitute the module-level runner. + docker = run_docker if run_docker is not None else _run_docker + run_path = Path(run_dir) + recorded = read_stage_snapshots(run_path) + if stages is None: + selected = sorted( + stage + for stage, entry in recorded.items() + if isinstance(entry, dict) and isinstance(entry.get("exported"), dict) + ) + if not selected: + raise SnapshotImportError( + f"no exported stage snapshots under {run_path} — every " + f"recorded stage ({sorted(recorded) or 'none'}) is ephemeral: " + "the images died with the run's cleanup. Re-run the " + "evaluation with --keep-snapshots to retain importable tars." + ) + else: + selected = list(stages) + missing = [stage for stage in selected if stage not in recorded] + if missing: + raise SnapshotImportError( + f"stage(s) {missing!r} were not captured by this run — it " + f"recorded {sorted(recorded)!r}" + ) + return [ + _load_one(run_path, stage, recorded[stage], run_docker=docker) + for stage in selected + ] diff --git a/src/benchflow/task/prompts.py b/src/benchflow/task/prompts.py index c4463d38c..de2676fb4 100644 --- a/src/benchflow/task/prompts.py +++ b/src/benchflow/task/prompts.py @@ -7,13 +7,16 @@ from typing import Any, Literal, cast from benchflow._types import Scene, Turn +from benchflow.branch_stage import AUTO_STAGES, normalize_stages from benchflow.contracts.user import BaseUser, DocumentNudgeUser, ModelDocumentNudgeUser from benchflow.task.document import TaskDocument PromptComposition = Literal["append", "replace"] PromptPartKind = Literal["base", "role", "scene", "turn"] TeamHandoffKind = Literal["none", "sequential-shared"] -BranchExecutionKind = Literal["none", "option-kinds-preserved", "unsupported"] +BranchExecutionKind = Literal[ + "none", "option-kinds-preserved", "forked-snapshot", "unsupported" +] _APPEND_DEFAULT_ORDER: tuple[PromptPartKind, ...] = ("base", "role", "scene", "turn") _REPLACE_DEFAULT_ORDER: tuple[PromptPartKind, ...] = ("turn", "scene", "role", "base") @@ -76,11 +79,20 @@ class _TeamHandoffPolicy: @dataclass(frozen=True) class CompiledUserRuntime: - """Concrete user runtime compiled from document-declared user metadata.""" + """Concrete user runtime compiled from document-declared user metadata. + + ``snapshot_stages`` is the stage-capture request a + ``branch_execution: forked-snapshot`` declaration compiles to (rollout + branching RFC §3.2): the launch policy asks the rollout to snapshot those + boundaries so an ablation/branch (or a ``--cut-stage`` continuation) can + fork them later. Empty for every other declaration — and on the + fail-closed path, so an unsupported task never requests capture. + """ contract: UserRuntimeContract user: BaseUser | None max_rounds: int | None + snapshot_stages: frozenset[str] = frozenset() @dataclass(frozen=True) @@ -241,19 +253,26 @@ def compile_document_user_runtime(document: TaskDocument) -> CompiledUserRuntime branch_execution_supported = False unsupported.append("benchflow.nudges.branch_execution must be a string") elif branch_execution is not None: - if branch_execution != "option-kinds-preserved": + if branch_execution not in _BRANCH_EXECUTION_VALUES: branch_execution_supported = False unsupported.append( "benchflow.nudges.branch_execution supports only " - "option-kinds-preserved; forked branch execution is not implemented" + "option-kinds-preserved and forked-snapshot" ) elif not branchable: branch_execution_supported = False unsupported.append( "benchflow.nudges.branch_execution requires branchable: true" ) + snapshot_stages, branch_stage_issues = _compile_branch_stages( + nudges, branch_execution=branch_execution + ) + if branch_stage_issues: + branch_execution_supported = False + unsupported.extend(branch_stage_issues) branch_execution_kind = _branch_execution_kind( branchable, + declared=branch_execution, supported=branch_execution_supported, ) confirmation_policy = _confirmation_policy_tier(nudges.get("confirmation_policy")) @@ -384,6 +403,7 @@ def compile_document_user_runtime(document: TaskDocument) -> CompiledUserRuntime ), user=user, max_rounds=max_rounds, + snapshot_stages=snapshot_stages, ) @@ -631,13 +651,62 @@ def _handoff_kind(policy: _TeamHandoffPolicy | None) -> TeamHandoffKind: return "sequential-shared" +#: The declarable ``benchflow.nudges.branch_execution`` values. Both are real: +#: ``option-kinds-preserved`` is the ACP ask_user slice (option IDs and kinds +#: survive the bridge), ``forked-snapshot`` is the Environment/sandbox +#: snapshot branch engine (rollout-branching RFC) — the task requests stage +#: capture so a branch/ablation can fork the recorded boundaries. +_BRANCH_EXECUTION_VALUES = frozenset({"option-kinds-preserved", "forked-snapshot"}) + + +def _compile_branch_stages( + nudges: dict[str, Any], *, branch_execution: str | None +) -> tuple[frozenset[str], list[str]]: + """The stage-capture request a forked-snapshot declaration compiles to. + + ``benchflow.nudges.branch_stages`` optionally narrows (or, for + ``post-research``, extends — declaring it says the harness will + ``mark_stage()`` it) the boundaries the task asks the rollout to + snapshot; it is validated against the branch-stage taxonomy and rejected + without ``branch_execution: forked-snapshot``, the declaration that gives + it meaning. With no explicit list, forked-snapshot requests the + auto-capturable boundaries (``env-ready``/``pre-verify``/``post-verify``) + — the failure-cascade stages the lifecycle can capture on its own. + """ + raw = nudges.get("branch_stages") + if raw is None: + if branch_execution == "forked-snapshot": + return frozenset(AUTO_STAGES), [] + return frozenset(), [] + if branch_execution != "forked-snapshot": + return frozenset(), [ + "benchflow.nudges.branch_stages requires branch_execution: forked-snapshot" + ] + if ( + not isinstance(raw, list) + or not raw + or not all(isinstance(stage, str) for stage in raw) + ): + return frozenset(), [ + "benchflow.nudges.branch_stages must be a non-empty list of " + "branch stage names" + ] + try: + return normalize_stages(raw), [] + except ValueError as exc: + return frozenset(), [f"benchflow.nudges.branch_stages: {exc}"] + + def _branch_execution_kind( branchable: bool, *, + declared: str | None = None, supported: bool = True, ) -> BranchExecutionKind: if not supported: return "unsupported" + if declared == "forked-snapshot": + return "forked-snapshot" if branchable: return "option-kinds-preserved" return "none" diff --git a/src/benchflow/task/runtime_capabilities.py b/src/benchflow/task/runtime_capabilities.py index e624d5441..f0054b695 100644 --- a/src/benchflow/task/runtime_capabilities.py +++ b/src/benchflow/task/runtime_capabilities.py @@ -17,6 +17,7 @@ from benchflow.rewards.rubric_config import criteria_aggregate_policy_from_rubric from benchflow.sandbox._compose import compose_definition_path from benchflow.sandbox.providers import ( + CONTAINER_SNAPSHOT_PROVIDERS, NO_NETWORK_UNSUPPORTED_PROVIDERS, SANDBOX_PROVIDER_SET, SINGLE_CONTAINER_PROVIDERS, @@ -443,6 +444,7 @@ def _append_document_user_issues( ) -> None: runtime = user_runtime or compile_document_user_runtime(document) if runtime.contract.status != "unsupported": + _append_forked_snapshot_issues(unsupported, runtime=runtime, sandbox=sandbox) return reason = runtime.contract.reason or "document user runtime is unsupported" if document.user: @@ -453,6 +455,39 @@ def _append_document_user_issues( _issue(unsupported, path="benchflow.nudges", reason=reason, sandbox=sandbox) +def _append_forked_snapshot_issues( + unsupported: list[UnsupportedTaskFeature], + *, + runtime: CompiledUserRuntime, + sandbox: str, +) -> None: + """Fail a forked-snapshot declaration closed on a snapshot-less backend. + + ``branch_execution: forked-snapshot`` compiles to a stage-capture request + the rollout honors with container-level snapshots (the branch engine's + sandbox layer); a backend whose sandboxes cannot snapshot would fail the + run at the first captured boundary, so the gate says so before launch. + Registry-level fact only — the runtime ``Sandbox.supports_snapshot`` gate + stays the final authority (it also catches per-strategy holes such as + Daytona DinD). + """ + if runtime.contract.branch_execution != "forked-snapshot": + return + if sandbox in CONTAINER_SNAPSHOT_PROVIDERS: + return + _issue( + unsupported, + path="benchflow.nudges.branch_execution", + reason=( + "forked-snapshot needs a container-snapshot-capable sandbox " + f"({', '.join(sorted(CONTAINER_SNAPSHOT_PROVIDERS))}); " + f"{sandbox} sandboxes cannot snapshot, so the requested stage " + "captures would fail closed at the first boundary" + ), + sandbox=sandbox, + ) + + def _append_prompt_policy_issues( unsupported: list[UnsupportedTaskFeature], *, diff --git a/src/benchflow/task/verifier_core.py b/src/benchflow/task/verifier_core.py index 70a786b26..de5afb73a 100644 --- a/src/benchflow/task/verifier_core.py +++ b/src/benchflow/task/verifier_core.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import json import logging import math @@ -291,6 +292,53 @@ async def verify(self) -> VerifierResult: return await self._verify_llm_judge() return await self._verify_test_script() + async def _verifier_outputs_are_mounted(self, service: str) -> bool: + """Whether the verifier's output dir is already visible on the host. + + ``True`` skips both the pre-run clear and the post-run download: the + container writes straight into the rollout directory. Getting it wrong + in that direction is silent and expensive — the verifier's + ``reward.txt`` stays inside the container and the rollout is scored + from a file that was never written — so the answer is taken from the + *live* container when the backend can supply it + (``Sandbox.has_host_mount``), not from the static ``is_mounted`` + declaration, which a ``restore()`` between start and verify can + invalidate. Anything unknown answers ``False``: a redundant download + costs a round trip, a skipped one costs the score. + """ + # Only the agent's ``main`` container ever has the rollout dir + # bind-mounted; a target service (#248) never does. + if service != "main" or not getattr(self._sandbox, "is_mounted", False): + return False + live_check = getattr(self._sandbox, "has_host_mount", None) + if not asyncio.iscoroutinefunction(live_check): + # Backend supplies no live check (or a stub that is not one): the + # static declaration is all there is. + return True + try: + mounted = bool( + await live_check( + host_dir=self._rollout_paths.verifier_dir, + container_dir=str(SandboxPaths.verifier_dir), + ) + ) + except Exception as e: + self._logger.warning( + "Could not confirm the verifier output mount (%s); downloading " + "the verifier directory instead", + e, + ) + return False + if not mounted: + self._logger.warning( + "The sandbox declares host-mounted verifier outputs but the " + "live container has no mount of %s at %s — downloading the " + "verifier directory instead", + self._rollout_paths.verifier_dir, + SandboxPaths.verifier_dir, + ) + return mounted + def _selected_verifier_document(self) -> VerifierDocument | None: paths = getattr(self._task, "paths", None) verifier_dir = getattr(paths, "tests_dir", None) @@ -358,9 +406,7 @@ async def _verify_test_script( DB modifications — instead of only the agent workspace (#248). """ service = self._task.config.verifier.service - verifier_outputs_are_mounted = service == "main" and getattr( - self._sandbox, "is_mounted", False - ) + verifier_outputs_are_mounted = await self._verifier_outputs_are_mounted(service) if not verifier_outputs_are_mounted: try: await clear_verifier_output_dir( @@ -745,9 +791,7 @@ async def _verify_reward_kit( ) -> VerifierResult: """Run a verifier-scoped Reward Kit runner inside the sandbox.""" service = self._task.config.verifier.service - verifier_outputs_are_mounted = service == "main" and getattr( - self._sandbox, "is_mounted", False - ) + verifier_outputs_are_mounted = await self._verifier_outputs_are_mounted(service) if not verifier_outputs_are_mounted: try: await clear_verifier_output_dir( diff --git a/tests/continue_run/_helpers.py b/tests/continue_run/_helpers.py index 4682b3936..fb7064e99 100644 --- a/tests/continue_run/_helpers.py +++ b/tests/continue_run/_helpers.py @@ -61,6 +61,7 @@ def write_run_folder( task_name: str = "demo-task", prompts: list[str] | None = None, timeout_sec: int = 3600, + stage_snapshots: dict[str, dict[str, Any]] | None = None, ) -> Path: """Materialize a synthetic run folder benchflow continue can load.""" root.mkdir(parents=True, exist_ok=True) @@ -97,4 +98,12 @@ def write_run_folder( ) + "\n" ) + if stage_snapshots is not None: + # The recorded stage registry a stage-named cut (--cut-stage) resolves + # against — the shape benchflow.branch_lineage.write_stage_snapshots + # leaves in a real run folder. + (root / "stage_snapshots.json").write_text( + json.dumps({"schema_version": 1, "stages": stage_snapshots}, indent=2) + + "\n" + ) return root diff --git a/tests/continue_run/test_orchestrator.py b/tests/continue_run/test_orchestrator.py index de67ade4d..5c6a49bd2 100644 --- a/tests/continue_run/test_orchestrator.py +++ b/tests/continue_run/test_orchestrator.py @@ -139,11 +139,10 @@ def test_live_forwarder_build_kwargs_resolves_route_offline(): assert kwargs["temperature"] == 0.5 -def test_stitched_trajectory_recorded_prefix_plus_live_suffix(tmp_path): - original = tmp_path / "orig.jsonl" - original.write_text('{"a": 1}\n{"b": 2}\n') +def test_stitched_trajectory_recorded_prefix_plus_live_suffix(): + recorded_lines = ['{"a": 1}', '{"b": 2}'] live = [exchange(completion(content="LIVE"))] - lines = stitched_trajectory_lines(original, live) + lines = stitched_trajectory_lines(recorded_lines, live) assert len(lines) == 3 assert json.loads(lines[0]) == {"a": 1} last = json.loads(lines[2]) @@ -151,11 +150,9 @@ def test_stitched_trajectory_recorded_prefix_plus_live_suffix(tmp_path): def test_write_stitched_trajectory_creates_file(tmp_path): - original = tmp_path / "orig.jsonl" - original.write_text('{"a": 1}\n') rollout_dir = tmp_path / "rollout" out = write_stitched_trajectory( - rollout_dir, original, [exchange(completion(content="L"))] + rollout_dir, ['{"a": 1}'], [exchange(completion(content="L"))] ) assert out == rollout_dir / "trajectory" / "llm_trajectory.jsonl" assert len(out.read_text().strip().splitlines()) == 2 diff --git a/tests/continue_run/test_replay_cut_point.py b/tests/continue_run/test_replay_cut_point.py new file mode 100644 index 000000000..459807700 --- /dev/null +++ b/tests/continue_run/test_replay_cut_point.py @@ -0,0 +1,1171 @@ +"""Replay cut-point tests. + +Guards the replay cut-point ("feat(continue): replay cut-point — replay first +K exchanges, then go live"; docs/rollout-branching-rfc.md WS-3; +FrontierPhysics#73). PR number to be added on submission. + +``max_exchanges`` replays at most the first K recorded exchanges, then switches +the proxy to live passthrough exactly as if the recording had ended there. The +router exposes cut-point accounting (``n_replayed_exchanges`` + +``served_request_digest``/``recorded_request_digest`` + the workspace digest +taken as the run crosses the cut), a cut continue-run's ``source_provenance`` +gains a ``cut_point`` block (configured at build time, reconciled with the +served counts post-run in host proxy mode), and ``stage_tags`` + ``cut_stage`` +name a cut by recorded stage — ``stage_tags[stage]`` is the 1-based count of +exchanges that had completed when the stage closed. Unit tests against the +existing continue_run test doubles — no Docker, no API keys. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar + +import pytest +from typer.testing import CliRunner + +from benchflow.cli.main import app +from benchflow.continue_run.orchestrator import ( + build_agent_env, + build_rollout_config, + continue_run, + cut_point_provenance, + resolve_cut_point, + served_cut_point, + stage_tags_from_run, + stitched_trajectory_lines, + summarize_llm_trajectory_usage, + update_continued_metadata, + write_continuation_artifacts, + write_stitched_trajectory, +) +from benchflow.continue_run.replay_proxy import ( + REQUEST_DIGEST_BASIS, + ReplayCutPointError, + ReplayRouter, + comparable_request_digest, + validate_max_exchanges, +) +from benchflow.continue_run.run_folder import load_run_folder + +from ._helpers import completion, exchange, write_run_folder + +runner = CliRunner() + + +def _recorded(n: int = 3): + """n exchanges with distinct contents and distinct request bodies.""" + return [ + exchange(completion(content=f"turn-{i}"), n_request_messages=i + 1) + for i in range(n) + ] + + +def _request(i: int) -> dict: + """The agent request matching ``_recorded``'s i-th turn (no divergence).""" + return {"messages": [{"role": "user"}] * (i + 1)} + + +# ── ReplayRouter: max_exchanges cut ─────────────────────────────────────── + + +def test_cut_serves_exactly_k_then_live(): + recorded = _recorded(3) + live_requests = [] + + def forwarder(req): + live_requests.append(req) + return completion(content="LIVE") + + router = ReplayRouter(recorded, live_forwarder=forwarder, max_exchanges=2) + + r1 = router.next_response(_request(0)) + assert r1.source == "replay" + assert r1.body["choices"][0]["message"]["content"] == "turn-0" + assert router.exhausted is False + + r2 = router.next_response(_request(1)) + assert r2.source == "replay" + assert r2.body["choices"][0]["message"]["content"] == "turn-1" + # the cut behaves like the natural end of the recording + assert router.exhausted is True + + r3 = router.next_response(_request(2)) + assert r3.source == "live" + assert r3.body["choices"][0]["message"]["content"] == "LIVE" + # exactly K recorded turns served; the third recording is never replayed + assert router.n_replayed_exchanges == 2 + assert len(live_requests) == 1 + assert len(router.live_exchanges) == 1 + assert router.divergences == 0 + + +def test_cut_without_forwarder_errors_like_natural_exhaustion(): + router = ReplayRouter(_recorded(3), live_forwarder=None, max_exchanges=1) + assert router.next_response(_request(0)).source == "replay" + result = router.next_response(_request(1)) + assert result.source == "error" + assert result.status == 503 + assert result.body["error"]["type"] == "replay_exhausted" + + +def test_cut_at_n_equals_full_prefix_behavior(): + recorded = _recorded(2) + cut = ReplayRouter( + recorded, live_forwarder=lambda req: completion(content="L"), max_exchanges=2 + ) + full = ReplayRouter(recorded, live_forwarder=lambda req: completion(content="L")) + for router in (cut, full): + sources = [router.next_response(_request(i)).source for i in range(3)] + assert sources == ["replay", "replay", "live"] + assert cut.n_replayed_exchanges == full.n_replayed_exchanges == 2 + assert cut.recorded_request_digest == full.recorded_request_digest + assert cut.served_request_digest == full.served_request_digest + + +@pytest.mark.parametrize("bad", [0, -1, 4]) +def test_invalid_cut_fails_closed_at_configuration_time(bad): + with pytest.raises(ReplayCutPointError, match=r"max_exchanges must be in 1\.\.3"): + ReplayRouter(_recorded(3), max_exchanges=bad) + + +def test_validate_max_exchanges_none_means_full_prefix(): + assert validate_max_exchanges(None, 5) == 5 + assert validate_max_exchanges(3, 5) == 3 + + +# ── cut-point accounting ────────────────────────────────────────────────── + + +def test_cut_point_digests_match_independent_sha256(): + """The served digest is of the ACTUAL incoming request at the cut. + + Guards "fix(continue): replay divergence detection compares actual content + and records workspace digests": the block's digest used to hash only the + *recorded* request, so a diverged replay hashed to the recorded value and + looked faithful. Both sides are now digested and named separately; here + the incoming request equals the recorded one, so the two digests agree — + and match an independently computed sha256 of the comparable projection. + """ + recorded = _recorded(3) + router = ReplayRouter( + recorded, live_forwarder=lambda req: completion(content="L"), max_exchanges=2 + ) + assert router.served_request_digest is None # nothing replayed yet + assert router.recorded_request_digest is None + for i in range(3): + router.next_response(_request(i)) + + expected = ( + "sha256:" + + hashlib.sha256( + json.dumps( + {"messages": recorded[1].request.body["messages"]}, + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + ) + assert router.n_replayed_exchanges == 2 + assert router.recorded_request_digest == expected + assert router.served_request_digest == expected + assert comparable_request_digest(recorded[1].request.body) == expected + + +def test_served_digest_differs_from_recorded_when_replay_diverged(): + """The half the recorded-only digest could not express: the agent's actual + request at the cut is what ``served_request_digest`` hashes.""" + recorded = _recorded(1) + router = ReplayRouter(recorded) + router.next_response({"messages": [{"role": "user", "content": "CHANGED"}]}) + assert router.served_request_digest != router.recorded_request_digest + assert router.recorded_request_digest == comparable_request_digest( + recorded[0].request.body + ) + + +def test_cut_point_digests_deterministic_across_runs(): + recorded = _recorded(3) + digests = [] + for _ in range(2): + router = ReplayRouter( + recorded, + live_forwarder=lambda req: completion(content="L"), + max_exchanges=2, + ) + for i in range(2): + router.next_response(_request(i)) + digests.append((router.served_request_digest, router.recorded_request_digest)) + assert digests[0] == digests[1] + assert digests[0][0] is not None + + +def test_natural_end_exposes_cut_point_accounting_too(): + recorded = _recorded(2) + router = ReplayRouter(recorded, live_forwarder=lambda req: completion(content="L")) + for i in range(3): + router.next_response(_request(i)) + assert router.n_replayed_exchanges == 2 + assert router.recorded_request_digest == comparable_request_digest( + recorded[1].request.body + ) + assert router.served_request_digest == comparable_request_digest(_request(1)) + + +# ── provenance: the cut_point block in source_provenance ────────────────── + + +def _build_config(tmp_path, run, **kwargs): + task = tmp_path / "real-task" + task.mkdir(exist_ok=True) + return build_rollout_config( + run, + task_path=task, + live_model="gemini-3.1-flash-lite-preview", + agent_env=build_agent_env("http://host:1/v1"), + timeout=123, + output_dir=tmp_path / "out", + rollout_name="demo-task__continued", + **kwargs, + ) + + +def test_provenance_gains_cut_point_block(tmp_path): + folder = write_run_folder(tmp_path / "run", exchanges=_recorded(3)) + run = load_run_folder(folder) + cfg = _build_config(tmp_path, run, max_exchanges=2) + cut = cfg.source_provenance["cut_point"] + assert cut["n_replayed_exchanges"] == 2 + assert cut["recorded_request_digest"] == comparable_request_digest( + run.exchanges[1].request.body + ) + assert cut["request_digest_basis"] == REQUEST_DIGEST_BASIS + # config-time blocks are computed from the recording, and say so — sandbox + # proxy mode (truncated upload, no live router to read back) keeps this + # basis in the final artifacts. Everything only the live run could + # observe is recorded honestly as null, never fabricated. + assert cut["accounting"] == "configured" + assert cut["served_request_digest"] is None + assert cut["divergences"] is None + assert cut["workspace_digest"] is None + assert "no live sandbox is reachable" in cut["workspace_digest_reason"] + assert "branch_stage" not in cut + # existing provenance fields are preserved alongside the new block + assert cfg.source_provenance["kind"] == "benchflow-continue" + assert cfg.source_provenance["n_recorded_exchanges"] == 3 + + +def test_provenance_natural_end_documents_cut_point(tmp_path): + folder = write_run_folder(tmp_path / "run", exchanges=_recorded(3)) + run = load_run_folder(folder) + cfg = _build_config(tmp_path, run) # no max_exchanges + cut = cfg.source_provenance["cut_point"] + assert cut["n_replayed_exchanges"] == 3 + assert cut["recorded_request_digest"] == comparable_request_digest( + run.exchanges[2].request.body + ) + assert cut["accounting"] == "configured" + + +# ── host-mode reconciliation: the served cut_point block ────────────────── + + +def test_served_cut_point_records_what_the_router_actually_served(): + """A run that went live before reaching the configured cut is visible. + + The router's served counters had no production readers — the block now + records n_replayed_exchanges plus the served/recorded request digests as + observed (accounting "served") and the requested configured_max_exchanges, + so a served-vs-configured divergence is detectable in artifacts. + """ + recorded = _recorded(3) + router = ReplayRouter( + recorded, live_forwarder=lambda req: completion(content="L"), max_exchanges=3 + ) + # the agent only ever replays 2 of the configured 3 exchanges + for i in range(2): + router.next_response(_request(i)) + + block = served_cut_point(router, configured_max_exchanges=3) + + reason = block.pop("workspace_digest_reason") + assert "never crossed the cut point" in reason + assert block == { + "n_replayed_exchanges": 2, + "served_request_digest": comparable_request_digest(_request(1)), + "recorded_request_digest": comparable_request_digest(recorded[1].request.body), + "request_digest_basis": REQUEST_DIGEST_BASIS, + "accounting": "served", + "divergences": [], + "workspace_digest": None, + "configured_max_exchanges": 3, + } + + +def test_served_cut_point_natural_end_and_stage_shape(): + recorded = _recorded(2) + router = ReplayRouter(recorded, live_forwarder=lambda req: completion(content="L")) + for i in range(3): + router.next_response(_request(i)) + + block = served_cut_point(router, branch_stage="post-research") + + # the cut WAS crossed here, but no workspace hook was configured — the + # reason says which of the two happened, never a fabricated digest + reason = block.pop("workspace_digest_reason") + assert "no live sandbox was reachable" in reason + assert block == { + "n_replayed_exchanges": 2, + "served_request_digest": comparable_request_digest(_request(1)), + "recorded_request_digest": comparable_request_digest(recorded[1].request.body), + "request_digest_basis": REQUEST_DIGEST_BASIS, + "accounting": "served", + "divergences": [], + "workspace_digest": None, + "branch_stage": "post-research", + } + + +def test_served_cut_point_carries_the_divergence_events(): + """Divergence accounting reaches the artifact, with the exchange index. + + Guards "fix(continue): replay divergence detection compares actual content + and records workspace digests": a divergence annotates (the continuation + still runs — fidelity caveats are recorded, not hidden, per RFC §3.5) and + the served block carries every event so the artifact is truthful about it. + """ + recorded = _recorded(2) + router = ReplayRouter(recorded, live_forwarder=lambda req: completion(content="L")) + router.next_response(_request(0)) # faithful + router.next_response({"messages": [{"role": "user", "content": "x"}] * 2}) + + block = served_cut_point(router) + + assert block["divergences"] == router.divergence_events + (event,) = block["divergences"] + assert event["exchange_index"] == 1 + assert event["served_request_digest"] == block["served_request_digest"] + assert event["recorded_request_digest"] == block["recorded_request_digest"] + assert block["served_request_digest"] != block["recorded_request_digest"] + + +# ── workspace digest at the cut point ───────────────────────────────────── + + +def test_workspace_digest_is_captured_once_as_the_run_crosses_the_cut(): + """Guards "fix(continue): replay divergence detection compares actual + content and records workspace digests": the RFC §3.5 workspace digest is + taken exactly once, at the first live-leg request — the moment the + replay-rebuilt workspace is complete — and lands in the served block.""" + calls: list[int] = [] + payload = {"digest": "sha256:" + "a" * 64, "basis": "b", "root": "/app"} + + def hook(): + calls.append(1) + return dict(payload) + + router = ReplayRouter( + _recorded(3), + live_forwarder=lambda req: completion(content="L"), + max_exchanges=1, + workspace_digest_fn=hook, + ) + router.next_response(_request(0)) + assert router.workspace_digest is None # cut not crossed yet + router.next_response(_request(1)) + router.next_response(_request(2)) + + assert calls == [1] # once, not per live request + assert router.workspace_digest == payload + block = served_cut_point(router) + assert block["workspace_digest"] == payload + assert "workspace_digest_reason" not in block + + +def test_workspace_digest_failure_is_recorded_never_fabricated(): + def hook(): + raise RuntimeError("sandbox is gone") + + router = ReplayRouter( + _recorded(1), + live_forwarder=lambda req: completion(content="L"), + workspace_digest_fn=hook, + ) + router.next_response(_request(0)) + router.next_response(_request(1)) # crosses the cut; hook raises + + assert router.workspace_digest is None + block = served_cut_point(router) + assert block["workspace_digest"] is None + assert "sandbox is gone" in block["workspace_digest_reason"] + + +async def test_compute_workspace_digest_runs_the_pipeline_in_the_sandbox(): + """The reusable helper execs the find|sort|sha256sum pipeline and parses + the marker-prefixed digest line, ignoring merged compose noise.""" + from benchflow.sandbox.protocol import ExecResult + from benchflow.sandbox.workspace_digest import ( + WORKSPACE_DIGEST_BASIS, + compute_workspace_digest, + ) + + commands: list[str] = [] + + class FakeSandbox: + async def exec(self, command, timeout_sec=None): + commands.append(command) + return ExecResult( + return_code=0, + stdout=( + "Found orphan containers for this project\n" + "BFWSDIGEST:" + "a" * 64 + " -\n" + ), + stderr="", + ) + + payload = await compute_workspace_digest(FakeSandbox()) + + assert payload == { + "digest": "sha256:" + "a" * 64, + "basis": WORKSPACE_DIGEST_BASIS, + "root": "/app", + } + (command,) = commands + # The null-safe, fail-closed form (PR #1046 second review, P1-A): names + # travel NUL-terminated and no stage hides behind a trailing pipe. + assert "find . -type f -print0" in command + assert "xargs -0 -r sha256sum" in command + assert "find . -type f | sort" not in command + assert "cd /app" in command + + +async def test_compute_workspace_digest_fails_closed_on_bad_output(): + from benchflow.sandbox.protocol import ExecResult + from benchflow.sandbox.workspace_digest import compute_workspace_digest + + class FailingSandbox: + async def exec(self, command, timeout_sec=None): + return ExecResult(return_code=1, stdout="", stderr="sh: cd: /app") + + class GarbageSandbox: + async def exec(self, command, timeout_sec=None): + return ExecResult( + return_code=0, stdout="BFWSDIGEST:not-a-digest\n", stderr="" + ) + + with pytest.raises(RuntimeError, match="failed"): + await compute_workspace_digest(FailingSandbox()) + with pytest.raises(RuntimeError, match="not a sha256"): + await compute_workspace_digest(GarbageSandbox()) + + +# Real-shell workspace-digest semantics (PR #1046 second review, P1-A). +# +# The finding: the digest pipeline fed filenames through newline-separated +# ``find | sort | xargs``, so a legal name like ``file with spaces.txt`` was +# word-split into several arguments; the inner ``sha256sum`` failed, the +# trailing ``| sha256sum`` succeeded anyway (a pipeline's status is its last +# command's), and two materially different workspaces recorded the *same* +# successful digest. These tests run the command in a genuine ``/bin/sh`` over +# a real temp dir — the shell, ``find``, ``sort`` and ``xargs`` under test are +# never faked. Only the digest *tools* may be shimmed on hosts without GNU +# coreutils (macOS): ``shasum -a 256`` is a real sha256 and BSD +# ``stat -f '%N %Lp'`` is a real stat — the shims translate flag spelling, +# not semantics. + + +def _digest_tool_bin(tmp_path: Path) -> Path | None: + """A PATH dir supplying ``sha256sum``/GNU-style ``stat``, if needed. + + Returns ``None`` when the host tools already speak the pipeline's dialect; + otherwise writes real-tool shims into ``tmp_path`` and returns it. Skips + the calling test when no real tool exists to shim. + """ + import shutil + import stat as stat_mod + import subprocess + + bin_dir = tmp_path / "digest-tools-bin" + needed = False + + def _shim(name: str, body: str) -> None: + path = bin_dir / name + bin_dir.mkdir(exist_ok=True) + path.write_text(f"#!/bin/sh\n{body}\n") + path.chmod(path.stat().st_mode | stat_mod.S_IXUSR) + + if shutil.which("sha256sum") is None: + shasum = shutil.which("shasum") + if shasum is None: + pytest.skip("host has neither sha256sum nor shasum") + _shim("sha256sum", f'exec "{shasum}" -a 256 "$@"') + needed = True + + probe = tmp_path / "digest-tools-probe" + probe.write_text("") + gnu_stat = ( + subprocess.run( + ["stat", "-c", "%n %a", str(probe)], capture_output=True + ).returncode + == 0 + ) + if not gnu_stat: + bsd_stat = ( + subprocess.run( + ["/usr/bin/stat", "-f", "%N %Lp", str(probe)], capture_output=True + ).returncode + == 0 + ) + if not bsd_stat: + pytest.skip("host has neither GNU nor BSD stat") + _shim( + "stat", + 'if [ "$1" = "-c" ] && [ "$2" = "%n %a" ]; then\n' + " shift 2\n" + ' exec /usr/bin/stat -f "%N %Lp" "$@"\n' + "fi\n" + 'exec /usr/bin/stat "$@"', + ) + needed = True + return bin_dir if needed else None + + +class _LocalShellSandbox: + """``exec`` = a real ``/bin/sh -c`` on the host — no faked shell semantics.""" + + def __init__(self, extra_path: Path | None) -> None: + self._extra_path = extra_path + + async def exec(self, command, timeout_sec=None): + import os + import subprocess + + from benchflow.sandbox.protocol import ExecResult + + env = dict(os.environ) + if self._extra_path is not None: + env["PATH"] = f"{self._extra_path}:{env.get('PATH', '')}" + proc = subprocess.run( + ["/bin/sh", "-c", command], + capture_output=True, + text=True, + env=env, + timeout=timeout_sec, + ) + return ExecResult( + return_code=proc.returncode, stdout=proc.stdout, stderr=proc.stderr + ) + + +async def test_workspace_digest_distinguishes_trees_differing_inside_a_spaced_name( + tmp_path: Path, +) -> None: + """The reviewer's exact repro: two workspaces whose only difference is the + CONTENT of ``file with spaces.txt`` (0 bytes vs 37) must record different + digests — the word-split pipeline recorded the same successful digest for + both. Guards ``fix(sandbox): workspace digests are null-safe and fail + closed``.""" + from benchflow.sandbox.workspace_digest import compute_workspace_digest + + ws_a = tmp_path / "ws_a" + ws_b = tmp_path / "ws_b" + ws_a.mkdir() + ws_b.mkdir() + (ws_a / "file with spaces.txt").write_text("") + (ws_b / "file with spaces.txt").write_text("x" * 37) + sandbox = _LocalShellSandbox(_digest_tool_bin(tmp_path)) + + digest_a = await compute_workspace_digest(sandbox, root=str(ws_a)) + digest_b = await compute_workspace_digest(sandbox, root=str(ws_b)) + + assert digest_a["digest"] != digest_b["digest"], ( + "materially different workspaces recorded the same digest — the " + "pipeline is still word-splitting filenames" + ) + # And the digest is still deterministic: same tree, same value. + digest_a_again = await compute_workspace_digest(sandbox, root=str(ws_a)) + assert digest_a_again["digest"] == digest_a["digest"] + + +async def test_workspace_digest_newline_filename_never_succeeds_wrongly( + tmp_path: Path, +) -> None: + """A filename containing a newline either digests correctly (different + contents => different digests) or fails closed with a reason — it must + never record the same successful digest for materially different trees. + Guards ``fix(sandbox): workspace digests are null-safe and fail closed``.""" + from benchflow.sandbox.workspace_digest import compute_workspace_digest + + name = "bad\nname.txt" + ws_a = tmp_path / "ws_a" + ws_b = tmp_path / "ws_b" + ws_a.mkdir() + ws_b.mkdir() + (ws_a / name).write_text("alpha") + (ws_b / name).write_text("beta") + sandbox = _LocalShellSandbox(_digest_tool_bin(tmp_path)) + + digests: list[str | None] = [] + for root in (ws_a, ws_b): + try: + payload = await compute_workspace_digest(sandbox, root=str(root)) + except RuntimeError: + digests.append(None) # fail-closed with a reason: acceptable + else: + digests.append(payload["digest"]) + if digests[0] is not None and digests[1] is not None: + assert digests[0] != digests[1], ( + "materially different workspaces recorded the same digest for a " + "newline-bearing filename" + ) + + +def test_update_continued_metadata_reconciles_cut_point_in_both_files(tmp_path): + """Host mode's post-run patch replaces the configured block with served.""" + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + configured = { + "n_replayed_exchanges": 3, + "recorded_request_digest": "sha256:configured", + "accounting": "configured", + } + (rollout_dir / "config.json").write_text( + json.dumps({"model": None, "source": {"cut_point": configured}}) + ) + (rollout_dir / "result.json").write_text( + json.dumps({"model": None, "source": {"cut_point": configured}}) + ) + traj = tmp_path / "llm_trajectory.jsonl" + traj.write_text("") + served = { + "n_replayed_exchanges": 2, + "served_request_digest": "sha256:served", + "recorded_request_digest": "sha256:recorded", + "accounting": "served", + "configured_max_exchanges": 3, + } + + update_continued_metadata( + rollout_dir, + live_model="live-model", + usage=summarize_llm_trajectory_usage(traj, n_recorded=2), + environment="docker", + cut_point=served, + ) + + config = json.loads((rollout_dir / "config.json").read_text()) + result = json.loads((rollout_dir / "result.json").read_text()) + assert config["source"]["cut_point"] == served + assert result["source"]["cut_point"] == served + + +# ── stage-named cuts ────────────────────────────────────────────────────── + + +def test_stage_named_cut_resolves_and_records_branch_stage(): + """stage_tags[stage] = the 1-based count of exchanges completed when the + stage closed; a cut at that stage replays exactly that many.""" + tags = {"env-ready": 1, "post-research": 2} + resolved, stage = resolve_cut_point( + None, stage_tags=tags, cut_stage="post-research" + ) + assert (resolved, stage) == (2, "post-research") + block = cut_point_provenance( + _recorded(3), max_exchanges=resolved, branch_stage=stage + ) + assert block["branch_stage"] == "post-research" + assert block["n_replayed_exchanges"] == 2 + + +def test_unknown_cut_stage_names_available_stages(): + tags = {"env-ready": 1, "post-research": 2} + with pytest.raises(ReplayCutPointError, match="env-ready, post-research"): + resolve_cut_point(None, stage_tags=tags, cut_stage="pre-verify") + + +def test_cut_stage_without_stage_tags_fails_closed(): + with pytest.raises(ReplayCutPointError, match="stage_tags"): + resolve_cut_point(None, cut_stage="post-research") + + +def test_cut_stage_and_max_exchanges_are_exclusive(): + with pytest.raises(ReplayCutPointError, match="not both"): + resolve_cut_point(2, stage_tags={"env-ready": 1}, cut_stage="env-ready") + + +@pytest.mark.parametrize("bad_tag", [0, -3]) +def test_stage_tag_below_one_fails_closed(bad_tag): + """A stage tag is a 1-based completed-exchange count — 0/negative is a + caller bug, rejected as ReplayCutPointError at resolve time.""" + with pytest.raises(ReplayCutPointError, match="1-based"): + resolve_cut_point( + None, stage_tags={"env-ready": bad_tag}, cut_stage="env-ready" + ) + + +def test_resolve_cut_point_passthrough_without_stage(): + assert resolve_cut_point(None) == (None, None) + assert resolve_cut_point(4) == (4, None) + + +# ── stage-named cuts resolve from the run folder's recorded registry ────── +# +# Guards "feat(continue): stage-named cut points": the reviewer-named gap was +# that stage-tagged cuts existed only as an SDK stage_tags override — `bench +# eval continue` exposed only numeric --max-exchanges, and nothing read the +# indices a run actually recorded. --cut-stage now resolves the exchange +# index from the run folder's stage_snapshots.json (written by the stage +# capture path since "feat(branch): record stage markers with trajectory +# exchange indices"). + + +def _registry_entry(exchanges_completed) -> dict: + """One stage_snapshots.json entry, in the shape write_stage_snapshots emits.""" + return { + "environment_ref": None, + "sandbox_ref": "bf-snap-1", + "layers": ["sandbox"], + "exchanges_completed": exchanges_completed, + } + + +def test_run_folder_loads_the_recorded_stage_registry(tmp_path): + folder = write_run_folder( + tmp_path / "run", + exchanges=_recorded(3), + stage_snapshots={ + "post-research": _registry_entry(2), + "pre-verify": _registry_entry(None), + }, + ) + run = load_run_folder(folder) + assert run.recorded_stages == ["post-research", "pre-verify"] + # Only stages with a usable index become tags; a null index stays visible + # in the registry but cannot name a cut. + assert run.stage_exchange_tags == {"post-research": 2} + + +def test_cut_stage_resolves_the_recorded_exchange_index(tmp_path): + """A recording with a marked stage resolves --cut-stage to that index.""" + folder = write_run_folder( + tmp_path / "run", + exchanges=_recorded(3), + stage_snapshots={"post-research": _registry_entry(2)}, + ) + run = load_run_folder(folder) + + tags = stage_tags_from_run(run, "post-research") + + assert resolve_cut_point(None, stage_tags=tags, cut_stage="post-research") == ( + 2, + "post-research", + ) + + +def test_cut_stage_with_no_recorded_stages_fails_closed(tmp_path): + folder = write_run_folder(tmp_path / "run", exchanges=_recorded(2)) + run = load_run_folder(folder) + with pytest.raises(ReplayCutPointError, match="recorded no stage snapshots"): + stage_tags_from_run(run, "post-research") + + +def test_cut_stage_unrecorded_stage_lists_what_was_recorded(tmp_path): + folder = write_run_folder( + tmp_path / "run", + exchanges=_recorded(2), + stage_snapshots={ + "env-ready": _registry_entry(0), + "post-research": _registry_entry(1), + }, + ) + run = load_run_folder(folder) + with pytest.raises( + ReplayCutPointError, match="env-ready, post-research" + ) as excinfo: + stage_tags_from_run(run, "pre-verify") + assert "unknown cut stage 'pre-verify'" in str(excinfo.value) + + +def test_cut_stage_recorded_without_an_index_fails_closed(tmp_path): + """A stage the gateway could not index is an honest null, not a guess.""" + folder = write_run_folder( + tmp_path / "run", + exchanges=_recorded(2), + stage_snapshots={"post-research": _registry_entry(None)}, + ) + run = load_run_folder(folder) + with pytest.raises(ReplayCutPointError, match="without an exchange index"): + stage_tags_from_run(run, "post-research") + + +def test_cut_stage_before_the_first_exchange_fails_closed(tmp_path): + """env-ready closes before any exchange — there is no prefix to replay.""" + folder = write_run_folder( + tmp_path / "run", + exchanges=_recorded(2), + stage_snapshots={"env-ready": _registry_entry(0)}, + ) + run = load_run_folder(folder) + with pytest.raises(ReplayCutPointError, match="before the first LLM exchange"): + stage_tags_from_run(run, "env-ready") + + +# ── stitched trajectory: prefix = the K parsed (replayed) exchanges ─────── + + +def test_stitched_prefix_truncated_at_cut(): + recorded_lines = ['{"a": 1}', '{"b": 2}', '{"c": 3}'] + lines = stitched_trajectory_lines( + recorded_lines, [exchange(completion(content="L"))], max_recorded=2 + ) + assert len(lines) == 3 + assert json.loads(lines[0]) == {"a": 1} + assert json.loads(lines[1]) == {"b": 2} + last = json.loads(lines[2]) + assert last["response"]["body"]["choices"][0]["message"]["content"] == "L" + + +def test_stitched_prefix_counts_parsed_exchanges_not_raw_lines(tmp_path): + """A malformed recorded line is never replayed — and never stitched. + + Replay counts *parsed* exchanges (load_llm_exchanges skips malformed + lines), so the stitched prefix must be the raw lines of the first K + parsed exchanges. Truncating by raw file-line index used to embed the + malformed (never-replayed) line and drop a replayed one, and shifted the + recorded/live usage-accounting boundary. + """ + recorded = _recorded(3) + folder = write_run_folder(tmp_path / "run", exchanges=recorded) + traj = folder / "trajectory" / "llm_trajectory.jsonl" + good_lines = traj.read_text().splitlines() + traj.write_text( + "\n".join([good_lines[0], "this is not json", *good_lines[1:]]) + "\n" + ) + run = load_run_folder(folder) + assert run.n_recorded_exchanges == 3 # malformed line skipped at load + + live = [exchange(completion(content="LIVE"))] + stitched = write_stitched_trajectory( + tmp_path / "rollout", run.exchange_lines, live, max_recorded=2 + ) + + lines = stitched.read_text().splitlines() + # exactly the 2 replayed exchanges' verbatim lines + the live suffix + assert len(lines) == 3 + assert lines[0] == good_lines[0] + assert lines[1] == good_lines[1] + assert "this is not json" not in stitched.read_text() + assert ( + json.loads(lines[2])["response"]["body"]["choices"][0]["message"]["content"] + == "LIVE" + ) + # the recorded/live usage split lands on the true replay boundary + usage = summarize_llm_trajectory_usage(stitched, n_recorded=2) + assert usage.recorded_total_tokens == 4 # 2 replayed exchanges x 2 tokens + assert usage.live_total_tokens == 2 # 1 live exchange x 2 tokens + + +# ── CLI: --max-exchanges reaches the orchestrator ───────────────────────── + + +def _patch_continue_run(monkeypatch, tmp_path, captured): + import benchflow.continue_run.orchestrator as orch + + async def fake_continue_run(folder, **kwargs): + captured.update(kwargs) + return SimpleNamespace( + rollout_dir=tmp_path / "r", + n_recorded=2, + n_live=1, + divergences=0, + rewards={"reward": 1.0}, + error=None, + ) + + monkeypatch.setattr(orch, "continue_run", fake_continue_run) + + +def test_cli_max_exchanges_reaches_orchestrator(tmp_path, monkeypatch): + captured: dict = {} + _patch_continue_run(monkeypatch, tmp_path, captured) + res = runner.invoke( + app, ["eval", "continue", str(tmp_path), "--max-exchanges", "2"] + ) + assert res.exit_code == 0, res.output + assert captured["max_exchanges"] == 2 + + +def test_cli_max_exchanges_defaults_to_all_recorded(tmp_path, monkeypatch): + captured: dict = {} + _patch_continue_run(monkeypatch, tmp_path, captured) + res = runner.invoke(app, ["eval", "continue", str(tmp_path)]) + assert res.exit_code == 0, res.output + assert captured["max_exchanges"] is None + + +def test_cli_out_of_range_cut_exits_clean(tmp_path): + folder = write_run_folder(tmp_path / "run", exchanges=_recorded(2)) + res = runner.invoke(app, ["eval", "continue", str(folder), "--max-exchanges", "99"]) + assert res.exit_code == 1 + assert "max_exchanges must be in 1..2" in res.output + assert "Traceback (most recent call last)" not in res.output + + +def test_cli_cut_stage_reaches_orchestrator(tmp_path, monkeypatch): + """Guards "feat(continue): stage-named cut points": the CLI exposes the + stage-named cut, not only numeric --max-exchanges.""" + captured: dict = {} + _patch_continue_run(monkeypatch, tmp_path, captured) + res = runner.invoke( + app, ["eval", "continue", str(tmp_path), "--cut-stage", "post-research"] + ) + assert res.exit_code == 0, res.output + assert captured["cut_stage"] == "post-research" + assert captured["max_exchanges"] is None + + +def test_cli_cut_stage_without_recorded_stages_exits_clean(tmp_path): + folder = write_run_folder(tmp_path / "run", exchanges=_recorded(2)) + res = runner.invoke( + app, ["eval", "continue", str(folder), "--cut-stage", "post-research"] + ) + assert res.exit_code == 1 + assert "recorded no stage snapshots" in res.output + assert "Traceback (most recent call last)" not in res.output + + +# ── one replay basis: stitched prefix, usage split and provenance agree ─── + + +class _AbandoningRollout: + """A run whose agent asks for fewer exchanges than the cut configured. + + Stands in for the real failure: the agent errors (or simply stops asking) + after turn 1 of a 2-exchange cut. ``run()`` drives the live router the way + the agent's proxy requests would, then writes the config/result pair a real + Rollout leaves behind — including the config-time ``cut_point`` block the + orchestrator is expected to reconcile. + """ + + served: ClassVar[int] = 1 + + def __init__(self, config) -> None: + self._config = config + self._rollout_dir = str( + Path(config.jobs_dir) / str(config.job_name) / str(config.rollout_name) + ) + + @classmethod + async def create(cls, config): + return cls(config) + + async def run(self): + for i in range(self.served): + _ROUTERS[0].next_response(_request(i)) + directory = Path(self._rollout_dir) + directory.mkdir(parents=True, exist_ok=True) + source = dict(self._config.source_provenance) + for name in ("config.json", "result.json"): + (directory / name).write_text( + json.dumps({"model": None, "source": source, "agent_result": {}}) + ) + return SimpleNamespace(rewards=None, error="Agent connection lost") + + +_ROUTERS: list[ReplayRouter] = [] + + +class _CapturingProxy: + """A ReplayProxy stand-in — no socket, but the same router seam.""" + + def __init__(self, router, **kwargs) -> None: + self.router = router + _ROUTERS.append(router) + + def start(self): + return self + + @property + def base_url(self) -> str: + return "http://127.0.0.1:1/v1" + + def stop(self) -> None: + return None + + +def _host_mode_run(tmp_path, monkeypatch, recorded): + """Wire continue_run's host path onto the fakes above.""" + import benchflow.continue_run.orchestrator as orch + import benchflow.rollout as rollout_module + + _ROUTERS.clear() + folder = write_run_folder(tmp_path / "run", exchanges=recorded) + tasks_dir = tmp_path / "tasks" + (tasks_dir / "demo-task").mkdir(parents=True) + monkeypatch.setattr(orch, "ReplayProxy", _CapturingProxy) + monkeypatch.setattr( + orch, "_host_proxy_binding", lambda env: ("127.0.0.1", "127.0.0.1") + ) + monkeypatch.setattr(rollout_module, "Rollout", _AbandoningRollout) + return folder, tasks_dir + + +async def test_host_mode_stitches_and_bills_only_the_exchanges_served( + tmp_path, monkeypatch +): + """An agent that stops short of the cut leaves consistent artifacts. + + Guards "fix(continue): stitch and account for the exchanges actually + served". In host proxy mode the stitched prefix, the recorded-vs-live token + split and the ``cut_point`` block were built on two different bases: the + first two on the *configured* K, the last on what the router had actually + served. An agent that exited after 1 of a configured 2 exchanges therefore + got a trajectory containing a recorded response it never received, that + response's tokens billed as replayed, and provenance that disagreed with + both — artifacts no experiment can be run on. + """ + recorded = _recorded(3) + folder, tasks_dir = _host_mode_run(tmp_path, monkeypatch, recorded) + + result = await continue_run( + folder, + tasks_dir=tasks_dir, + output_dir=tmp_path / "continued", + proxy_mode="host", + replay_only=True, + max_exchanges=2, + ) + + # the agent consumed one exchange, so one exchange is the whole trajectory + stitched = (result.rollout_dir / "trajectory" / "llm_trajectory.jsonl").read_text() + lines = stitched.splitlines() + assert len(lines) == 1 + assert ( + json.loads(lines[0])["response"]["body"]["choices"][0]["message"]["content"] + == "turn-0" + ) + assert "turn-1" not in stitched + + payload = json.loads((result.rollout_dir / "result.json").read_text()) + # the usage split is on the same basis: 1 replayed exchange x 2 tokens + assert payload["agent_result"]["usage_details"] == { + "source": "stitched_llm_trajectory", + "recorded_total_tokens": 2, + "live_total_tokens": 0, + } + assert payload["agent_result"]["total_tokens"] == 2 + # and so is the provenance — served, with the request K still visible + cut = payload["source"]["cut_point"] + reason = cut.pop("workspace_digest_reason") + assert "never crossed the cut point" in reason + assert cut == { + "n_replayed_exchanges": 1, + "served_request_digest": comparable_request_digest(_request(0)), + "recorded_request_digest": comparable_request_digest(recorded[0].request.body), + "request_digest_basis": REQUEST_DIGEST_BASIS, + "accounting": "served", + "divergences": [], + "workspace_digest": None, + "configured_max_exchanges": 2, + } + assert result.n_recorded == 1 + + +async def test_host_mode_cut_stage_resolves_from_the_run_folder(tmp_path, monkeypatch): + """End to end: ``cut_stage`` alone configures the cut the registry names. + + Guards "feat(continue): stage-named cut points" — no SDK ``stage_tags`` + override anywhere: the run folder's recorded ``stage_snapshots.json`` + (post-research closed at exchange 2) is the only source, and the + continuation's provenance records both the resolved K and the stage name. + """ + recorded = _recorded(3) + monkeypatch.setattr(_AbandoningRollout, "served", 2) + folder, tasks_dir = _host_mode_run(tmp_path, monkeypatch, recorded) + (folder / "stage_snapshots.json").write_text( + json.dumps( + {"schema_version": 1, "stages": {"post-research": _registry_entry(2)}} + ) + ) + + result = await continue_run( + folder, + tasks_dir=tasks_dir, + output_dir=tmp_path / "continued", + proxy_mode="host", + replay_only=True, + cut_stage="post-research", + ) + + payload = json.loads((result.rollout_dir / "result.json").read_text()) + cut = payload["source"]["cut_point"] + assert cut["n_replayed_exchanges"] == 2 + assert cut["configured_max_exchanges"] == 2 + assert cut["branch_stage"] == "post-research" + assert result.n_recorded == 2 + + +async def test_host_mode_full_consumption_is_unchanged(tmp_path, monkeypatch): + """The control: when the agent consumes the cut, nothing moves. + + The other half of "fix(continue): stitch and account for the exchanges + actually served" — served and configured coincide for every run that + reached its cut, which is the ordinary case, and those artifacts must read + exactly as they did before. + """ + recorded = _recorded(3) + monkeypatch.setattr(_AbandoningRollout, "served", 2) + folder, tasks_dir = _host_mode_run(tmp_path, monkeypatch, recorded) + + result = await continue_run( + folder, + tasks_dir=tasks_dir, + output_dir=tmp_path / "continued", + proxy_mode="host", + replay_only=True, + max_exchanges=2, + ) + + stitched = (result.rollout_dir / "trajectory" / "llm_trajectory.jsonl").read_text() + assert len(stitched.splitlines()) == 2 + payload = json.loads((result.rollout_dir / "result.json").read_text()) + assert payload["agent_result"]["usage_details"]["recorded_total_tokens"] == 4 + assert payload["source"]["cut_point"]["n_replayed_exchanges"] == 2 + assert result.n_recorded == 2 + + +def test_continuation_artifacts_are_written_on_one_basis(tmp_path): + """Stitch, usage split and provenance come from a single prefix length. + + The structural half of "fix(continue): stitch and account for the exchanges + actually served": the two proxy modes pick different bases (served vs + configured) but each writes all three artifacts through this one call, so a + future caller cannot reintroduce a per-artifact basis. Here the basis is + the sandbox one — configured, truncated upload, no live router — and the + block says so. + """ + recorded = _recorded(3) + folder = write_run_folder(tmp_path / "run", exchanges=recorded) + run = load_run_folder(folder) + rollout_dir = tmp_path / "rollout" + rollout_dir.mkdir() + (rollout_dir / "result.json").write_text(json.dumps({"model": None})) + + stitched = write_continuation_artifacts( + rollout_dir, + run, + [exchange(completion(content="LIVE"))], + n_recorded=2, + cut_point=cut_point_provenance(run.exchanges, max_exchanges=2), + live_model="gemini-3.1-flash-lite-preview", + ) + + assert len(stitched.read_text().splitlines()) == 3 + payload = json.loads((rollout_dir / "result.json").read_text()) + assert payload["agent_result"]["usage_details"]["recorded_total_tokens"] == 4 + assert payload["agent_result"]["usage_details"]["live_total_tokens"] == 2 + assert payload["source"]["cut_point"]["accounting"] == "configured" + assert payload["source"]["cut_point"]["n_replayed_exchanges"] == 2 diff --git a/tests/continue_run/test_replay_proxy.py b/tests/continue_run/test_replay_proxy.py index f7c1d55d5..b2ad8b3bf 100644 --- a/tests/continue_run/test_replay_proxy.py +++ b/tests/continue_run/test_replay_proxy.py @@ -13,6 +13,7 @@ ReplayRouter, completion_to_sse, ) +from benchflow.trajectories.types import LLMExchange, LLMRequest, LLMResponse from ._helpers import completion, exchange @@ -72,6 +73,82 @@ def test_divergence_strict_raises(): router.next_response({"messages": [{}, {}]}) +def test_same_count_content_change_is_detected_and_recorded(): + """Guards "fix(continue): replay divergence detection compares actual + content and records workspace digests": a prompt/content/tool change that + keeps the message count was invisible to the count-only heuristic. The + router now compares comparable-projection digests per replayed exchange + and records a divergence event carrying the exchange index and both + digests.""" + recorded = [exchange(completion(content="a"), n_request_messages=1)] + router = ReplayRouter(recorded) + + # same count (1 message), different content + router.next_response({"messages": [{"role": "user", "content": "DIFFERENT"}]}) + + assert router.divergences == 1 + (event,) = router.divergence_events + assert event["exchange_index"] == 0 + assert event["n_messages_served"] == event["n_messages_recorded"] == 1 + assert event["served_request_digest"] != event["recorded_request_digest"] + assert event["served_request_digest"].startswith("sha256:") + + +def test_same_count_tool_definition_change_is_detected(): + """The tools half of the same fix: a changed tool schema at an unchanged + message count is a real divergence.""" + recorded = [ + LLMExchange( + request=LLMRequest( + body={ + "messages": [{"role": "user"}], + "tools": [{"function": {"name": "bash"}}], + } + ), + response=LLMResponse(status_code=200, body=completion(content="a")), + ) + ] + router = ReplayRouter(recorded) + router.next_response( + {"messages": [{"role": "user"}], "tools": [{"function": {"name": "python"}}]} + ) + assert router.divergences == 1 + + +def test_same_count_content_divergence_strict_raises(): + recorded = [exchange(completion(content="a"), n_request_messages=1)] + router = ReplayRouter(recorded, strict_divergence=True) + with pytest.raises(ReplayDivergenceError): + router.next_response({"messages": [{"role": "user", "content": "changed"}]}) + + +def test_transport_fields_do_not_false_positive_divergence(): + """The recorded request is a normalized projection: the live request's + placeholder model, stream flag and sampling params are transport, not + content, and must not flag every exchange as divergent.""" + recorded = [exchange(completion(content="a"), n_request_messages=2)] + router = ReplayRouter(recorded) + router.next_response( + { + "messages": [{"role": "user"}] * 2, + "model": "openai/replay", + "stream": True, + "temperature": 0.0, + } + ) + assert router.divergences == 0 + assert router.divergence_events == [] + + +def test_count_mismatch_records_a_divergence_event_too(): + recorded = [exchange(completion(content="a"), n_request_messages=3)] + router = ReplayRouter(recorded) + router.next_response({"messages": [{}, {}]}) + (event,) = router.divergence_events + assert event["n_messages_served"] == 2 + assert event["n_messages_recorded"] == 3 + + def test_recorded_failure_passed_through(): recorded = [exchange({"error": {"message": "boom"}}, status=500)] router = ReplayRouter(recorded) diff --git a/tests/test_ablate_cli.py b/tests/test_ablate_cli.py new file mode 100644 index 000000000..9f81d9e31 --- /dev/null +++ b/tests/test_ablate_cli.py @@ -0,0 +1,2512 @@ +"""Regression tests for ``bench eval ablate`` and its ablation library. + +Guards "feat(cli): bench eval ablate — stage-level ablation over branch +children" (docs/rollout-branching-rfc.md §5 / WS-4c; FrontierPhysics#73). PR +number to be added on submission. + +The branch machinery underneath (stage snapshots, per-child deltas, the +skill-mode fresh-rollout child) already had unit coverage; what it had no +surface for was *running an experiment*. These tests pin that surface: arm +specs map onto exactly the deltas the engine executes, a request the engine +would reject fails closed before the parent run costs anything, the report is +deterministic, per-arm errors are isolated (the arms that ran keep their +rewards) and exit 1, and the attribution line stays an observation rather than +a causal claim. + +Section 6 guards "feat(ablate): per-test attribution so a scalar tie cannot +hide a behavioral difference" — the measured case where both arms of a skills +ablation scored 0.00 while their two sub-tests flipped in opposite directions, +and the tool printed "no difference in this comparison". + +Unit tests against fakes and a patched engine — no Docker, Daytona, or API +keys. +""" + +from __future__ import annotations + +import asyncio +import json +import re +from pathlib import Path +from types import SimpleNamespace +from typing import Any, ClassVar, cast + +import click +import pytest +import typer +from typer.testing import CliRunner + +from benchflow.ablate import ( + ARM_KIND_CONFIG, + ARM_KIND_ENV, + ARM_KIND_INJECT, + ARM_KIND_SKILL_MODE, + AblationReport, + AblationRequest, + AblationSpecError, + ArmOutcome, + attribute, + differing_tests, + parse_arm, + parse_arms, + resolve_ablation_task, + run_ablation, + sub_test_attribution, + validate_arms_for_stage, + write_ablation_report, +) +from benchflow.branch import UNSCORED_KEY +from benchflow.branch_delta import BranchDelta +from benchflow.cli.main import app +from benchflow.rollout_branch import CHILD_WALL_CLOCK_KEY +from benchflow.trajectories.tree import RolloutTree + +runner = CliRunner() + + +def _flat(text: str) -> str: + """Collapse Rich's terminal wrapping so a message asserts as one sentence.""" + return re.sub(r"\s+", " ", text) + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_CLI_MD = _REPO_ROOT / "docs" / "reference" / "cli.md" +_DOC_HEADER = "### bench eval ablate" + + +def _task_dir(tmp_path: Path, name: str = "task") -> Path: + task = tmp_path / name + task.mkdir(parents=True) + (task / "task.toml").write_text('version = "1.0"\n', encoding="utf-8") + (task / "instruction.md").write_text("solve it\n", encoding="utf-8") + return task + + +def _report( + *, + arms: list[ArmOutcome], + parent_reward: float | None = 1.0, + value: float | None = 0.5, + error: str | None = None, +) -> AblationReport: + return AblationReport( + task_id="demo", + task_path="/tasks/demo", + stage="env-ready", + snapshot_layers=["sandbox"], + agent="claude-agent-acp", + model="claude-sonnet", + sandbox="docker", + arms=arms, + parent_reward=parent_reward, + parent_run_dir="/out/ablation/demo", + value=value, + error=error, + ) + + +def _skill_outcomes(with_skill: float, no_skill: float) -> list[ArmOutcome]: + return [ + ArmOutcome( + name="with-skill", + kind=ARM_KIND_SKILL_MODE, + delta=BranchDelta(skill_mode="with-skill").provenance_dict(), + delta_execution="fresh-rollout", + reward=with_skill, + wall_clock_sec=61.0, + node_id="n1", + artifacts="/out/ablation/demo/branches/root/children/n1", + ), + ArmOutcome( + name="no-skill", + kind=ARM_KIND_SKILL_MODE, + delta=BranchDelta(skill_mode="no-skill").provenance_dict(), + delta_execution="fresh-rollout", + reward=no_skill, + wall_clock_sec=44.0, + node_id="n2", + artifacts="/out/ablation/demo/branches/root/children/n2", + ), + ] + + +def _patched_run(monkeypatch, report: AblationReport) -> list[AblationRequest]: + """Patch the engine and capture the request the CLI built for it.""" + seen: list[AblationRequest] = [] + + async def fake_run_ablation(request: AblationRequest) -> AblationReport: + seen.append(request) + return report + + monkeypatch.setattr("benchflow.ablate.run_ablation", fake_run_ablation) + return seen + + +# 1. Arm specs -> BranchDelta (the library-level mapping) + + +def test_arm_specs_map_onto_the_deltas_the_engine_executes(tmp_path: Path) -> None: + """Each v1 arm kind lowers to exactly one executable BranchDelta field. + + ``skill_mode`` and ``injected_prompt`` are the only two fields the branch + engine executes (RFC §3.3); an arm that lowered to anything else would + fail closed at fork time, after a full parent run. + """ + plan = tmp_path / "oracle-plan.md" + plan.write_text("1. read the spec\n2. patch the module\n", encoding="utf-8") + + arms = parse_arms(f"with-skill,no-skill,inject:{plan}") + + assert [arm.name for arm in arms] == ["with-skill", "no-skill", f"inject:{plan}"] + assert [arm.kind for arm in arms] == [ + ARM_KIND_SKILL_MODE, + ARM_KIND_SKILL_MODE, + ARM_KIND_INJECT, + ] + assert arms[0].delta == BranchDelta(skill_mode="with-skill") + assert arms[1].delta == BranchDelta(skill_mode="no-skill") + assert arms[2].delta.injected_prompt == plan.read_text() + assert arms[2].delta.skill_mode is None + assert arms[2].source == str(plan) + # The injected text is hash-only in provenance (#908) — never the raw plan. + provenance = arms[2].delta.provenance_dict() + assert provenance["injected_prompt_sha256"].startswith("sha256:") + assert "read the spec" not in json.dumps(provenance) + + +def test_whitespace_around_arms_is_tolerated_but_an_empty_entry_is_not() -> None: + """A dropped arm would publish a table with a missing comparison.""" + assert [arm.name for arm in parse_arms(" with-skill , no-skill ")] == [ + "with-skill", + "no-skill", + ] + with pytest.raises(AblationSpecError, match="empty arm"): + parse_arms("with-skill,,no-skill") + + +def test_single_and_duplicate_arm_requests_fail_closed() -> None: + """A fork needs >= 2 children, and the report keys arms by name.""" + with pytest.raises(AblationSpecError, match=">= 2 children"): + parse_arms("with-skill") + with pytest.raises(AblationSpecError, match="duplicate arm"): + parse_arms("no-skill,no-skill") + + +def test_injection_arm_needs_a_readable_non_empty_file(tmp_path: Path) -> None: + empty = tmp_path / "empty.md" + empty.write_text(" \n", encoding="utf-8") + with pytest.raises(AblationSpecError, match="cannot read its injection file"): + parse_arm(f"inject:{tmp_path / 'missing.md'}") + with pytest.raises(AblationSpecError, match="empty injection file"): + parse_arm(f"inject:{empty}") + with pytest.raises(AblationSpecError, match="names no file"): + parse_arm("inject:") + + +def test_skill_arms_are_rejected_away_from_env_ready() -> None: + """The engine's own gate (skills are deployed by install_agent()), applied + before the ablation pays for a full parent run.""" + arms = parse_arms("with-skill,no-skill") + assert validate_arms_for_stage(arms, "env-ready") == "env-ready" + with pytest.raises(AblationSpecError, match="needs --at-stage 'env-ready'"): + validate_arms_for_stage(arms, "pre-verify") + + +def test_injection_arms_run_at_env_ready_and_at_later_boundaries( + tmp_path: Path, +) -> None: + """An ``inject:`` arm at ``env-ready`` is supported, not rejected. + + WS-4c flagged that combination as unsound because the child was run in + place on a world where ``install_agent()`` had been rolled back. The + resolution routes it through the same fresh-rollout path the skills arms + use rather than banning it, so the pre-flight must keep accepting it — + a regression here would silently turn a supported ablation into a + spec error. + """ + plan = tmp_path / "plan.md" + plan.write_text("Follow this plan.", encoding="utf-8") + other = tmp_path / "other.md" + other.write_text("Follow that plan.", encoding="utf-8") + arms = parse_arms(f"inject:{plan},no-skill") + assert validate_arms_for_stage(arms, "env-ready") == "env-ready" + inject_only = parse_arms(f"inject:{plan},inject:{other}") + assert validate_arms_for_stage(inject_only, "env-ready") == "env-ready" + assert validate_arms_for_stage(inject_only, "pre-verify") == "pre-verify" + + +def test_config_arm_specs_map_onto_config_override_deltas(tmp_path: Path) -> None: + """``config:`` lowers to the config_override delta. + + Guards "feat(branch): execute config_override deltas as fresh rollouts + from env-ready": the arm parses through the same loader as the run-level + override (inline JSON or an ``@file`` ref), its commas stay content when + nested in JSON braces, and its provenance records the #790 sha + keys — + never the raw patch. + """ + inline = 'config:{"agent": {"timeout_sec": 7}, "metadata": {"x": 1}}' + overlay_file = tmp_path / "overlay.json" + overlay_file.write_text('{"agent": {"timeout_sec": 9}}', encoding="utf-8") + + arms = parse_arms(f"with-skill,{inline},config:@{overlay_file}") + + assert [arm.kind for arm in arms] == [ + ARM_KIND_SKILL_MODE, + ARM_KIND_CONFIG, + ARM_KIND_CONFIG, + ] + assert arms[1].name == inline # the spec survives the comma-aware split + assert arms[1].delta.config_override == { + "agent": {"timeout_sec": 7}, + "metadata": {"x": 1}, + } + assert arms[2].delta.config_override == {"agent": {"timeout_sec": 9}} + assert arms[2].source == str(overlay_file) + provenance = arms[1].delta.provenance_dict() + assert provenance["config_override_sha256"].startswith("sha256:") + assert provenance["config_override_keys"] == ["agent", "metadata"] + + +def test_braces_and_commas_inside_quoted_json_strings_stay_content() -> None: + """The arm splitter tracks JSON string state, not just brace depth. + + Guards the quote-aware ``_split_arm_specs`` fix for the review finding on + PR #1046: braces, brackets and commas inside a quoted inline-JSON string + value are content — ``{"metadata": {"note": "a{b,c]d"}}`` is one arm, and + an escaped quote (``\\"``) does not end the string. A bare quote at depth + zero (e.g. in an ``inject:`` path) keeps its historical literal meaning. + """ + # The close-braces inside the string zero out a depth-only counter, so the + # comma right after them split the spec mid-JSON before the fix. + spooky = 'config:{"metadata": {"close": "}}", "note": "a{b,c]d", "quote": "say \\"hi\\", ok"}}' + arms = parse_arms(f"no-skill,{spooky}") + assert [arm.kind for arm in arms] == [ARM_KIND_SKILL_MODE, ARM_KIND_CONFIG] + assert arms[1].delta.config_override == { + "metadata": {"close": "}}", "note": "a{b,c]d", "quote": 'say "hi", ok'} + } + + from benchflow.ablate import _split_arm_specs + + # A bare quote at depth zero never opens a string: the comma still splits. + assert _split_arm_specs('inject:some"quoted.md,no-skill') == [ + 'inject:some"quoted.md', + "no-skill", + ] + + +def test_malformed_config_arm_specs_fail_closed(tmp_path: Path) -> None: + """An unparsable, empty, or scorer-touching config arm dies at parse time + — before the parent run costs anything, exactly where the engine's own + allowlist gate would kill the fork.""" + with pytest.raises(AblationSpecError, match="carries no patch"): + parse_arm("config:") + with pytest.raises(AblationSpecError, match="cannot load its config patch"): + parse_arm(f"config:@{tmp_path / 'missing.yaml'}") + with pytest.raises(AblationSpecError, match="empty patch"): + parse_arm("config:{}") + with pytest.raises(AblationSpecError, match="may only patch"): + parse_arm('config:{"verifier": {"timeout_sec": 1}}') + + +def test_config_arms_are_rejected_away_from_env_ready() -> None: + """The engine's own stage gate (the child's setup() applies the patch), + applied before the ablation pays for a full parent run.""" + arms = parse_arms('no-skill,config:{"agent": {"timeout_sec": 7}}') + assert validate_arms_for_stage(arms, "env-ready") == "env-ready" + config_only = parse_arms( + 'config:{"agent": {"timeout_sec": 7}},config:{"agent": {"timeout_sec": 9}}' + ) + with pytest.raises(AblationSpecError, match="needs --at-stage 'env-ready'"): + validate_arms_for_stage(config_only, "pre-verify") + + +def test_env_arm_specs_map_onto_environment_ref_deltas() -> None: + """``env:`` lowers to the environment_ref delta verbatim. + + Guards "feat(branch): execute service-level environment_ref deltas from + env-ready": the ref is recorded exactly as written (the registry + content-addresses it at resolution), an empty ref fails at parse time, + and the arm is env-ready-only like every fresh-child delta. + """ + arms = parse_arms("no-skill,env:env0@outage") + assert arms[1].kind == ARM_KIND_ENV + assert arms[1].delta == BranchDelta(environment_ref="env0@outage") + assert arms[1].delta.provenance_dict()["environment_ref"] == "env0@outage" + with pytest.raises(AblationSpecError, match="names no environment"): + parse_arm("env:") + env_only = parse_arms("env:env0@prod,env:env0@outage") + assert validate_arms_for_stage(env_only, "env-ready") == "env-ready" + with pytest.raises(AblationSpecError, match="needs --at-stage 'env-ready'"): + validate_arms_for_stage(env_only, "post-verify") + + +async def test_an_env_arm_against_a_manifest_less_task_fails_before_the_parent_runs( + tmp_path: Path, monkeypatch +) -> None: + """The env arm's content gates need the parent's manifest, so they run in + run_ablation — but still *before* the parent run: a task that declares no + environment has no plane to swap, and the request dies without building a + single rollout.""" + built = _patch_rollout(monkeypatch) + request = _request(tmp_path, "no-skill,env:env0@outage") + + with pytest.raises(AblationSpecError, match="no environment manifest"): + await run_ablation(request) + + assert built == [] # nothing ran, nothing was paid for + + +def test_env_ready_ablation_requires_the_container_layer(tmp_path: Path) -> None: + """Every arm of an ``env-ready`` ablation re-installs the agent for itself, + so the stage snapshot has to carry the container layer — the mirror of the + engine's own gate, paid before a full parent run instead of after it.""" + plan = tmp_path / "plan.md" + plan.write_text("Follow this plan.", encoding="utf-8") + arms = parse_arms(f"inject:{plan},no-skill") + + with pytest.raises(AblationSpecError, match="sandbox"): + validate_arms_for_stage( + arms, "env-ready", snapshot_layers=frozenset({"environment"}) + ) + assert ( + validate_arms_for_stage( + arms, "env-ready", snapshot_layers=frozenset({"environment", "sandbox"}) + ) + == "env-ready" + ) + # The command's own default is the container layer, so the CLI path is + # never the one that trips this. + assert AblationRequest( + task_path=tmp_path, arms=arms, agent="claude-agent-acp" + ).snapshot_layers == frozenset({"sandbox"}) + + +def test_tasks_dir_must_resolve_to_exactly_one_task(tmp_path: Path) -> None: + """An ablation's axis is the arms; several tasks is a request it cannot + answer, not a batch to expand.""" + collection = tmp_path / "tasks" + _task_dir(collection, "alpha") + assert resolve_ablation_task(collection / "alpha") == collection / "alpha" + assert resolve_ablation_task(collection) == collection / "alpha" + _task_dir(collection, "beta") + with pytest.raises(AblationSpecError, match="holds 2 tasks"): + resolve_ablation_task(collection) + + +# 2. Attribution: verdicts are observations, not causal claims + + +def test_attribution_pairs_the_skill_arms_and_names_the_stage() -> None: + outcomes = _skill_outcomes(with_skill=1.0, no_skill=0.0) + attribute(outcomes, parent_reward=0.0, stage="env-ready") + + assert [arm.reference for arm in outcomes] == ["no-skill", "with-skill"] + assert outcomes[0].verdict == ( + "passes (1.00) where no-skill fails (0.00) at env-ready — this delta " + "decides the outcome when applied at env-ready (1 run per arm)" + ) + assert outcomes[1].verdict.startswith("fails (0.00) where with-skill passes (1.00)") + # No claim about boundaries this ablation never forked. + assert "at or before" not in outcomes[1].verdict + + +def test_attribution_reports_no_difference_and_falls_back_to_the_parent( + tmp_path: Path, +) -> None: + both_pass = _skill_outcomes(with_skill=1.0, no_skill=1.0) + attribute(both_pass, parent_reward=1.0, stage="env-ready") + assert "no difference in this comparison" in both_pass[0].verdict + + injected = ArmOutcome(name="inject:plan.md", kind=ARM_KIND_INJECT, reward=1.0) + attribute([injected], parent_reward=0.0, stage="pre-verify") + assert injected.reference == "parent" + assert "where parent fails (0.00) at pre-verify" in injected.verdict + + orphan = ArmOutcome(name="inject:plan.md", kind=ARM_KIND_INJECT, reward=0.5) + attribute([orphan], parent_reward=None, stage="pre-verify") + assert orphan.reference is None + assert "no counterpart arm or parent reward" in orphan.verdict + + +def test_report_json_is_deterministic_and_carries_arm_provenance( + tmp_path: Path, +) -> None: + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + + first = write_ablation_report(report, tmp_path / "out").read_text() + second = write_ablation_report(report, tmp_path / "out").read_text() + + assert first == second + payload = json.loads(first) + assert payload["task"]["id"] == "demo" + assert payload["stage"] == "env-ready" + assert payload["parent"]["reward"] == 1.0 + assert [arm["delta"]["skill_mode"] for arm in payload["arms"]] == [ + "with-skill", + "no-skill", + ] + assert [arm["status"] for arm in payload["arms"]] == ["pass", "fail"] + assert all(arm["delta_execution"] == "fresh-rollout" for arm in payload["arms"]) + assert report.has_errors is False + + +# 3. The engine path: arms map onto forked children, failures stay isolated + + +class _FakeRollout: + """A Rollout stand-in that forks children the way the branch engine does. + + Mirrors the engine contract WS-4a/WS-4b pinned: one child node per delta, + attached and run in order, each carrying its ``delta`` provenance, its + ``delta_execution`` and its ``reward``; a child that raises leaves an + attached node with no reward and no further children. + """ + + rewards: ClassVar[list[Any]] = [1.0, 0.0] + parent_rewards: ClassVar[dict | None] = {"reward": 1.0} + + def __init__(self, config) -> None: + self.config = config + self.tree = RolloutTree() + self.result = None + self.calls: list[str] = [] + self._rollout_dir = ( + Path(config.jobs_dir) / str(config.job_name) / str(config.rollout_name) + ) + + async def setup(self) -> None: + self.calls.append("setup") + + async def start(self) -> None: + self.calls.append("start") + + async def install_agent(self) -> None: + self.calls.append("install_agent") + + async def connect(self) -> None: + self.calls.append("connect") + + async def execute(self) -> None: + self.calls.append("execute") + + async def verify(self) -> dict | None: + self.calls.append("verify") + return self.parent_rewards + + async def cleanup(self) -> None: + self.calls.append("cleanup") + + async def branch_at_stage(self, stage, n, *, deltas=None) -> float: + self.calls.append(f"branch_at_stage:{stage}") + assert deltas is not None and len(deltas) == n + returns: list[float] = [] + for delta, outcome in zip(deltas, self.rewards, strict=True): + child = self.tree.attach(self.tree.root) + child.state["delta"] = delta.provenance_dict() + if delta.skill_mode is not None: + child.state["delta_execution"] = "fresh-rollout" + child.state[CHILD_WALL_CLOCK_KEY] = 12.5 + if isinstance(outcome, Exception): + raise outcome + child.state["reward"] = float(outcome) + returns.append(float(outcome)) + return sum(returns) / len(returns) + + +def _patch_rollout(monkeypatch, cls: type = _FakeRollout) -> list[Any]: + """Patch the Rollout the ablation builds; return the instances it built.""" + built: list[Any] = [] + + class _Capturing(cls): # type: ignore[valid-type, misc] + def __init__(self, config) -> None: + super().__init__(config) + built.append(self) + + monkeypatch.setattr("benchflow.rollout.Rollout", _Capturing) + return built + + +def _request(tmp_path: Path, arms: str = "with-skill,no-skill") -> AblationRequest: + return AblationRequest( + task_path=_task_dir(tmp_path), + arms=parse_arms(arms), + agent="claude-agent-acp", + model="claude-sonnet", + out_dir=tmp_path / "out", + ) + + +async def test_run_ablation_forks_one_child_per_arm_and_reads_their_rewards( + tmp_path: Path, monkeypatch +) -> None: + built = _patch_rollout(monkeypatch) + request = _request(tmp_path) + + report = await run_ablation(request) + + # Not Rollout.run(): the branch has to happen while the sandbox is still + # up, so cleanup is last and the fork comes before it. + assert built[0].calls == [ + "setup", + "start", + "install_agent", + "connect", + "execute", + "verify", + "branch_at_stage:env-ready", + "cleanup", + ] + assert [arm.name for arm in report.arms] == ["with-skill", "no-skill"] + assert [arm.reward for arm in report.arms] == [1.0, 0.0] + assert [arm.status for arm in report.arms] == ["pass", "fail"] + assert [arm.delta["skill_mode"] for arm in report.arms] == [ + "with-skill", + "no-skill", + ] + assert all(arm.wall_clock_sec == 12.5 for arm in report.arms) + assert report.arms[0].artifacts.endswith("branches/root/children/n1") + assert report.parent_reward == 1.0 + assert report.value == 0.5 + assert report.has_errors is False + assert "decides the outcome when applied at env-ready" in report.arms[0].verdict + + +async def test_the_ablation_parent_states_the_no_skill_mode_the_gate_requires( + tmp_path: Path, monkeypatch +) -> None: + """The parent an ablation runs is ``no-skill`` by statement, not by default. + + Guards the fix from "fix(branch): gate skill deltas on the parent's own + skill mode". The arms fork the parent's own ``env-ready`` image, and a + ``with-skill`` parent bakes its pack into that image — a ``no-skill`` arm + would restore the pack and still be labelled ``no-skill``. The branch + engine now refuses that fork, so this command was correct only for as long + as ``RolloutConfig.from_legacy`` happened to default ``skill_mode`` to + ``no-skill``. It says so itself now, and a change to that default cannot + quietly move every ablation to the wrong side of the gate. + """ + built = _patch_rollout(monkeypatch) + + await run_ablation(_request(tmp_path)) + + assert built[0].config.skill_mode == "no-skill" + assert built[0].config.recorded_skill_mode == "no-skill" + + +async def test_run_ablation_cleans_up_and_isolates_a_failing_arm( + tmp_path: Path, monkeypatch +) -> None: + """A child failure propagates out of branch_at_stage (it is never scored + 0.0), so the arm that raised carries the error, the arms after it report as + skipped, and the arm that already ran keeps its reward.""" + monkeypatch.setattr(_FakeRollout, "rewards", [1.0, RuntimeError("no snapshot")]) + built = _patch_rollout(monkeypatch) + request = _request(tmp_path, "with-skill,no-skill") + + report = await run_ablation(request) + + assert built[0].calls[-1] == "cleanup" + assert [arm.status for arm in report.arms] == ["pass", "error"] + assert report.arms[0].reward == 1.0 + assert "no snapshot" in report.arms[1].error + assert report.arms[1].verdict == "errored before scoring — no reward to attribute" + assert report.value is None + assert report.has_errors is True + + +class _UnscoredRollout(_FakeRollout): + """A fork whose children ran but produced no verifier reward. + + The engine contract this mirrors is pinned directly in + ``tests/test_branch_skill_delta.py`` and ``tests/test_rollout_branch.py``: + a child whose ``verify()`` yielded nothing gets ``UNSCORED_KEY`` (the + reason) on its node and *no* ``reward`` key, and the branch returns ``None`` + because there is nothing to average. This is the live-run failure that + reported two confident ``0.00``s. + """ + + reason: ClassVar[str] = ( + "branch child regex-email-parser produced no verifier reward — " + "No reward file found at /out/verifier/reward.txt or reward.json" + ) + + async def branch_at_stage(self, stage, n, *, deltas=None) -> float | None: + self.calls.append(f"branch_at_stage:{stage}") + assert deltas is not None and len(deltas) == n + for delta in deltas: + child = self.tree.attach(self.tree.root) + child.state["delta"] = delta.provenance_dict() + child.state["delta_execution"] = "fresh-rollout" + child.state[CHILD_WALL_CLOCK_KEY] = 118.7 + child.state[UNSCORED_KEY] = self.reason + return None + + +async def test_an_unscored_child_is_an_arm_error_never_a_fabricated_zero( + tmp_path: Path, monkeypatch +) -> None: + """The regression that made a real ablation lie. + + Both children ran and neither was scored. Every arm must report ``error`` + with the reason, carry a ``null`` reward in ``ablation.json``, claim + nothing in its verdict (never "no difference" between two arms that were + never scored), leave V undefined, and make the command exit non-zero. + """ + _patch_rollout(monkeypatch, _UnscoredRollout) + + report = await run_ablation(_request(tmp_path)) + + assert [arm.status for arm in report.arms] == ["error", "error"] + assert [arm.reward for arm in report.arms] == [None, None] + assert all("No reward file found" in arm.error for arm in report.arms) + assert all("no difference" not in arm.verdict for arm in report.arms) + assert all("0.00" not in arm.verdict for arm in report.arms) + assert report.value is None + assert report.has_errors is True + + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert [arm["reward"] for arm in payload["arms"]] == [None, None] + assert [arm["status"] for arm in payload["arms"]] == ["error", "error"] + + +def test_cli_exits_1_and_prints_no_score_for_an_unscored_arm( + tmp_path: Path, monkeypatch +) -> None: + """The table must not show a score for an arm that was never scored.""" + arms = _skill_outcomes(with_skill=0.0, no_skill=0.0) + for arm in arms: + arm.reward = None + arm.error = _UnscoredRollout.reason + report = _report(arms=arms, value=None) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 1 + assert "produced no verifier reward" in _flat(result.stderr) + assert "no difference" not in _flat(result.output) + payload = json.loads((tmp_path / "out" / "ablation.json").read_text()) + assert [arm["reward"] for arm in payload["arms"]] == [None, None] + + +async def test_run_ablation_records_a_branch_that_never_forked( + tmp_path: Path, monkeypatch +) -> None: + """A stage that was never captured fails before any child attaches: the + report carries the error and every arm reports as skipped rather than 0.0.""" + + class _NoStage(_FakeRollout): + async def branch_at_stage(self, stage, n, *, deltas=None) -> float: + raise LookupError(f"no snapshot recorded at stage {stage!r}") + + _patch_rollout(monkeypatch, _NoStage) + + report = await run_ablation(_request(tmp_path)) + + assert "no snapshot recorded at stage 'env-ready'" in report.error + assert [arm.status for arm in report.arms] == ["skipped", "skipped"] + assert report.arms[0].verdict == "not run — an earlier arm errored" + assert report.has_errors is True + + +async def test_run_ablation_keeps_the_arms_when_the_parent_leg_fails( + tmp_path: Path, monkeypatch +) -> None: + """Attributing a failed run is the point (RFC §1): a parent that died after + the boundary is recorded, and the arms still fork from the snapshot.""" + + class _ParentFails(_FakeRollout): + async def execute(self) -> None: + raise RuntimeError("agent session died") + + built = _patch_rollout(monkeypatch, _ParentFails) + + report = await run_ablation(_request(tmp_path)) + + assert built[0].calls[-1] == "cleanup" + assert "agent session died" in report.parent_error + assert report.parent_reward is None + assert [arm.reward for arm in report.arms] == [1.0, 0.0] + assert report.has_errors is False + + +async def test_run_ablation_raises_when_the_boundary_is_never_reached( + tmp_path: Path, monkeypatch +) -> None: + class _StartFails(_FakeRollout): + async def start(self) -> None: + raise RuntimeError("sandbox never came up") + + _patch_rollout(monkeypatch, _StartFails) + + with pytest.raises(Exception, match="nothing to branch"): + await run_ablation(_request(tmp_path)) + + +# Request-global settings (PR #1046 second review, P2-A): with Gemini and +# --reasoning-effort high the parent built the environment, snapshotted, +# installed and connected the agent, and only then hit the ACP effort +# rejection — then ablation restored a child and failed the same way again. +# Layer 1: what is knowable statically fails before anything is provisioned. +# Layer 2: a parent that failed with a request-global error never has its +# doomed configuration retried in branch children. + + +async def test_an_unsupported_reasoning_effort_fails_before_any_provisioning( + tmp_path: Path, monkeypatch +) -> None: + """gemini declares no ACP effort config option, so --reasoning-effort is + decidable from the request alone: the rejection must land before any + sandbox exists — no Rollout is even constructed. Guards ``fix(ablate): + validate request-global settings before provisioning``.""" + built = _patch_rollout(monkeypatch) + request = AblationRequest( + task_path=_task_dir(tmp_path), + arms=parse_arms("with-skill,no-skill"), + agent="gemini", + model="gemini-3.1-pro-preview", + reasoning_effort="high", + out_dir=tmp_path / "out", + ) + + with pytest.raises(AblationSpecError, match="not supported by agent 'gemini'"): + await run_ablation(request) + + assert built == [] # nothing provisioned, nothing to clean up + + +async def test_a_request_global_parent_failure_is_never_retried_in_children( + tmp_path: Path, monkeypatch +) -> None: + """An agent session rejecting a global option (effort/model) after + provisioning is not task-attributable: every branch child would restore + the snapshot, reinstall the agent, and fail identically. The engine must + skip the fork entirely and the report must say why. Guards ``fix(ablate): + ... never retry them in children``.""" + from benchflow.acp.runtime import ACPRequestGlobalError + + class _EffortRejected(_FakeRollout): + async def connect(self) -> None: + self.calls.append("connect") + raise ACPRequestGlobalError( + "request-global setting rejected: reasoning_effort='high' was " + "requested for agent 'gemini', but that agent does not " + "declare an ACP effort config option" + ) + + built = _patch_rollout(monkeypatch, _EffortRejected) + + report = await run_ablation(_request(tmp_path)) + + assert "does not declare an ACP effort" in report.parent_error + assert not any(c.startswith("branch_at_stage") for c in built[0].calls) + assert built[0].calls[-1] == "cleanup" + assert [arm.status for arm in report.arms] == ["skipped", "skipped"] + assert [arm.reward for arm in report.arms] == [None, None] + assert "request-global" in report.error + assert "not attempted" in report.error + assert report.has_errors is True + + +async def test_a_task_attributable_parent_failure_still_forks_the_arms( + tmp_path: Path, monkeypatch +) -> None: + """The complement of the request-global skip: an ordinary post-boundary + parent failure keeps the documented RFC §1 behavior — the arms still fork + from the recorded snapshot (same shape as + test_run_ablation_keeps_the_arms_when_the_parent_leg_fails, pinned here + against the new skip's classifier over-matching).""" + + class _ExecuteDies(_FakeRollout): + async def execute(self) -> None: + raise RuntimeError("ACP error -32603: agent crashed mid-task") + + built = _patch_rollout(monkeypatch, _ExecuteDies) + + report = await run_ablation(_request(tmp_path)) + + assert "agent crashed mid-task" in report.parent_error + assert any(c.startswith("branch_at_stage") for c in built[0].calls) + assert [arm.reward for arm in report.arms] == [1.0, 0.0] + + +# 4. The CLI surface + + +def test_bad_arm_stage_and_empty_arms_fail_without_a_traceback( + tmp_path: Path, +) -> None: + task = _task_dir(tmp_path) + for flags, expected in ( + (["--arms", "with-skill,turbo"], "unknown arm 'turbo'"), + (["--arms", ""], "--arms is empty"), + (["--arms", "with-skill"], "at least two arms"), + (["--arms", "with-skill,config:"], "carries no patch"), + ( + ["--arms", 'with-skill,config:{"verifier": {"timeout_sec": 1}}'], + "may only patch", + ), + (["--at-stage", "mid-flight"], "unknown --at-stage 'mid-flight'"), + (["--at-stage", "post-research"], "cannot be captured by this command"), + (["--at-stage", "pre-verify"], "needs --at-stage 'env-ready'"), + ): + result = runner.invoke( + app, ["eval", "ablate", "--tasks-dir", str(task), *flags] + ) + assert result.exit_code == 1, result.output + assert expected in _flat(result.stderr) + assert "Traceback (most recent call last)" not in result.output + # Nothing ran: no report directory was created for a rejected request. + assert not (tmp_path / "out").exists() + + +def test_missing_tasks_dir_fails_closed(tmp_path: Path) -> None: + result = runner.invoke( + app, ["eval", "ablate", "--tasks-dir", str(tmp_path / "nope")] + ) + assert result.exit_code == 1, result.output + assert "is not a directory" in _flat(result.stderr) + assert "Traceback (most recent call last)" not in result.output + + +def test_table_renders_arms_rewards_and_the_attribution_line( + tmp_path: Path, monkeypatch +) -> None: + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + seen = _patched_run(monkeypatch, report) + task = _task_dir(tmp_path) + # Rich wraps cells to the terminal width; widen it so the attribution line + # is asserted as one sentence instead of as wrap-dependent fragments. + monkeypatch.setenv("COLUMNS", "300") + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(task), + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 0, result.output + out = result.output + assert "with-skill" in out and "no-skill" in out + assert "1.00" in out and "0.00" in out + assert "pass" in out and "fail" in out + assert "61s" in out and "44s" in out + assert "decides the outcome when applied at env-ready" in _flat(out) + assert (tmp_path / "out" / "ablation.json").is_file() + # The CLI built the request from its flags — defaults included. + assert [arm.name for arm in seen[0].arms] == ["with-skill", "no-skill"] + assert seen[0].stage == "env-ready" + assert seen[0].task_path == task + + +def test_json_output_parses_and_carries_the_arm_provenance( + tmp_path: Path, monkeypatch +) -> None: + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + task = _task_dir(tmp_path) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(task), + "--out-dir", + str(tmp_path / "out"), + "--json", + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert payload["stage"] == "env-ready" + assert payload["task"]["id"] == "demo" + assert payload["parent"]["reward"] == 1.0 + assert payload["value"] == 0.5 + assert [arm["name"] for arm in payload["arms"]] == ["with-skill", "no-skill"] + assert [arm["delta"]["skill_mode"] for arm in payload["arms"]] == [ + "with-skill", + "no-skill", + ] + assert [arm["reward"] for arm in payload["arms"]] == [1.0, 0.0] + assert payload["report_path"] == str(tmp_path / "out" / "ablation.json") + # stdout stays machine-readable: the Rich table never lands on it. + assert "─" not in result.output + + +def test_exit_code_1_when_an_arm_errors(tmp_path: Path, monkeypatch) -> None: + arms = _skill_outcomes(with_skill=1.0, no_skill=0.0) + arms[1].reward = None + arms[1].error = "sandbox 'ModalSandbox' does not implement snapshot/restore" + report = _report(arms=arms, value=None) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + task = _task_dir(tmp_path) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(task), + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 1 + assert "does not implement snapshot/restore" in _flat(result.stderr) + assert "Traceback (most recent call last)" not in result.output + # The report is still written — the arm that did run keeps its reward. + payload = json.loads((tmp_path / "out" / "ablation.json").read_text()) + assert payload["arms"][0]["reward"] == 1.0 + assert payload["arms"][1]["status"] == "error" + + +def test_engine_errors_surface_as_one_line_not_a_traceback( + tmp_path: Path, monkeypatch +) -> None: + async def fake_run_ablation(request: AblationRequest) -> AblationReport: + raise AblationSpecError("bench eval ablate needs an ACP agent") + + monkeypatch.setattr("benchflow.ablate.run_ablation", fake_run_ablation) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--agent", + "oracle", + ], + ) + + assert result.exit_code == 1 + assert "needs an ACP agent" in _flat(result.stderr) + assert "Traceback (most recent call last)" not in result.output + + +# 5. Docs: the flag table in cli.md must match the parser + + +def _doc_section(header: str) -> str: + doc = _CLI_MD.read_text() + index = doc.index(header) + level = len(header) - len(header.lstrip("#")) + match = re.search(rf"\n#{{2,{level}}} ", doc[index + len(header) :]) + end = index + len(header) + match.start() if match else len(doc) + return doc[index:end] + + +def test_ablate_flags_and_cli_md_are_set_equal() -> None: + """Same bidirectional guard tests/test_cli_docs_drift.py holds `bench eval + run` to (#731), extended to this command: a new flag cannot land + undocumented, and a documented flag cannot rot out of the parser.""" + command = cast("click.Group", typer.main.get_command(app)).commands["eval"] + ablate = cast("click.Group", command).commands["ablate"] + cli = { + opt + for param in ablate.params + for opt in getattr(param, "opts", []) + if opt.startswith("--") + } - {"--help"} + doc = set(re.findall(r"`(--[a-z0-9-]+)`", _doc_section(_DOC_HEADER))) + assert cli == doc, ( + "bench eval ablate CLI↔cli.md flag drift:\n" + f" in CLI but UNDOCUMENTED: {sorted(cli - doc)}\n" + f" documented but NOT in CLI: {sorted(doc - cli)}" + ) + + +def test_ablate_is_documented_next_to_the_rfc() -> None: + """The command's docs must point at the design it implements, and the RFC's + validation plan must name the command that produces its evidence.""" + section = _doc_section(_DOC_HEADER) + assert "rollout-branching-rfc" in section + rfc = (_REPO_ROOT / "docs" / "rollout-branching-rfc.md").read_text() + assert "bench eval ablate" in rfc + + +# 6. Per-test (sub-outcome) attribution: a scalar tie cannot hide a difference + + +def _ctrf(tests: dict[str, str]) -> str: + """A CTRF report (pytest-json-ctrf shape) naming ``{node id -> status}``.""" + return json.dumps( + { + "reportFormat": "CTRF", + "results": { + "tool": {"name": "pytest"}, + "tests": [ + {"name": name, "status": status} for name, status in tests.items() + ], + }, + } + ) + + +def _tested_outcomes( + with_skill: dict[str, str] | None, no_skill: dict[str, str] | None +) -> list[ArmOutcome]: + """The measured lake-warming pair: both arms 0.00, per-test maps attached.""" + arms = _skill_outcomes(with_skill=0.0, no_skill=0.0) + arms[0].tests = with_skill + arms[1].tests = no_skill + return arms + + +class _CtrfRollout(_FakeRollout): + """A fork whose children leave a CTRF report where their kind leaves one. + + A **fresh-rollout** child (every child of ``env-ready``) has a run + directory of its own, and that run directory *is* its artifact directory, + so ``/verifier/ctrf.json`` is where its verifier writes. An + **in-place** child (``pre-verify`` / ``post-verify``) has no run directory: + it writes through the parent's bind mounts, and the branch engine archives + what it wrote under ``/mounted/verifier/`` — the layout + :mod:`benchflow.branch_artifacts` owns, which is why both paths are asked + for here rather than spelled out. + + ``per_child`` is indexed by fork order; ``None`` writes no report at all + (the verifier emitted no per-test data). + """ + + rewards: ClassVar[list[Any]] = [0.0, 0.0] + per_child: ClassVar[list[dict[str, str] | None]] = [ + {"::T::test_trend_result": "passed", "::T::test_dominant_factor": "failed"}, + {"::T::test_trend_result": "failed", "::T::test_dominant_factor": "passed"}, + ] + #: ``True`` runs the children the way a ``pre-verify`` fork does — in + #: place, with their output preserved under ``mounted/``. + in_place: ClassVar[bool] = False + + def _child_verifier_dir(self, child) -> Path: + from benchflow.branch_artifacts import child_mount_dir + from benchflow.branch_lineage import branch_child_dir + + if self.in_place: + return ( + child_mount_dir(self._rollout_dir, child.parent.id, child.id) + / "verifier" + ) + return ( + branch_child_dir(self._rollout_dir, child.parent.id, child.id) / "verifier" + ) + + async def branch_at_stage(self, stage, n, *, deltas=None): + value = await super().branch_at_stage(stage, n, deltas=deltas) + children = [node for node in self.tree.nodes() if "delta" in node.state] + for child, tests in zip(children, self.per_child, strict=True): + if tests is None: + continue + verifier = self._child_verifier_dir(child) + verifier.mkdir(parents=True, exist_ok=True) + (verifier / "ctrf.json").write_text(_ctrf(tests), encoding="utf-8") + return value + + +class _InPlaceCtrfRollout(_CtrfRollout): + """The ``pre-verify`` fork: children run in place, artifacts under ``mounted/``.""" + + in_place: ClassVar[bool] = True + + +async def test_per_test_outcomes_are_mined_from_each_child_verifier_report( + tmp_path: Path, monkeypatch +) -> None: + """Guards "feat(ablate): per-test attribution so a scalar tie cannot hide a + behavioral difference": every arm carries the outcome map its own verifier + reported, read through the CLI's existing CTRF parser (no second parser, + and the ``ClassName::test`` node ids are displayed by test name).""" + _patch_rollout(monkeypatch, _CtrfRollout) + + report = await run_ablation(_request(tmp_path)) + + assert [arm.reward for arm in report.arms] == [0.0, 0.0] + assert report.arms[0].tests == { + "test_dominant_factor": "failed", + "test_trend_result": "passed", + } + assert report.arms[1].tests == { + "test_dominant_factor": "passed", + "test_trend_result": "failed", + } + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert [arm["tests"] for arm in payload["arms"]] == [ + {"test_dominant_factor": "failed", "test_trend_result": "passed"}, + {"test_dominant_factor": "passed", "test_trend_result": "failed"}, + ] + + +def _in_place_request(tmp_path: Path) -> AblationRequest: + """A ``pre-verify`` ablation: two injection arms, children run in place.""" + first = tmp_path / "plan-a.md" + first.write_text("Warm the lake first.\n", encoding="utf-8") + second = tmp_path / "plan-b.md" + second.write_text("Cool the lake first.\n", encoding="utf-8") + return AblationRequest( + task_path=_task_dir(tmp_path), + arms=parse_arms(f"inject:{first},inject:{second}"), + agent="claude-agent-acp", + model="claude-sonnet", + stage="pre-verify", + out_dir=tmp_path / "out", + ) + + +async def test_per_test_outcomes_are_mined_from_an_in_place_childs_mounted_output( + tmp_path: Path, monkeypatch +) -> None: + """An in-place branch child keeps its per-test attribution. + + Guards the fix from "fix(ablate): read per-test outcomes from preserved + in-place child artifacts", reproduced live by @Galius5136 at ``--at-stage + pre-verify``. A child that runs in place has no rollout directory of its + own — its verifier output is the ``mounted/`` archive of what it wrote to + the parent's shared bind mounts (the artifact isolation from "fix(branch): + branch children no longer clobber the parent's artifacts"). The reader + looked only at ``/verifier/ctrf.json``, so the report fell back to + scalar-only attribution while the per-test data sat one directory down. + """ + _patch_rollout(monkeypatch, _InPlaceCtrfRollout) + + report = await run_ablation(_in_place_request(tmp_path)) + + assert [arm.reward for arm in report.arms] == [0.0, 0.0] + assert report.arms[0].tests == { + "test_dominant_factor": "failed", + "test_trend_result": "passed", + } + assert report.arms[1].tests == { + "test_dominant_factor": "passed", + "test_trend_result": "failed", + } + section = sub_test_attribution(report.arms) + assert section["arms_without_tests"] == [] + assert section["summary"] == ( + "scalar rewards tie, but 2 sub-test outcome(s) differ: " + "test_dominant_factor, test_trend_result" + ) + + +async def test_a_fresh_childs_own_verifier_output_wins_over_its_mounted_copy( + tmp_path: Path, monkeypatch +) -> None: + """Both roots can exist; the child's own run directory is authoritative. + + Guards the ordering in "fix(ablate): read per-test outcomes from preserved + in-place child artifacts". A fresh-rollout child writes through the + restored parent mounts *and* downloads its artifacts into its own run + directory, so the ``mounted/`` archive can hold a second, staler copy — + reading that one instead would attribute the wrong outcomes to the arm. + """ + from benchflow.branch_artifacts import child_mount_dir + + class _BothRoots(_CtrfRollout): + async def branch_at_stage(self, stage, n, *, deltas=None): + value = await super().branch_at_stage(stage, n, deltas=deltas) + for child in (node for node in self.tree.nodes() if "delta" in node.state): + mounted = ( + child_mount_dir(self._rollout_dir, child.parent.id, child.id) + / "verifier" + ) + mounted.mkdir(parents=True, exist_ok=True) + (mounted / "ctrf.json").write_text( + _ctrf({"::T::test_stale": "failed"}), encoding="utf-8" + ) + return value + + _patch_rollout(monkeypatch, _BothRoots) + + report = await run_ablation(_request(tmp_path)) + + assert [sorted(arm.tests or {}) for arm in report.arms] == [ + ["test_dominant_factor", "test_trend_result"], + ["test_dominant_factor", "test_trend_result"], + ] + + +def test_only_the_tests_whose_outcome_differs_are_attributed() -> None: + """Tying tests are noise for attribution: they are counted, never listed. + + Guards the differences table of "feat(ablate): per-test attribution …" — + a table that reprinted every test would bury the two rows that carry the + signal. + """ + arms = _tested_outcomes( + {"test_a": "passed", "test_b": "failed", "test_same": "passed"}, + {"test_a": "failed", "test_b": "passed", "test_same": "passed"}, + ) + + differences = differing_tests(arms) + section = sub_test_attribution(arms) + + assert differences == [ + {"test": "test_a", "outcomes": {"no-skill": "failed", "with-skill": "passed"}}, + {"test": "test_b", "outcomes": {"no-skill": "passed", "with-skill": "failed"}}, + ] + assert section["tying_tests"] == ["test_same"] + assert section["arms_with_tests"] == ["with-skill", "no-skill"] + assert section["arms_without_tests"] == [] + assert section["scalar_tie"] is True + + +def test_a_test_reported_for_one_arm_only_is_a_difference_not_a_failure() -> None: + """A test the other arm never reported is recorded as ``None``. + + Guards "feat(ablate): per-test attribution …" against inventing rows: the + honest reading of a missing entry is "not observed for that arm", never + "failed there". + """ + arms = _tested_outcomes( + {"test_a": "passed"}, {"test_a": "passed", "test_b": "failed"} + ) + + assert differing_tests(arms) == [ + {"test": "test_b", "outcomes": {"no-skill": "failed", "with-skill": None}} + ] + + +def test_scalar_tie_with_differing_sub_tests_says_so_in_the_verdict() -> None: + """The exact regression: two arms at 0.00 whose sub-tests flip oppositely. + + Guards "feat(ablate): per-test attribution so a scalar tie cannot hide a + behavioral difference". The measured lake-warming ablation printed "no + difference in this comparison" while ``test_trend_result`` flipped with the + skill pack and ``test_dominant_factor`` flipped against it. That wording + must be unreachable whenever a sub-test difference was observed. + """ + arms = _tested_outcomes( + {"test_trend_result": "passed", "test_dominant_factor": "failed"}, + {"test_trend_result": "failed", "test_dominant_factor": "passed"}, + ) + + attribute(arms, parent_reward=0.0, stage="env-ready") + + assert arms[0].verdict == ( + "matches no-skill at env-ready (both 0.00) — scalar rewards tie, but 2 " + "sub-test outcome(s) differ: test_dominant_factor, test_trend_result" + ) + assert all("no difference in this comparison" not in a.verdict for a in arms) + assert sub_test_attribution(arms)["summary"] == ( + "scalar rewards tie, but 2 sub-test outcome(s) differ: " + "test_dominant_factor, test_trend_result" + ) + + +def test_a_genuine_tie_keeps_todays_no_difference_wording() -> None: + """Sub-tests that were observed *and* tie are not a difference. + + The other half of "feat(ablate): per-test attribution …": the qualification + fires on measured disagreement only, so an ablation where nothing moved + still reads as it always has. + """ + arms = _tested_outcomes({"test_a": "failed"}, {"test_a": "failed"}) + + attribute(arms, parent_reward=0.0, stage="env-ready") + + assert arms[0].verdict == ( + "matches no-skill at env-ready (both 0.00) — no difference in this comparison" + ) + assert sub_test_attribution(arms)["summary"] == ( + "2 arms reported per-test outcomes and all 1 tie — no sub-test " + "difference in this comparison" + ) + + +async def test_a_child_without_per_test_data_degrades_to_scalar_only( + tmp_path: Path, monkeypatch +) -> None: + """A verifier that emits no CTRF report yields no rows, and says so. + + Guards the degradation path of "feat(ablate): per-test attribution …": + ``tests`` stays ``None`` (not ``{}``, not a fabricated row), the section + names the arms it could not read, and the scalar verdict is left alone — + nothing was measured to qualify it with. + """ + monkeypatch.setattr(_CtrfRollout, "per_child", [{"::T::test_a": "passed"}, None]) + _patch_rollout(monkeypatch, _CtrfRollout) + + report = await run_ablation(_request(tmp_path)) + + assert report.arms[0].tests == {"test_a": "passed"} + assert report.arms[1].tests is None + section = sub_test_attribution(report.arms) + assert section["arms_with_tests"] == ["with-skill"] + assert section["arms_without_tests"] == ["no-skill"] + assert section["differing_tests"] == [] + assert section["summary"] == ( + "scalar-only attribution — 1 of 2 arms reported per-test outcomes, so " + "no sub-test comparison was made" + ) + assert "no difference in this comparison" in report.arms[0].verdict + + +def test_test_attribution_is_deterministic_and_carries_no_wall_clock( + tmp_path: Path, +) -> None: + """Same input, byte-identical report — the determinism guarantee of + "feat(ablate): per-test attribution …" extended to the new section: test + names sort, arm order follows the request, and nothing is stamped.""" + arms = _tested_outcomes( + {"test_z": "failed", "test_a": "passed"}, + {"test_a": "failed", "test_z": "passed"}, + ) + report = _report(arms=arms, parent_reward=0.0, value=0.0) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + + first = write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + second = write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + + assert first == second + payload = json.loads(first) + assert [ + entry["test"] for entry in payload["test_attribution"]["differing_tests"] + ] == [ + "test_a", + "test_z", + ] + assert list(payload["arms"][0]["tests"]) == ["test_a", "test_z"] + + +def test_table_lists_the_differing_sub_tests_and_omits_the_tying_ones( + tmp_path: Path, monkeypatch +) -> None: + """The console half of "feat(ablate): per-test attribution …". + + A reader of the table must be able to see the behavioral difference the + two 0.00s hide, and must not have to scan tying rows to find it. + """ + arms = _tested_outcomes( + { + "test_trend_result": "passed", + "test_dominant_factor": "failed", + "test_boring": "passed", + }, + { + "test_trend_result": "failed", + "test_dominant_factor": "passed", + "test_boring": "passed", + }, + ) + report = _report(arms=arms, parent_reward=0.0, value=0.0) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + monkeypatch.setenv("COLUMNS", "300") + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 0, result.output + out = _flat(result.output) + assert "Sub-test outcomes that differ (2)" in out + assert "test_trend_result" in out and "test_dominant_factor" in out + assert "test_boring" not in out + assert "1 test(s) tie across the arms and are omitted." in out + assert "scalar rewards tie, but 2 sub-test outcome(s) differ" in out + assert "no difference in this comparison" not in out + + +def test_json_output_carries_the_per_test_maps_and_the_difference_section( + tmp_path: Path, monkeypatch +) -> None: + """``--json`` stays machine-readable and complete for "feat(ablate): + per-test attribution …": the per-arm maps and the differences section both + round-trip through stdout.""" + arms = _tested_outcomes( + {"test_trend_result": "passed", "test_dominant_factor": "failed"}, + {"test_trend_result": "failed", "test_dominant_factor": "passed"}, + ) + report = _report(arms=arms, parent_reward=0.0, value=0.0) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--out-dir", + str(tmp_path / "out"), + "--json", + ], + ) + + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert payload["arms"][0]["tests"] == { + "test_dominant_factor": "failed", + "test_trend_result": "passed", + } + section = payload["test_attribution"] + assert section["scalar_tie"] is True + assert [entry["test"] for entry in section["differing_tests"]] == [ + "test_dominant_factor", + "test_trend_result", + ] + assert section["differing_tests"][0]["outcomes"] == { + "no-skill": "passed", + "with-skill": "failed", + } + assert "─" not in result.output + + +# 7. The task's declared environment reaches the parent and every arm + + +_ABLATE_MANIFEST = """\ +environment: + name: ablate-declared-env + image: example/ablate-declared:latest +""" + + +def _manifest_task_dir( + tmp_path: Path, + *, + declares: str = "environment.yaml", + write_manifest: bool = True, +) -> Path: + """A task whose ``task.md`` declares ``benchflow.environment.manifest``.""" + task = _task_dir(tmp_path) + (task / "task.md").write_text( + "---\n" + 'schema_version: "1.3"\n' + "task:\n" + " name: benchflow/ablate-manifest-demo\n" + " description: A task that declares its own environment\n" + "benchflow:\n" + " environment:\n" + f" manifest: {declares}\n" + "---\n" + "Solve it.\n", + encoding="utf-8", + ) + if write_manifest: + (task / "environment.yaml").write_text(_ABLATE_MANIFEST, encoding="utf-8") + return task + + +async def test_the_task_declared_environment_reaches_the_parent_and_the_arms( + tmp_path: Path, monkeypatch +) -> None: + """A manifest-backed task is ablated in the world it declares. + + Guards "fix(ablate): resolve a task-declared environment manifest for + parent and arms". ``bench eval`` resolves ``benchflow.environment.manifest`` + per task before building the rollout config; ``bench eval ablate`` left + ``environment_manifest=None``, so the parent — and therefore every arm + forked from its snapshot — ran without the task's required image, + services, provisioning and readiness gates, and the report compared a + different environment than a normal evaluation of the same task. + + Both halves are asserted: the parent config the command builds, and the + child config the branch engine derives from it for a fresh ``env-ready`` + child (:func:`benchflow.branch_skill.make_fresh_child_runner`, the + production path for every arm at that boundary). + """ + from types import SimpleNamespace + + import benchflow.branch_skill as branch_skill + from benchflow.branch_skill import make_fresh_child_runner + + task = _manifest_task_dir(tmp_path) + built = _patch_rollout(monkeypatch) + request = AblationRequest( + task_path=task, + arms=parse_arms("with-skill,no-skill"), + agent="claude-agent-acp", + model="claude-sonnet", + out_dir=tmp_path / "out", + ) + + await run_ablation(request) + + manifest = built[0].config.environment_manifest + assert manifest is not None + assert (manifest.name, manifest.image) == ( + "ablate-declared-env", + "example/ablate-declared:latest", + ) + + child_configs: list[Any] = [] + + async def _capture_fresh_child(rollout, config, *, prompts=None): + child_configs.append(config) + return 1.0 + + monkeypatch.setattr(branch_skill, "run_fresh_child", _capture_fresh_child) + tree = RolloutTree() + child = tree.attach(tree.root) + for arm in request.arms: + run_child = make_fresh_child_runner( + cast(Any, SimpleNamespace(_config=built[0].config)), + delta=arm.delta, + parent=tree.root, + branch_stage="env-ready", + run_dir=built[0]._rollout_dir, + ) + await run_child(child) + + assert [cfg.environment_manifest for cfg in child_configs] == [manifest, manifest] + + +async def test_a_task_without_a_declared_environment_still_binds_nothing( + tmp_path: Path, monkeypatch +) -> None: + """The control for the fix: no declaration, no manifest — nothing invented. + + Guards "fix(ablate): resolve a task-declared environment manifest for + parent and arms" from over-reaching: a task that declares no environment + must keep running on the plain sandbox, exactly as before. + """ + built = _patch_rollout(monkeypatch) + + await run_ablation(_request(tmp_path)) + + assert built[0].config.environment_manifest is None + + +async def test_an_unresolvable_declared_environment_fails_before_the_parent_runs( + tmp_path: Path, monkeypatch +) -> None: + """A declaration that cannot be resolved is fatal, not silently dropped. + + Guards "fix(ablate): resolve a task-declared environment manifest for + parent and arms". Degrading to ``None`` would run the whole experiment — + parent and arms — in a world the task says is wrong and report it as if it + were the right one, so the request fails closed before the parent run costs + anything. + """ + task = _manifest_task_dir(tmp_path, declares="missing.yaml", write_manifest=False) + built = _patch_rollout(monkeypatch) + request = AblationRequest( + task_path=task, + arms=parse_arms("with-skill,no-skill"), + agent="claude-agent-acp", + model="claude-sonnet", + out_dir=tmp_path / "out", + ) + + with pytest.raises(AblationSpecError, match="environment manifest"): + await run_ablation(request) + + assert built == [] + + +# 7b. The parent runs the canonical evaluation configuration + + +async def test_the_parent_config_is_the_canonical_plan_with_the_ablation_overlaid( + tmp_path: Path, monkeypatch +) -> None: + """The parent's config carries what a plain eval run of the task would. + + Guards "fix(ablate): resolve the canonical eval plan and overlay only the + ablation axis" (the PR #1046 review finding): run_ablation used to + hand-roll a reduced RolloutConfig, so a real E2E parent published + ``task_digest: null`` and ``reasoning_effort: null`` — the run lost its + attribution to exact task content and its effort control. The request now + resolves through ``build_eval_plan`` + ``task_rollout_config`` (the same + two stages as ``bench eval run``), and only the ablation-owned fields are + overlaid on top. + """ + import dataclasses + + from benchflow._utils.benchmark_repos import task_source_provenance + from benchflow._utils.task_authoring import task_digest + from benchflow.usage_tracking import UsageTrackingConfig + + built = _patch_rollout(monkeypatch) + request = dataclasses.replace(_request(tmp_path), reasoning_effort="MAX") + + await run_ablation(request) + + config = built[0].config + # Canonical controls and provenance, from the task fixture through the plan. + assert config.task_digest == task_digest(request.task_path) + assert config.task_digest.startswith("sha256:") + assert config.reasoning_effort == "max" # normalized, not passed through raw + assert config.primary_reasoning_effort == "max" + assert config.model == "claude-sonnet" + assert config.source_provenance == task_source_provenance( + None, Path(request.task_path) + ) + assert isinstance(config.usage_tracking, UsageTrackingConfig) + assert config.agent_idle_timeout == 600 # the plan's normalized default + # The overlay: only what the ablation owns. + assert config.skill_mode == "no-skill" + assert config.snapshot_stages == frozenset({"env-ready"}) + assert config.snapshot_layers == frozenset({"sandbox"}) + assert config.job_name == "ablation" + assert config.rollout_name == request.task_path.name + + +async def test_task_digest_and_reasoning_effort_flow_into_every_arms_config( + tmp_path: Path, monkeypatch +) -> None: + """The child configs inherit the canonical fields, not nulls. + + The second half of "fix(ablate): resolve the canonical eval plan and + overlay only the ablation axis": a fresh ``env-ready`` child derives its + config from the parent's (``child_skill_config`` is a + ``dataclasses.replace``), so the E2E nulls in the *child* configs were the + parent's nulls showing through. Asserted on the production path every arm + takes at that boundary (:func:`benchflow.branch_skill.make_fresh_child_runner`). + """ + import dataclasses + from types import SimpleNamespace + + import benchflow.branch_skill as branch_skill + from benchflow._utils.task_authoring import task_digest + from benchflow.branch_skill import make_fresh_child_runner + + built = _patch_rollout(monkeypatch) + request = dataclasses.replace(_request(tmp_path), reasoning_effort="high") + + await run_ablation(request) + + child_configs: list[Any] = [] + + async def _capture_fresh_child(rollout, config, *, prompts=None): + child_configs.append(config) + return 1.0 + + monkeypatch.setattr(branch_skill, "run_fresh_child", _capture_fresh_child) + tree = RolloutTree() + child = tree.attach(tree.root) + for arm in request.arms: + run_child = make_fresh_child_runner( + cast(Any, SimpleNamespace(_config=built[0].config)), + delta=arm.delta, + parent=tree.root, + branch_stage="env-ready", + run_dir=built[0]._rollout_dir, + ) + await run_child(child) + + digest = task_digest(request.task_path) + assert [cfg.task_digest for cfg in child_configs] == [digest, digest] + assert [cfg.reasoning_effort for cfg in child_configs] == ["high", "high"] + + +async def test_a_bad_reasoning_effort_or_sandbox_fails_before_the_parent_runs( + tmp_path: Path, monkeypatch +) -> None: + """Plan validation is the ablation's validation, fail-closed and free. + + Resolving the canonical plan buys its gates too: a typo'd effort or an + unknown sandbox dies as a spec error before any rollout is built, with the + same message ``bench eval run`` prints — not deep in the parent run. + """ + import dataclasses + + built = _patch_rollout(monkeypatch) + request = _request(tmp_path) + + with pytest.raises(AblationSpecError, match="reasoning_effort"): + await run_ablation(dataclasses.replace(request, reasoning_effort="turbo")) + with pytest.raises(AblationSpecError, match="Invalid --sandbox"): + await run_ablation(dataclasses.replace(request, sandbox="warpdrive")) + + assert built == [] # nothing ran, nothing was paid for + + +def test_cli_reasoning_effort_reaches_the_request(tmp_path: Path, monkeypatch) -> None: + """--reasoning-effort threads to AblationRequest — same flag as eval run.""" + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + seen = _patched_run(monkeypatch, report) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--reasoning-effort", + "max", + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 0, result.output + assert seen[0].reasoning_effort == "max" + + +# 8. Ablation self-description: the bound environment and the stage snapshot + + +def _request_with_manifest( + tmp_path: Path, + task: Path, + *, + arms: str = "with-skill,no-skill", + environment_manifest: Path | str | None = None, +) -> AblationRequest: + return AblationRequest( + task_path=task, + arms=parse_arms(arms), + agent="claude-agent-acp", + model="claude-sonnet", + out_dir=tmp_path / "out", + environment_manifest=environment_manifest, + ) + + +async def test_an_explicit_environment_manifest_beats_the_task_declared_one( + tmp_path: Path, monkeypatch +) -> None: + """--environment-manifest wins over task.md, same precedence as the run path. + + Guards "feat(ablate): bind an explicit environment manifest and stamp the + bound environment". ``bench eval run`` consults + ``manifest_from_task_document`` only when no explicit manifest was bound + (``benchflow.evaluation``); the ablation must mirror that, or the same + flag would mean two different worlds on the two commands. + """ + from benchflow._utils.content_address import sha256_prefixed + + task = _manifest_task_dir(tmp_path) # declares environment.yaml + override = tmp_path / "override.yaml" + override.write_text( + "environment:\n name: ablate-flag-env\n image: example/flag:latest\n", + encoding="utf-8", + ) + built = _patch_rollout(monkeypatch) + + report = await run_ablation( + _request_with_manifest(tmp_path, task, environment_manifest=override) + ) + + manifest = built[0].config.environment_manifest + assert (manifest.name, manifest.image) == ("ablate-flag-env", "example/flag:latest") + assert report.environment == { + "name": "ablate-flag-env", + "ref": str(override), + "env_hash": sha256_prefixed(override.read_bytes()), + "image": "example/flag:latest", + "base_image": None, + } + + +async def test_the_task_declared_environment_is_stamped_deterministically( + tmp_path: Path, monkeypatch +) -> None: + """The report says which world the arms compared in, machine-independently. + + Guards "feat(ablate): bind an explicit environment manifest and stamp the + bound environment" — the reviewer's open question. The stamp's ``ref`` is + the ``task.md`` declaration verbatim (never the resolved absolute path, + which would make ablation.json differ per machine) and ``env_hash`` is the + manifest's sha256 content address, the same identity the registry uses. + """ + from benchflow._utils.content_address import sha256_prefixed + + task = _manifest_task_dir(tmp_path) + _patch_rollout(monkeypatch) + + report = await run_ablation(_request_with_manifest(tmp_path, task)) + + assert report.environment == { + "name": "ablate-declared-env", + "ref": "environment.yaml", + "env_hash": sha256_prefixed((task / "environment.yaml").read_bytes()), + "image": "example/ablate-declared:latest", + "base_image": None, + } + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert payload["environment"]["ref"] == "environment.yaml" + assert str(tmp_path) not in json.dumps(payload["environment"]) + + +async def test_a_manifest_less_ablation_stamps_no_environment( + tmp_path: Path, monkeypatch +) -> None: + """No bound world, no stamp — nothing invented.""" + _patch_rollout(monkeypatch) + + report = await run_ablation(_request(tmp_path)) + + assert report.environment is None + assert all(arm.environment is None for arm in report.arms) + + +async def test_an_unresolvable_explicit_manifest_fails_before_the_parent_runs( + tmp_path: Path, monkeypatch +) -> None: + """The flag fails closed exactly like a broken task declaration.""" + built = _patch_rollout(monkeypatch) + request = _request_with_manifest( + tmp_path, + _task_dir(tmp_path), + environment_manifest=tmp_path / "nope.yaml", + ) + + with pytest.raises(AblationSpecError, match="--environment-manifest"): + await run_ablation(request) + + assert built == [] + + +def test_the_environment_manifest_flag_reaches_the_request( + tmp_path: Path, monkeypatch +) -> None: + """The CLI threads --environment-manifest through to the request — the + same flag name and value forms as ``bench eval run``.""" + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + seen = _patched_run(monkeypatch, report) + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--environment-manifest", + "env0@prod", + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 0, result.output + assert seen[0].environment_manifest == Path("env0@prod") + + +_ABLATE_PROD_MANIFEST = """\ +environment: + name: ablate-prod + base_image: example/claws:latest + owns_lifecycle: false + services: + - name: gmail + command: serve-gmail + port: 9001 + - name: gcal + command: serve-gcal + port: 9002 +""" + +_ABLATE_OUTAGE_MANIFEST = """\ +environment: + name: ablate-outage + base_image: example/claws:latest + owns_lifecycle: false + services: + - name: gmail + command: serve-gmail + port: 9001 +""" + + +async def test_an_env_arms_swapped_environment_is_stamped_on_its_row( + tmp_path: Path, monkeypatch +) -> None: + """An ``env:`` arm's row names the world its delta swapped in. + + Guards "feat(ablate): bind an explicit environment manifest and stamp the + bound environment": the top-level stamp is the parent's world, and the one + arm that ran a *different* manifest carries its own stamp — name, the ref + exactly as the arm spec wrote it, and the child manifest's content hash — + while every inheriting arm carries none. + """ + from benchflow._utils.content_address import sha256_prefixed + + task = _task_dir(tmp_path) + prod = tmp_path / "prod.yaml" + prod.write_text(_ABLATE_PROD_MANIFEST, encoding="utf-8") + outage = tmp_path / "outage.yaml" + outage.write_text(_ABLATE_OUTAGE_MANIFEST, encoding="utf-8") + (task / "task.md").write_text( + "---\n" + 'schema_version: "1.3"\n' + "task:\n" + " name: benchflow/ablate-env-arm-demo\n" + " description: A task with a swappable service set\n" + "benchflow:\n" + " environment:\n" + f" manifest: {prod}\n" + "---\n" + "Solve it.\n", + encoding="utf-8", + ) + _patch_rollout(monkeypatch) + + report = await run_ablation( + _request_with_manifest(tmp_path, task, arms=f"no-skill,env:{outage}") + ) + + assert report.environment is not None + assert report.environment["name"] == "ablate-prod" + assert report.arms[0].environment is None + assert report.arms[1].environment == { + "name": "ablate-outage", + "ref": str(outage), + "env_hash": sha256_prefixed(outage.read_bytes()), + "image": None, + "base_image": "example/claws:latest", + } + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert payload["arms"][1]["environment"]["name"] == "ablate-outage" + + +# 9. Stage UX: the snapshot refs travel with the report, and the post-research +# rejection teaches the SDK path + + +class _StageRefRollout(_FakeRollout): + """A fork that records its stage registry the way the real engine does.""" + + async def branch_at_stage(self, stage, n, *, deltas=None): + from benchflow.branch import StageSnapshot + from benchflow.sandbox.protocol import SandboxImage + + self._stage_snapshots = { + stage: StageSnapshot( + environment_ref=None, + sandbox_ref=SandboxImage(provider="fake", ref="bf-snap-root-1"), + stage=stage, + ) + } + return await super().branch_at_stage(stage, n, deltas=deltas) + + +async def test_the_branched_stages_snapshot_refs_are_stamped_into_the_report( + tmp_path: Path, monkeypatch +) -> None: + """ablation.json records the stage's roll-back handles. + + Guards "feat(ablate): bind an explicit environment manifest and stamp the + bound environment" (the stage-UX half): the committed sandbox image ref of + the branched boundary is what a reader needs to restore that world and + re-branch it by hand later, so it travels in the report — same refs as the + parent run's ``stage_snapshots.json``, no second source of truth. + """ + _patch_rollout(monkeypatch, _StageRefRollout) + + report = await run_ablation(_request(tmp_path)) + + assert report.stage_snapshot == { + "environment_ref": None, + "sandbox_ref": "bf-snap-root-1", + "layers": ["sandbox"], + # The fake registry records no capture-time exchange index, so the + # report carries the honest null ("feat(branch): record stage markers + # with trajectory exchange indices"). + "exchanges_completed": None, + # Without --keep-snapshots the committed image dies with the run's + # cleanup, and the report says so (see the retention tests below). + "ephemeral": True, + "exported": None, + } + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert payload["stage_snapshot"]["sandbox_ref"] == "bf-snap-root-1" + + +# 10. Snapshot retention (RFC §3.6): --keep-snapshots exports, and a handle +# cleanup destroyed is recorded as ephemeral, never as restorable + + +class _ExportingStageRefRollout(_StageRefRollout): + """A stage-ref rollout whose sandbox can docker-save its snapshot image.""" + + def __init__(self, config) -> None: + super().__init__(config) + rollout = self + + class _Sandbox: + async def export_image(self, ref: str, target_path) -> None: + # Recorded in the rollout's call log so the test can prove the + # export happened BEFORE cleanup destroyed the image. + rollout.calls.append(f"export:{ref}") + Path(target_path).write_bytes(b"fake-docker-save-tar") + + self.env = _Sandbox() + + +async def test_keep_snapshots_exports_the_tar_before_cleanup_and_records_it( + tmp_path: Path, monkeypatch +) -> None: + """Guards "feat(ablate): --keep-snapshots export; ephemeral handles + recorded truthfully" (the durable half): with the flag, the branched + stage's sandbox image is docker-saved to /snapshots/.tar + *before* cleanup destroys it, and ablation.json records the tar's path + and content sha256 with ephemeral: false.""" + import dataclasses + import hashlib + + built = _patch_rollout(monkeypatch, _ExportingStageRefRollout) + request = dataclasses.replace(_request(tmp_path), keep_snapshots=True) + + report = await run_ablation(request) + + # The export ran between the fork and cleanup — the one window in which + # the committed image still exists. + calls = built[0].calls + assert calls.index("export:bf-snap-root-1") < calls.index("cleanup") + tar_path = Path(request.out_dir) / "snapshots" / "bf-snap-root-1.tar" + assert tar_path.read_bytes() == b"fake-docker-save-tar" + assert report.stage_snapshot["ephemeral"] is False + assert report.stage_snapshot["exported"] == { + "path": str(tar_path), + "sha256": "sha256:" + hashlib.sha256(b"fake-docker-save-tar").hexdigest(), + # The fake's bytes are not a real `docker save` tar, so the image id + # is the honest null ("feat(rollout): stage snapshots record their + # lifetime; --keep-snapshots on bench eval run with a tested import + # path" — the id is read from the tar's manifest for verification at + # import time, never guessed). + "image_id": None, + } + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert payload["stage_snapshot"]["exported"]["path"] == str(tar_path) + + +async def test_without_keep_snapshots_the_handle_is_recorded_ephemeral( + tmp_path: Path, monkeypatch +) -> None: + """Guards "feat(ablate): --keep-snapshots export; ephemeral handles + recorded truthfully" (the truthful half): the real E2E report recorded + bf-snap-…, and docker image inspect immediately confirmed the image no + longer existed — cleanup runs before serialization. The default report + now says the handle is ephemeral so a reader knows it does not resolve.""" + _patch_rollout(monkeypatch, _StageRefRollout) + + report = await run_ablation(_request(tmp_path)) + + assert report.stage_snapshot["ephemeral"] is True + assert report.stage_snapshot["exported"] is None + assert "export_error" not in report.stage_snapshot + payload = json.loads( + write_ablation_report(report, tmp_path / "out").read_text(encoding="utf-8") + ) + assert payload["stage_snapshot"]["ephemeral"] is True + assert payload["stage_snapshot"]["exported"] is None + + +async def test_a_failed_export_is_recorded_and_does_not_cost_the_rewards( + tmp_path: Path, monkeypatch +) -> None: + """The failure path of "feat(ablate): --keep-snapshots export; ephemeral + handles recorded truthfully": a backend that cannot export (here: no + sandbox attached) records export_error and keeps the handle ephemeral — + the arms' rewards and the report survive.""" + import dataclasses + + _patch_rollout(monkeypatch, _StageRefRollout) # no .env on this fake + request = dataclasses.replace(_request(tmp_path), keep_snapshots=True) + + report = await run_ablation(request) + + snap = report.stage_snapshot + assert snap["ephemeral"] is True + assert snap["exported"] is None + assert "--keep-snapshots" in snap["export_error"] + assert [arm.reward for arm in report.arms] == [1.0, 0.0] + assert not (Path(request.out_dir) / "snapshots").exists() + + +def test_cli_keep_snapshots_reaches_the_request(tmp_path: Path, monkeypatch) -> None: + """The flag half of "feat(ablate): --keep-snapshots export; ephemeral + handles recorded truthfully": --keep-snapshots reaches + AblationRequest.keep_snapshots, and its absence means False.""" + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + seen = _patched_run(monkeypatch, report) + + base = [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path, "with-flag")), + "--out-dir", + str(tmp_path / "out-a"), + ] + result = runner.invoke(app, [*base, "--keep-snapshots"]) + assert result.exit_code == 0, result.output + assert seen[0].keep_snapshots is True + + base[3] = str(_task_dir(tmp_path, "without-flag")) + result = runner.invoke(app, base) + assert result.exit_code == 0, result.output + assert seen[1].keep_snapshots is False + + +def test_the_table_prints_the_environment_and_the_snapshot_refs( + tmp_path: Path, monkeypatch +) -> None: + """The self-description reaches the human surface, not only the JSON.""" + report = _report(arms=_skill_outcomes(with_skill=1.0, no_skill=0.0)) + report.environment = { + "name": "ablate-prod", + "ref": "env0@prod", + "env_hash": "sha256:abc123", + "image": None, + "base_image": "example/claws:latest", + } + report.stage_snapshot = { + "environment_ref": None, + "sandbox_ref": "bf-snap-root-1", + "layers": ["sandbox"], + } + attribute(report.arms, parent_reward=report.parent_reward, stage=report.stage) + _patched_run(monkeypatch, report) + monkeypatch.setenv("COLUMNS", "300") + + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--out-dir", + str(tmp_path / "out"), + ], + ) + + assert result.exit_code == 0, result.output + out = _flat(result.output) + assert "Environment: ablate-prod (env0@prod, sha256:abc123)" in out + assert "Stage snapshot: sandbox=bf-snap-root-1" in out + + +def test_post_research_rejection_names_the_sdk_branching_path() -> None: + """The rejection teaches how post-research branching IS done. + + Guards "feat(ablate): bind an explicit environment manifest and stamp the + bound environment" (the stage-UX half): the CLI keeps refusing the stage + it cannot capture, but the message now states the working recipe — + ``mark_stage`` at the cut point, then ``branch_at_stage`` — instead of a + bare pointer at "the Python API". + """ + arms = parse_arms("with-skill,no-skill") + with pytest.raises(AblationSpecError) as excinfo: + validate_arms_for_stage(arms, "post-research") + message = str(excinfo.value) + assert "cannot be captured by this command" in message + assert "rollout.mark_stage('post-research')" in message + assert "rollout.branch_at_stage('post-research'" in message + + +# 11. post-research capture via a research-end marker file +# ("feat(ablate): post-research capture via a research-end marker file") + + +def _inject_arms(tmp_path: Path) -> str: + """Two injection arms — the only arm kind executable at post-research.""" + plan = tmp_path / "oracle-plan.md" + plan.write_text("follow the oracle plan\n", encoding="utf-8") + decoy = tmp_path / "decoy-plan.md" + decoy.write_text("follow the decoy plan\n", encoding="utf-8") + return f"inject:{plan},inject:{decoy}" + + +def test_post_research_with_a_marker_passes_preflight(tmp_path: Path) -> None: + """--mark-research-end-on lifts the pre-flight rejection the reviewer named: + 'bench eval ablate rejects post-research'.""" + arms = parse_arms(_inject_arms(tmp_path)) + assert ( + validate_arms_for_stage( + arms, "post-research", research_end_marker="/app/PLAN.md" + ) + == "post-research" + ) + + +def test_post_research_rejection_names_the_marker_flag() -> None: + """Without the trigger the rejection now teaches the CLI recipe too.""" + arms = parse_arms("with-skill,no-skill") + with pytest.raises(AblationSpecError, match="--mark-research-end-on"): + validate_arms_for_stage(arms, "post-research") + + +def test_marker_on_a_non_marked_stage_fails_closed(tmp_path: Path) -> None: + """The marker defines the research-end boundary and nothing else.""" + arms = parse_arms(_inject_arms(tmp_path)) + with pytest.raises(AblationSpecError, match="only applies to --at-stage"): + validate_arms_for_stage(arms, "env-ready", research_end_marker="/app/PLAN.md") + + +class _ResearchMarkerRollout(_FakeRollout): + """A parent whose agent materializes PLAN.md mid-``execute()``. + + The sandbox double answers the watcher's ``test -e`` from + ``marker_present``, which ``execute()`` flips *before* yielding the event + loop — so the watcher observes the file while the agent is still running, + exactly the mid-run capture the wire adds. + """ + + def __init__(self, config) -> None: + super().__init__(config) + self.marker_present = False + self.marker_commands: list[str] = [] + self._stage_snapshots: dict[str, Any] = {} + rollout = self + + class _Env: + async def exec(self, cmd, *, user="root", timeout_sec=30): + rollout.marker_commands.append(cmd) + found = rollout.marker_present and "PLAN.md" in cmd + return SimpleNamespace(return_code=0 if found else 1) + + self.env = _Env() + + async def execute(self) -> None: + self.calls.append("execute") + self.marker_present = True # the agent just wrote its plan... + await asyncio.sleep(0.05) # ...and keeps running while the watcher polls + + async def mark_stage(self, name: str): + from benchflow.branch import StageSnapshot + from benchflow.sandbox.protocol import SandboxImage + + self.calls.append(f"mark_stage:{name}") + snap = StageSnapshot( + environment_ref=None, + sandbox_ref=SandboxImage(provider="fake", ref="bf-snap-research"), + stage=name, + meta={"exchanges_completed": 3}, + ) + self._stage_snapshots[name] = snap + return snap + + +async def test_marker_file_appearing_mid_run_captures_post_research_and_branches( + tmp_path: Path, monkeypatch +) -> None: + """PLAN.md appears at exchange k → the post-research snapshot exists and + branch_at_stage('post-research') forks the arms — the reviewer-named + end-to-end gap for `bench eval ablate --at-stage post-research`.""" + built = _patch_rollout(monkeypatch, _ResearchMarkerRollout) + request = AblationRequest( + task_path=_task_dir(tmp_path), + arms=parse_arms(_inject_arms(tmp_path)), + agent="claude-agent-acp", + model="claude-sonnet", + stage="post-research", + out_dir=tmp_path / "out", + mark_research_end_on="/app/PLAN.md", + ) + + report = await run_ablation(request) + + rollout = built[0] + # The watcher marked the stage while the agent was still executing — + # before verify(), from a `test -e` on the marker path. + assert "mark_stage:post-research" in rollout.calls + assert rollout.calls.index("mark_stage:post-research") < rollout.calls.index( + "verify" + ) + assert any("test -e /app/PLAN.md" in cmd for cmd in rollout.marker_commands) + assert "post-research" in rollout._stage_snapshots + # ...and the branch really forked there, one child per arm, rewards read. + assert "branch_at_stage:post-research" in rollout.calls + assert report.stage == "post-research" + assert [arm.reward for arm in report.arms] == [1.0, 0.0] + assert report.error is None + + +class _NeverMarkedRollout(_FakeRollout): + """A parent whose workspace never grows the marker file.""" + + def __init__(self, config) -> None: + super().__init__(config) + self._stage_snapshots: dict[str, Any] = {} + + class _Env: + async def exec(self, cmd, *, user="root", timeout_sec=30): + return SimpleNamespace(return_code=1) + + self.env = _Env() + + async def branch_at_stage(self, stage, n, *, deltas=None) -> float: + raise AssertionError("must not fork a stage that was never captured") + + +async def test_marker_never_appearing_reports_the_missing_capture( + tmp_path: Path, monkeypatch +) -> None: + """The unacceptable outcome would be silently forking a different stage — + instead the report says the marker never appeared and skips every arm.""" + built = _patch_rollout(monkeypatch, _NeverMarkedRollout) + request = AblationRequest( + task_path=_task_dir(tmp_path), + arms=parse_arms(_inject_arms(tmp_path)), + agent="claude-agent-acp", + model="claude-sonnet", + stage="post-research", + out_dir=tmp_path / "out", + mark_research_end_on="/app/PLAN.md", + ) + + report = await run_ablation(request) + + assert built[0].calls[-1] == "cleanup" + assert "never appeared" in report.error + assert "/app/PLAN.md" in report.error + assert [arm.status for arm in report.arms] == ["skipped", "skipped"] + assert report.has_errors is True + + +async def test_watch_research_end_polls_until_the_file_exists() -> None: + """The watcher itself: polls, ignores transient exec failures, marks once.""" + from benchflow.ablate import watch_research_end + + class _Rollout: + def __init__(self) -> None: + self.polls = 0 + self.marked: list[str] = [] + self._agent_cwd = "/app" + rollout = self + + class _Env: + async def exec(self, cmd, *, user="root", timeout_sec=30): + rollout.polls += 1 + if rollout.polls == 1: + raise RuntimeError("transient exec failure") + return SimpleNamespace(return_code=0 if rollout.polls >= 3 else 1) + + self.env = _Env() + + async def mark_stage(self, name: str) -> None: + self.marked.append(name) + + rollout = _Rollout() + + marked = await watch_research_end(rollout, "PLAN.md", poll_interval=0.001) + + assert marked is True + assert rollout.polls == 3 + assert rollout.marked == ["post-research"] + + +def test_cli_rejects_post_research_without_the_marker_flag(tmp_path: Path) -> None: + result = runner.invoke( + app, + [ + "eval", + "ablate", + "--tasks-dir", + str(_task_dir(tmp_path)), + "--at-stage", + "post-research", + "--arms", + _inject_arms(tmp_path), + ], + ) + assert result.exit_code == 1, result.output + assert "--mark-research-end-on" in _flat(result.stderr) + assert "Traceback (most recent call last)" not in result.output + + +class _FileWritingStageRefRollout(_ExportingStageRefRollout): + """A stage-ref rollout that also leaves capture-time stage_snapshots.json, + the way the real engine's ``capture_stage`` does.""" + + async def branch_at_stage(self, stage, n, *, deltas=None): + value = await super().branch_at_stage(stage, n, deltas=deltas) + from benchflow.branch_lineage import write_stage_snapshots + + self._rollout_dir.mkdir(parents=True, exist_ok=True) + write_stage_snapshots( + run_dir=self._rollout_dir, snapshots=self._stage_snapshots + ) + return value + + +async def test_retention_annotates_the_parent_runs_stage_snapshots_file( + tmp_path: Path, monkeypatch +) -> None: + """Guards "feat(rollout): stage snapshots record their lifetime; + --keep-snapshots on bench eval run with a tested import path" (the + ablate/plain consistency half): the parent run's stage_snapshots.json is + the on-disk twin of ``report.stage_snapshot``, so the retention decision + — exported where, or ephemeral — must land in both, teaching readers one + schema.""" + import dataclasses + + _patch_rollout(monkeypatch, _FileWritingStageRefRollout) + request = dataclasses.replace(_request(tmp_path), keep_snapshots=True) + + report = await run_ablation(request) + + run_dir = Path(report.parent_run_dir) + recorded = json.loads((run_dir / "stage_snapshots.json").read_text())["stages"][ + report.stage + ] + assert recorded == report.stage_snapshot + assert recorded["ephemeral"] is False + assert recorded["exported"]["path"] == str( + Path(request.out_dir) / "snapshots" / "bf-snap-root-1.tar" + ) + + +async def test_without_keep_snapshots_the_parent_file_is_marked_ephemeral( + tmp_path: Path, monkeypatch +) -> None: + """The default half of the same consistency guard: no flag, so the file + entry says the ref died with the run — never a bare ref.""" + _patch_rollout(monkeypatch, _FileWritingStageRefRollout) + + report = await run_ablation(_request(tmp_path)) + + run_dir = Path(report.parent_run_dir) + recorded = json.loads((run_dir / "stage_snapshots.json").read_text())["stages"][ + report.stage + ] + assert recorded == report.stage_snapshot + assert recorded["ephemeral"] is True + assert recorded["exported"] is None diff --git a/tests/test_acp_model_config_dispatch.py b/tests/test_acp_model_config_dispatch.py index 26f53f805..974c7b718 100644 --- a/tests/test_acp_model_config_dispatch.py +++ b/tests/test_acp_model_config_dispatch.py @@ -113,6 +113,154 @@ async def test_codex_litellm_alias_uses_bare_model_for_set_model(tmp_path): mock_acp.set_config_option.assert_not_awaited() +@pytest.mark.asyncio +async def test_codex_off_catalog_model_owned_by_launch_config_skips_set_model( + tmp_path, +): + """Guards the fix for codex-acp@1.6.0 catalog-validated set_model. + + 1.6.0 rejects ``session/set_model`` for any model absent from its built-in + catalog ("Unknown model gpt-5.4-mini[medium]", verified live 2026-08-21) — + even when that model is already the session's current model via the + ``CODEX_CONFIG`` injection ``apply_codex_provider_config`` writes for the + LiteLLM gateway route. When the requested model maps to no advertised + ``model[effort]`` variant and the session already runs BenchFlow's own + injected route, the runtime must skip the doomed call instead of failing + the whole rollout.""" + import json + + mock_acp = _make_mocks( + config_options=[{"id": "model"}], + model_state={ + "availableModels": [ + {"modelId": "gpt-5.6-sol[medium]"}, + {"modelId": "gpt-5.5[medium]"}, + ], + "currentModelId": "benchflow-us-openai-gpt-5.4-mini[medium]", + }, + ) + await _connect( + mock_acp, + agent="codex-acp", + model="us-openai/gpt-5.4-mini", + tmp_path=tmp_path, + agent_env={ + LITELLM_MODEL_VIA_ENV: "1", + LITELLM_MODEL_ALIAS_ENV: "benchflow-us-openai-gpt-5.4-mini", + "CODEX_CONFIG": json.dumps( + { + "model": "benchflow-us-openai-gpt-5.4-mini", + "model_provider": "benchflow-litellm", + } + ), + }, + ) + + mock_acp.set_model.assert_not_awaited() + mock_acp.set_config_option.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_codex_off_catalog_skip_satisfies_effort_carried_by_launch_config( + tmp_path, +): + """A requested effort that the CODEX_CONFIG-injected current model already + carries is satisfied by the skip — the effort step must not fail closed + for an effort that is in place.""" + import json + + mock_acp = _make_mocks( + config_options=[{"id": "model"}], + model_state={ + "availableModels": [{"modelId": "gpt-5.6-sol[medium]"}], + "currentModelId": "benchflow-us-openai-gpt-5.4-mini[xhigh]", + }, + ) + await _connect( + mock_acp, + agent="codex-acp", + model="us-openai/gpt-5.4-mini", + tmp_path=tmp_path, + reasoning_effort="xhigh", + agent_env={ + LITELLM_MODEL_VIA_ENV: "1", + "CODEX_CONFIG": json.dumps({"model": "benchflow-us-openai-gpt-5.4-mini"}), + }, + ) + + mock_acp.set_model.assert_not_awaited() + mock_acp.set_config_option.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_codex_off_catalog_skip_with_unsatisfied_effort_fails_closed( + tmp_path, +): + """When the launch config owns an off-catalog model but does NOT carry the + requested effort, there is no channel left to deliver it — the existing + effort step must still fail closed rather than silently drop it.""" + import json + + mock_acp = _make_mocks( + config_options=[{"id": "model"}], + model_state={ + "availableModels": [{"modelId": "gpt-5.6-sol[medium]"}], + "currentModelId": "benchflow-us-openai-gpt-5.4-mini[medium]", + }, + ) + with pytest.raises(RuntimeError, match="does not declare an ACP effort"): + await _connect( + mock_acp, + agent="codex-acp", + model="us-openai/gpt-5.4-mini", + tmp_path=tmp_path, + reasoning_effort="xhigh", + agent_env={ + LITELLM_MODEL_VIA_ENV: "1", + "CODEX_CONFIG": json.dumps( + {"model": "benchflow-us-openai-gpt-5.4-mini"} + ), + }, + ) + + mock_acp.set_model.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_codex_in_catalog_model_still_uses_set_model_despite_launch_config( + tmp_path, +): + """An in-catalog model keeps the set_model path even when CODEX_CONFIG + injected a gateway route: the skip is only for models set_model cannot + express (no advertised variant), so the 62cc7e41 bare→``model[effort]`` + mapping must not regress.""" + import json + + mock_acp = _make_mocks( + config_options=[{"id": "model"}], + model_state={ + "availableModels": [ + {"modelId": "gpt-5.4-mini[low]"}, + {"modelId": "gpt-5.4-mini[medium]"}, + ], + "currentModelId": "benchflow-us-openai-gpt-5.4-mini[medium]", + }, + ) + await _connect( + mock_acp, + agent="codex-acp", + model="us-openai/gpt-5.4-mini", + tmp_path=tmp_path, + agent_env={ + LITELLM_MODEL_VIA_ENV: "1", + "CODEX_CONFIG": json.dumps({"model": "benchflow-us-openai-gpt-5.4-mini"}), + }, + ) + + mock_acp.set_model.assert_awaited_once_with("gpt-5.4-mini[medium]") + mock_acp.set_config_option.assert_not_awaited() + + @pytest.mark.asyncio async def test_codex_with_model_option_still_uses_set_model(tmp_path): """codex-acp@1.6.0 advertises a 'model' config option whose values reject @@ -198,6 +346,109 @@ async def test_effort_without_effort_config_id_fails_closed(tmp_path): mock_acp.close.assert_awaited() +# Request-global classification (PR #1046 second review, P2-A): with Gemini +# and --reasoning-effort high the run built the environment, snapshotted, +# installed and connected the agent, and only then hit the effort rejection — +# which ablation then retried in a branch child. These pin the two halves of +# the fix's classification layer: a deterministic rejection of a global +# request setting (effort/model) surfaces as ACPRequestGlobalError and +# classifies as REQUEST_GLOBAL (non-retryable, never task-attributable), +# while a mere timeout on the same call proves nothing about compatibility +# and must NOT be branded request-global. + + +@pytest.mark.asyncio +async def test_unsupported_effort_is_a_request_global_rejection(tmp_path): + """The reviewer's exact shape: gemini declares no ACP effort config option, + so a requested effort is rejected — and the rejection must classify as + request-global, not as a retryable/task-attributable agent error.""" + from benchflow._utils.scoring import REQUEST_GLOBAL, classify_error + from benchflow._utils.text import describe_exception + from benchflow.acp.runtime import ACPRequestGlobalError + + mock_acp = _make_mocks(config_options=[]) + with pytest.raises( + ACPRequestGlobalError, match="does not declare an ACP effort" + ) as excinfo: + await _connect( + mock_acp, + agent="gemini", + model=None, + tmp_path=tmp_path, + reasoning_effort="high", + ) + + assert classify_error(describe_exception(excinfo.value)) == REQUEST_GLOBAL + + +@pytest.mark.asyncio +async def test_agent_rejected_effort_option_is_request_global(tmp_path): + """An agent *answering* set_config_option with a protocol error is a + deterministic rejection of the requested value — request-global.""" + from unittest.mock import AsyncMock + + from benchflow._utils.scoring import REQUEST_GLOBAL, classify_error + from benchflow._utils.text import describe_exception + from benchflow.acp.client import ACPError + from benchflow.acp.runtime import ACPRequestGlobalError + + mock_acp = _make_mocks(config_options=[{"id": "effort"}]) + mock_acp.set_config_option = AsyncMock( + side_effect=ACPError(-32602, "unsupported effort: xhigh") + ) + with pytest.raises(ACPRequestGlobalError) as excinfo: + await _connect( + mock_acp, + agent="claude-agent-acp", + model=None, + tmp_path=tmp_path, + reasoning_effort="xhigh", + ) + + assert classify_error(describe_exception(excinfo.value)) == REQUEST_GLOBAL + + +@pytest.mark.asyncio +async def test_timed_out_effort_option_is_not_request_global(tmp_path): + """A timeout on set_config_option is transport trouble, not evidence the + setting is unsupported — it must keep its ordinary (retryable) class.""" + from unittest.mock import AsyncMock + + from benchflow._utils.scoring import REQUEST_GLOBAL, classify_error + from benchflow._utils.text import describe_exception + from benchflow.acp.runtime import ACPRequestGlobalError + + mock_acp = _make_mocks(config_options=[{"id": "effort"}]) + mock_acp.set_config_option = AsyncMock(side_effect=TimeoutError()) + with pytest.raises(RuntimeError) as excinfo: + await _connect( + mock_acp, + agent="claude-agent-acp", + model=None, + tmp_path=tmp_path, + reasoning_effort="high", + ) + + assert not isinstance(excinfo.value, ACPRequestGlobalError) + assert classify_error(describe_exception(excinfo.value)) != REQUEST_GLOBAL + + +def test_reasoning_effort_preflight_matches_the_runtime_dispatch(): + """The static pre-flight must reject exactly what the runtime would reject + from registry facts alone: no acp_effort_config_id and no codex-style + effort-in-model-id. Agents the registry cannot vouch for are left to the + runtime's own fail-closed check.""" + from benchflow.acp.runtime import reasoning_effort_preflight_error + + assert reasoning_effort_preflight_error("gemini", None) is None + assert reasoning_effort_preflight_error("claude-agent-acp", "high") is None + assert reasoning_effort_preflight_error("codex-acp", "high") is None + assert reasoning_effort_preflight_error("no-such-agent", "high") is None + error = reasoning_effort_preflight_error("gemini", "high") + assert error is not None + assert "not supported by agent 'gemini'" in error + + @pytest.mark.asyncio async def test_env_owned_model_skips_advertised_model_option(tmp_path): """A manifest-shaped agent (supports_acp_set_model=False + a diff --git a/tests/test_branch_child_result.py b/tests/test_branch_child_result.py new file mode 100644 index 000000000..26d78c8e7 --- /dev/null +++ b/tests/test_branch_child_result.py @@ -0,0 +1,294 @@ +"""Regression tests for in-place branch-child result synthesis. + +Guards "feat(branch): synthesize a full result.json for in-place branch +children" (docs/rollout-branching-rfc.md §3.4). A fresh-rollout child +(``env-ready``) is a Rollout of its own and leaves the standard artifact set +for free; an in-place child (``pre-verify`` / ``post-verify`` / cursor branch) +continues the parent instance and used to leave only ``provenance.json`` / +``reward.json`` plus the ``mounted/`` archive — so "what happened in this arm" +required cross-reading ``tree.json``. Now every completed in-place child of a +run-dir-bearing rollout leaves its own ``result.json`` / ``timing.json`` / +trajectory, built from the child's OWN state (scoped to zero before it runs), +with no parent bleed-through: a field the child did not produce is null/absent, +never copied from the parent. + +Unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from benchflow.diagnostics import TransportClosedDiagnostic +from benchflow.environment.protocol import StateSnapshot +from benchflow.rollout import Rollout, RolloutConfig, Scene +from benchflow.task.paths import RolloutPaths + + +class FakeEnvironment: + """Environment-plane stand-in recording snapshot/restore calls.""" + + def __init__(self) -> None: + self.snapshots: list[StateSnapshot] = [] + self.restored: list[StateSnapshot] = [] + + async def snapshot(self) -> StateSnapshot: + snap = StateSnapshot(id=f"snap-{len(self.snapshots) + 1}", path="/tmp/x") + self.snapshots.append(snap) + return snap + + async def restore(self, snap: StateSnapshot) -> None: + self.restored.append(snap) + + +def _rollout(tmp_path: Path) -> Rollout: + rollout = Rollout( + RolloutConfig(task_path=tmp_path / "task", scenes=[Scene.single(agent="dummy")]) + ) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + rollout._rollout_dir = run_dir + RolloutPaths(rollout_dir=run_dir).mkdir() + return rollout + + +async def _run_parent_step(rollout: Rollout, monkeypatch) -> None: + """One real execute() on the parent, so the fork has linear history.""" + + async def fake_execute_prompts(*_a, **_kw): + return [{"role": "agent", "text": "child-step"}], 1 + + monkeypatch.setattr(rollout._planes, "execute_prompts", fake_execute_prompts) + rollout._acp_client = object() + await rollout.execute(["parent-prompt"]) + + +def _fake_child_phases(monkeypatch, verify_outcomes: list[dict | None]) -> None: + """Class-level connect/disconnect/verify fakes writing child-own state. + + ``verify()`` mutates the instance the way the real phase does — assigning + ``_rewards`` / ``_verifier_error`` outright and writing into ``_timing`` — + with per-child values drawn from ``verify_outcomes`` in fork order. + """ + outcomes = iter(verify_outcomes) + + async def fake_connect(self): + self._acp_client = object() + + async def fake_disconnect(self): + # Mirror the real disconnect: the session counters rewind, so a fresh + # child session's cumulative trajectory is read from its own start. + self._acp_client = None + self._session_tool_count = 0 + self._session_traj_count = 0 + + async def fake_verify(self): + rewards = next(outcomes) + self._timing["verification"] = 99.0 + self._verifier_error = "child verifier: exit 0" + self._diagnostics.set( + TransportClosedDiagnostic(raw_message="child", transport_diagnosis="child") + ) + self._rewards = rewards + return self._rewards + + monkeypatch.setattr(Rollout, "connect", fake_connect) + monkeypatch.setattr(Rollout, "disconnect", fake_disconnect) + monkeypatch.setattr(Rollout, "verify", fake_verify) + + +def _child_dir(rollout: Rollout, parent, child) -> Path: + return rollout._rollout_dir / "branches" / parent.id / "children" / child.id + + +async def test_in_place_children_leave_their_own_result_json( + tmp_path: Path, monkeypatch +) -> None: + """Each completed in-place child leaves a result.json of its OWN run. + + The red case this guards: only fresh-rollout (env-ready) children were + first-class runs; an in-place ablation arm left no result.json at all. + Reward, verifier error, timing and trajectory must all be the child's own + — and the fields the child did not produce (the parent's error, the + parent's timing keys, token usage) must be null/absent, never inherited. + """ + rollout = _rollout(tmp_path) + await _run_parent_step(rollout, monkeypatch) + + # The parent's own result-bearing state, as a linear run leaves it. None + # of it may appear in a child's result. + rollout._verifier_error = "parent verifier: 1 test failed" + rollout._error = "parent agent: idle for 600s" + rollout._timing["agent_setup"] = 41.5 + timing_before = dict(rollout._timing) + parent = rollout._cursor + + _fake_child_phases(monkeypatch, [{"reward": 1.0}, {"reward": 0.0}]) + + assert await rollout.branch(2) == 0.5 + + children = parent.children + assert len(children) == 2 + for child, expected_reward in zip(children, [1.0, 0.0], strict=True): + child_dir = _child_dir(rollout, parent, child) + result = json.loads((child_dir / "result.json").read_text()) + # The child's own observations. + assert result["rewards"] == {"reward": expected_reward} + assert result["verifier_error"] == "child verifier: exit 0" + assert result["rollout_name"] == child.id + assert result["task_name"] == "task" + # No parent bleed-through: fields the child never produced are null. + assert result["error"] is None + assert result["agent_result"]["n_input_tokens"] is None + assert result["agent_result"]["usage_source"] == "unavailable" + # result.json names its own lineage, without tree.json. + assert result["source"]["kind"] == "benchflow-branch" + assert result["source"]["branch_stage"] == f"cursor:{parent.id}" + # The child's own timing: its verify wrote 99.0; the parent's + # agent_setup is the parent's, never copied. + timing = json.loads((child_dir / "timing.json").read_text()) + assert timing["verification"] == 99.0 + assert "agent_setup" not in timing + assert timing["total"] >= 0 + # The trajectory is the child's continuation only — one step, not the + # parent's history relabelled. + lines = ( + (child_dir / "trajectory" / "acp_trajectory.jsonl") + .read_text() + .strip() + .splitlines() + ) + assert len(lines) == 1 + # reward.json (the lineage artifact) still stands beside it. + assert json.loads((child_dir / "reward.json").read_text()) == { + "reward": expected_reward + } + + # The parent's reported state is still the parent's own (the isolation + # invariant this feature must not weaken). + assert rollout._timing == timing_before + assert rollout._verifier_error == "parent verifier: 1 test failed" + assert rollout._error == "parent agent: idle for 600s" + + +async def test_an_unscored_in_place_child_still_leaves_its_result( + tmp_path: Path, monkeypatch +) -> None: + """A child that ran but was never scored is evidenced, not invented. + + Its result.json exists with ``rewards: null`` (a missing score is not a + zero) and its own verifier error; reward.json is absent, exactly as the + lineage writer has always kept it. + """ + rollout = _rollout(tmp_path) + await _run_parent_step(rollout, monkeypatch) + parent = rollout._cursor + _fake_child_phases(monkeypatch, [{}, {}]) + + assert await rollout.branch(2) is None + + for child in parent.children: + child_dir = _child_dir(rollout, parent, child) + result = json.loads((child_dir / "result.json").read_text()) + assert result["rewards"] is None + assert result["verifier_error"] == "child verifier: exit 0" + assert not (child_dir / "reward.json").exists() + + +async def test_result_synthesis_failure_never_costs_the_reward( + tmp_path: Path, monkeypatch +) -> None: + """The writer is best-effort by contract, like every branch artifact. + + A result set that cannot be built (full disk, a field that will not + serialize) is logged and skipped — the fork still scores, the node still + carries its reward, and reward.json is still written. + """ + rollout = _rollout(tmp_path) + await _run_parent_step(rollout, monkeypatch) + parent = rollout._cursor + _fake_child_phases(monkeypatch, [{"reward": 1.0}, {"reward": 1.0}]) + + def boom(*_a, **_kw): + raise OSError("No space left on device") + + monkeypatch.setattr("benchflow.rollout._results._build_rollout_result", boom) + + assert await rollout.branch(2) == 1.0 + + for child in parent.children: + child_dir = _child_dir(rollout, parent, child) + assert not (child_dir / "result.json").exists() + assert json.loads((child_dir / "reward.json").read_text()) == {"reward": 1.0} + + +async def test_a_child_that_raises_hard_leaves_no_result_json( + tmp_path: Path, monkeypatch +) -> None: + """A hard child failure ends the fork before a result can be honest. + + The failing child's evidence is its ``mounted/`` archive and its + provenance (the partial-lineage guarantee); a synthesized result for a + child that died mid-phase would describe a run that never completed. + """ + rollout = _rollout(tmp_path) + await _run_parent_step(rollout, monkeypatch) + parent = rollout._cursor + + async def run_child(child): + raise RuntimeError("agent connection lost") + + with pytest.raises(RuntimeError, match="agent connection lost"): + await rollout.branch(2, run_child=run_child) + + child = parent.children[0] + child_dir = _child_dir(rollout, parent, child) + assert (child_dir / "provenance.json").is_file() + assert not (child_dir / "result.json").exists() + + +async def test_fresh_rollout_children_are_not_double_synthesized( + tmp_path: Path, monkeypatch +) -> None: + """An ``env-ready`` fresh child writes its own result — the engine must + not overwrite it with a synthesized in-place one.""" + import benchflow.rollout_branch as rollout_branch + from benchflow.sandbox.protocol import SandboxImage + from benchflow.trajectories.tree import Step + + class FakeSnapSandbox: + supports_snapshot = True + + async def snapshot(self, name: str | None = None) -> SandboxImage: + return SandboxImage(provider="fake", ref="bf-snap-1") + + async def restore(self, image: SandboxImage) -> None: + return None + + rollout = _rollout(tmp_path) + rollout._env = FakeSnapSandbox() + await rollout.mark_stage("env-ready", snapshot_layers={"environment", "sandbox"}) + rollout._cursor = rollout._tree.advance(rollout._cursor, Step(id="s1")) + + synthesized: list[str] = [] + real_write = rollout_branch.write_in_place_child_result + + def recording_write(*args, **kwargs): + synthesized.append(kwargs["child"].id) + return real_write(*args, **kwargs) + + monkeypatch.setattr(rollout_branch, "write_in_place_child_result", recording_write) + + def make_runner(rollout_arg, **kwargs): + async def _runner(child): + return 1.0 + + return _runner + + monkeypatch.setattr(rollout_branch, "make_fresh_child_runner", make_runner) + + assert await rollout.branch_at_stage("env-ready", 2) == 1.0 + assert synthesized == [] diff --git a/tests/test_branch_composed.py b/tests/test_branch_composed.py new file mode 100644 index 000000000..d5b60c329 --- /dev/null +++ b/tests/test_branch_composed.py @@ -0,0 +1,440 @@ +"""Regression tests for the composed checkpoint layer. + +Guards the composed-checkpoint layer ("feat(branch): compose sandbox + +environment checkpoints in the branch engine"; docs/rollout-branching-rfc.md +WS-1; FrontierPhysics#73). PR number to be added on submission. + +The composed checkpoint (RFC §3.1) layers the container snapshot +(``Sandbox.snapshot``) with the environment-state snapshot into one +``StageSnapshot``: environment first on checkpoint, sandbox first on restore. +The branch engine requests layers via ``snapshot_layers``; the default +``{"environment"}`` must keep the legacy environment-only path byte-for-byte. + +These are unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from benchflow.branch import ( + StageSnapshot, + checkpoint, + checkpoint_composed, + restore, + restore_composed, +) +from benchflow.environment.manifest import EnvironmentManifest +from benchflow.environment.manifest_env import ManifestEnvironment +from benchflow.environment.protocol import StateSnapshot +from benchflow.rollout import Rollout, RolloutConfig, Scene +from benchflow.sandbox.protocol import SandboxImage, SandboxSnapshotNotSupported +from benchflow.trajectories.tree import RolloutTree + + +class FakeEnv: + """Stateful in-memory Environment — snapshot copies state, restore rolls back. + + ``calls`` may be shared with a fake sandbox to record cross-layer order. + """ + + def __init__(self, calls: list[str] | None = None) -> None: + self.state: dict[str, str] = {"db": "initial"} + self._saved: dict[str, dict[str, str]] = {} + self.calls = calls if calls is not None else [] + self.snapshots: list[StateSnapshot] = [] + self.restored: list[StateSnapshot] = [] + + async def snapshot(self) -> StateSnapshot: + snap = StateSnapshot(id=f"env-snap-{len(self.snapshots) + 1}", path="/tmp/x") + self._saved[snap.id] = dict(self.state) + self.snapshots.append(snap) + self.calls.append("env.snapshot") + return snap + + async def restore(self, snap: StateSnapshot) -> None: + self.state = dict(self._saved[snap.id]) + self.restored.append(snap) + self.calls.append("env.restore") + + +class FakeSnapSandbox: + """Snapshot-capable Sandbox stand-in with in-memory filesystem state.""" + + supports_snapshot = True + + def __init__(self, calls: list[str] | None = None) -> None: + self.fs: dict[str, str] = {"/workspace": "clean"} + self._images: dict[str, dict[str, str]] = {} + self.calls = calls if calls is not None else [] + self.snapshots: list[SandboxImage] = [] + self.restored: list[SandboxImage] = [] + + async def snapshot(self, name: str | None = None) -> SandboxImage: + img = SandboxImage(provider="fake", ref=f"bf-snap-{len(self.snapshots) + 1}") + self._images[img.ref] = dict(self.fs) + self.snapshots.append(img) + self.calls.append("sandbox.snapshot") + return img + + async def restore(self, image: SandboxImage) -> None: + self.fs = dict(self._images[image.ref]) + self.restored.append(image) + self.calls.append("sandbox.restore") + + +class NoSnapshotSandbox: + """Sandbox without container snapshot — the capability gate must fail closed.""" + + supports_snapshot = False + + def __init__(self) -> None: + self.calls: list[str] = [] + + async def snapshot(self, name: str | None = None) -> SandboxImage: + self.calls.append("sandbox.snapshot") + raise SandboxSnapshotNotSupported("NoSnapshotSandbox does not support snapshot") + + async def restore(self, image: SandboxImage) -> None: + self.calls.append("sandbox.restore") + raise SandboxSnapshotNotSupported("NoSnapshotSandbox does not support restore") + + +# A manifest with no [environment.state] — a *stateless* environment whose +# snapshot()/restore() raise (the manifest_env fail-closed precedent, #387). +_STATELESS_MANIFEST = EnvironmentManifest.model_validate_toml( + """ +[environment] +name = "chi-bench" +image = "chi-bench:latest" +ports = [8020] +owns_lifecycle = true +""" +) + + +def _rollout(tmp_path: Path) -> Rollout: + return Rollout( + RolloutConfig(task_path=tmp_path / "task", scenes=[Scene.single(agent="dummy")]) + ) + + +# 1. Pure ops: composition order, fail-closed recording, both snapshot shapes + + +async def test_checkpoint_composed_snapshots_environment_before_sandbox(): + """RFC §3.1 order: environment.snapshot() first, then sandbox.snapshot().""" + tree = RolloutTree() + calls: list[str] = [] + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + + snap = await checkpoint_composed( + tree.root, environment=env, sandbox=sandbox, stage="env-ready" + ) + + assert calls == ["env.snapshot", "sandbox.snapshot"] + assert isinstance(snap, StageSnapshot) + assert snap.environment_ref is env.snapshots[0] + assert snap.sandbox_ref is sandbox.snapshots[0] + assert snap.stage == "env-ready" + assert tree.root.state["snapshot"] is snap + + +async def test_restore_composed_restores_sandbox_before_environment(): + """RFC §3.1 restore order is reversed: sandbox first, then environment.""" + tree = RolloutTree() + calls: list[str] = [] + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + await checkpoint_composed(tree.root, environment=env, sandbox=sandbox) + calls.clear() + + await restore_composed(tree.root, environment=env, sandbox=sandbox) + + assert calls == ["sandbox.restore", "env.restore"] + + +async def test_checkpoint_composed_needs_at_least_one_layer(): + tree = RolloutTree() + with pytest.raises(ValueError, match="at least one layer"): + await checkpoint_composed(tree.root) + + +async def test_checkpoint_composed_failure_records_nothing_on_the_node(): + """A failing layer propagates and never leaves a partial StageSnapshot. + + The environment layer snapshots fine, the sandbox layer raises — the node + must not carry a half-recorded roll-back point. + """ + tree = RolloutTree() + env, sandbox = FakeEnv(), NoSnapshotSandbox() + + with pytest.raises(SandboxSnapshotNotSupported): + await checkpoint_composed(tree.root, environment=env, sandbox=sandbox) + + assert "snapshot" not in tree.root.state + + +async def test_restore_composed_accepts_a_legacy_bare_state_snapshot(): + """Back-compat: checkpoint() consumers keep working through restore_composed. + + A node checkpointed by the legacy checkpoint() holds a bare StateSnapshot; + restore_composed must treat it as an environment-only checkpoint. + """ + tree = RolloutTree() + env = FakeEnv() + snap = await checkpoint(tree.root, env) + + await restore_composed(tree.root, environment=env) + + assert env.restored == [snap] + + +async def test_restore_composed_without_a_checkpoint_raises(): + tree = RolloutTree() + with pytest.raises(ValueError, match="no checkpoint"): + await restore_composed(tree.root, environment=FakeEnv()) + + +# 2. Pure ops: mismatched snapshot / live-object shapes fail closed + + +async def test_restore_composed_rejects_a_missing_live_object_for_a_layer(): + """A layer present in the checkpoint needs its live object — and no layer + is restored before the mismatch is caught.""" + tree = RolloutTree() + env, sandbox = FakeEnv(), FakeSnapSandbox() + await checkpoint_composed(tree.root, environment=env, sandbox=sandbox) + + with pytest.raises(ValueError, match="no live sandbox"): + await restore_composed(tree.root, environment=env) + with pytest.raises(ValueError, match="no live environment"): + await restore_composed(tree.root, sandbox=sandbox) + + assert sandbox.restored == [] + assert env.restored == [] + + +async def test_restore_composed_rejects_a_live_object_without_a_layer(): + """A live object for a layer the checkpoint never captured is a bug.""" + tree = RolloutTree() + env, sandbox = FakeEnv(), FakeSnapSandbox() + await checkpoint_composed(tree.root, sandbox=sandbox) + + with pytest.raises(ValueError, match="no environment layer"): + await restore_composed(tree.root, environment=env, sandbox=sandbox) + + assert sandbox.restored == [] + + +async def test_legacy_restore_rejects_a_composed_stage_snapshot(): + """A StageSnapshot node fails closed in legacy restore() — use restore_composed. + + A node checkpointed via checkpoint_composed() holds a StageSnapshot, not a + bare StateSnapshot; passing it to the legacy restore() used to explode deep + inside the environment (AttributeError on ``.path``). It must be rejected + with a clear ValueError *before* the environment is touched. + """ + tree = RolloutTree() + env, sandbox = FakeEnv(), FakeSnapSandbox() + await checkpoint_composed(tree.root, environment=env, sandbox=sandbox) + + with pytest.raises(ValueError, match="restore_composed"): + await restore(tree.root, env) + + assert env.restored == [] + assert env.calls == ["env.snapshot"] # checkpoint only — no restore call + + +async def test_restore_composed_rejects_a_sandbox_against_a_legacy_snapshot(): + """A legacy bare StateSnapshot is environment-only — a live sandbox is a + shape mismatch, not a silent no-op.""" + tree = RolloutTree() + env, sandbox = FakeEnv(), FakeSnapSandbox() + await checkpoint(tree.root, env) + + with pytest.raises(ValueError, match="no sandbox layer"): + await restore_composed(tree.root, environment=env, sandbox=sandbox) + + assert sandbox.restored == [] + assert env.restored == [] + + +# 3. Pure ops: the zero-delta invariant (RFC §5, T2 shrunk to unit scale) + + +async def test_zero_delta_checkpoint_restore_round_trips_both_layers(): + """checkpoint_composed + restore_composed is lossless on both layers.""" + tree = RolloutTree() + env, sandbox = FakeEnv(), FakeSnapSandbox() + env.state["db"] = "checkpointed" + sandbox.fs["/workspace"] = "checkpointed" + await checkpoint_composed(tree.root, environment=env, sandbox=sandbox) + + env.state["db"] = "dirty" + sandbox.fs["/workspace"] = "dirty" + await restore_composed(tree.root, environment=env, sandbox=sandbox) + + assert env.state == {"db": "checkpointed"} + assert sandbox.fs == {"/workspace": "checkpointed"} + + +# 4. Engine: snapshot_layers wiring through Rollout.branch() + + +async def test_default_snapshot_layers_keep_the_legacy_environment_only_path( + tmp_path: Path, +): + """Back-compat: the default never touches the sandbox layer and records + the legacy bare-StateSnapshot shape on the node.""" + rollout = _rollout(tmp_path) + env, sandbox = FakeEnv(), FakeSnapSandbox() + rollout._environment = env + rollout._env = sandbox + parent = rollout._cursor + + async def run_child(child): + return 1.0 + + value = await rollout.branch(2, run_child=run_child) + + assert value == 1.0 + assert sandbox.calls == [] + snap = parent.state["snapshot"] + assert isinstance(snap, StateSnapshot) + assert not isinstance(snap, StageSnapshot) + + +async def test_engine_composed_branch_orders_layers_per_rfc(tmp_path: Path): + """With both layers requested the engine checkpoints env-then-sandbox once + and restores sandbox-then-env once per child.""" + rollout = _rollout(tmp_path) + calls: list[str] = [] + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + rollout._environment = env + rollout._env = sandbox + parent = rollout._cursor + + async def run_child(child): + return 1.0 + + value = await rollout.branch( + 2, run_child=run_child, snapshot_layers={"environment", "sandbox"} + ) + + assert value == 1.0 + assert calls == [ + "env.snapshot", + "sandbox.snapshot", + "sandbox.restore", + "env.restore", + "sandbox.restore", + "env.restore", + ] + snap = parent.state["snapshot"] + assert isinstance(snap, StageSnapshot) + assert snap.environment_ref is env.snapshots[0] + assert snap.sandbox_ref is sandbox.snapshots[0] + + +async def test_sandbox_layer_requested_but_unsupported_fails_closed(tmp_path: Path): + """The capability gate fires before anything is snapshotted — no partial + node.state mutation, the existing #384 diagnostic pattern.""" + rollout = _rollout(tmp_path) + env = FakeEnv() + rollout._environment = env + rollout._env = NoSnapshotSandbox() + parent = rollout._cursor + + async def run_child(child): + return 1.0 + + with pytest.raises(RuntimeError, match="container-level snapshot/restore"): + await rollout.branch( + 2, run_child=run_child, snapshot_layers={"environment", "sandbox"} + ) + + assert "snapshot" not in parent.state + assert env.calls == [] + + +async def test_environment_layer_on_a_stateless_env_propagates_the_typed_error( + tmp_path: Path, +): + """A stateless ManifestEnvironment's snapshot() RuntimeError propagates, + annotated with the requested snapshot_layers; the sandbox layer (ordered + second) is never snapshotted and the node stays unmutated.""" + rollout = _rollout(tmp_path) + sandbox = FakeSnapSandbox() + rollout._environment = ManifestEnvironment( + _STATELESS_MANIFEST, sandbox=FakeSnapSandbox() + ) + rollout._env = sandbox + parent = rollout._cursor + + async def run_child(child): + return 1.0 + + with pytest.raises(RuntimeError, match="stateless") as excinfo: + await rollout.branch( + 2, run_child=run_child, snapshot_layers={"environment", "sandbox"} + ) + + assert any("snapshot_layers" in note for note in excinfo.value.__notes__) + assert "snapshot" not in parent.state + assert sandbox.snapshots == [] + + +async def test_sandbox_only_branching_on_a_stateless_environment(tmp_path: Path): + """snapshot_layers={"sandbox"} branches a stateless env end-to-end: the + children restore from the container image and environment + snapshot/restore are never called (they would raise if they were).""" + rollout = _rollout(tmp_path) + sandbox = FakeSnapSandbox() + rollout._environment = ManifestEnvironment( + _STATELESS_MANIFEST, sandbox=FakeSnapSandbox() + ) + rollout._env = sandbox + parent = rollout._cursor + + returns = iter([1.0, 0.0]) + + async def run_child(child): + return next(returns) + + value = await rollout.branch(2, run_child=run_child, snapshot_layers={"sandbox"}) + + assert value == 0.5 + # one container checkpoint, one restore per child, all from the same image + assert len(sandbox.snapshots) == 1 + assert sandbox.restored == [sandbox.snapshots[0], sandbox.snapshots[0]] + snap = parent.state["snapshot"] + assert isinstance(snap, StageSnapshot) + assert snap.environment_ref is None + assert snap.sandbox_ref is sandbox.snapshots[0] + assert parent.state["value"] == 0.5 + + +async def test_unknown_snapshot_layer_is_rejected(tmp_path: Path): + rollout = _rollout(tmp_path) + rollout._environment = FakeEnv() + + async def run_child(child): + return 0.0 + + with pytest.raises(ValueError, match="unknown snapshot_layers"): + await rollout.branch( + 2, run_child=run_child, snapshot_layers={"environment", "agent-session"} + ) + + +async def test_empty_snapshot_layers_are_rejected(tmp_path: Path): + rollout = _rollout(tmp_path) + rollout._environment = FakeEnv() + + async def run_child(child): + return 0.0 + + with pytest.raises(ValueError, match="at least one layer"): + await rollout.branch(2, run_child=run_child, snapshot_layers=set()) diff --git a/tests/test_branch_composed_docker.py b/tests/test_branch_composed_docker.py new file mode 100644 index 000000000..a6f34e638 --- /dev/null +++ b/tests/test_branch_composed_docker.py @@ -0,0 +1,665 @@ +"""T2 oracle-invariant proofs for composed checkpoints, against real Docker. + +Guards the composed-checkpoint layer ("feat(branch): compose sandbox + +environment checkpoints in the branch engine"; docs/rollout-branching-rfc.md +WS-1, T2 tier per WS-3/§5; FrontierPhysics#73). PR number to be added on +submission. + +The rollout-branching RFC's validation plan (docs/rollout-branching-rfc.md §5) +names three tiers; ``tests/test_branch_composed.py`` is T1 (fakes). This file +is the T2 tier: an executable proof on a live ``DockerSandbox`` that the +composed checkpoint→restore of RFC §3.1 is lossless — zero delta ⇒ identical +observable state — and that a known-bad delta is detectable: + +1. ``test_container_layer_zero_delta_round_trip`` — the container layer alone + (``Sandbox.snapshot``/``restore``, i.e. ``docker commit`` → ``docker run``) + round-trips a file-tree digest computed inside the sandbox, and a + post-restore mutation moves the digest again (negative control). +2. ``test_composed_two_layer_round_trip`` — ``checkpoint_composed`` / + ``restore_composed`` over a real ``ManifestEnvironment`` (``[environment. + state] kind=sqlite``) plus the sandbox layer: file digest AND a sqlite + query both return to their pre-checkpoint values, and the recorded + ``StageSnapshot`` carries both layer refs. +3. ``test_zero_delta_reward_proxy_invariant`` — a deterministic in-sandbox + verifier (the reward function in miniature) reads PASS before checkpoint, + FAIL after a destructive mutation, PASS again after restore: the + reward-equality argument of RFC §5 (a) executed end-to-end. + +Marking: ``pytest.mark.live`` — the repo's marker for tests needing a real +Docker daemon (pyproject registers it as "requires real Anthropic API and +Docker daemon"; this file needs Docker only, no API keys — there is no +narrower docker-only marker, and ``integration`` implies Gemini/Daytona +credentials, so ``live`` is the closest fit). The default addopts deselect it; +when selected without a reachable daemon, the module-scoped ``docker_prereqs`` +fixture skips every test with a specific reason instead of erroring. + +Construction path: ``DockerSandbox`` is built exactly the way +``benchflow.sandbox.setup._create_sandbox_environment`` builds it for +``sandbox_type="docker"`` (a tiny environment dir + ``SandboxConfig``), with a +throwaway Dockerfile (alpine + the ``sqlite3`` CLI that +``ManifestEnvironment``'s sqlite backup/restore shells out to). Teardown +always runs: compose down via ``Sandbox.stop`` (with its force-kill-by-label +fallback), a container sweep by compose-project label, and ``docker rmi`` of +every ``bf-snap--*`` image this file's snapshots created (the tag +pattern from ``DockerSandbox.snapshot``). The built ``bf__`` base image +is deliberately kept so repeat runs hit the Docker build cache. +""" + +from __future__ import annotations + +import json +import shutil +import subprocess +import uuid +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path + +import pytest + +from benchflow.branch import StageSnapshot, checkpoint_composed, restore_composed +from benchflow.environment.manifest import EnvironmentManifest, StateSpec +from benchflow.environment.manifest_env import ManifestEnvironment +from benchflow.environment.protocol import StateSnapshot +from benchflow.sandbox.docker import ( + DockerSandbox, + _sanitize_docker_compose_project_name, +) +from benchflow.task.config import SandboxConfig +from benchflow.task.paths import RolloutPaths, SandboxPaths +from benchflow.trajectories.tree import RolloutTree + +pytestmark = [pytest.mark.live] + +# One environment name for the whole file: the image builds once +# (``DockerSandbox._image_build_locks`` + the Docker build cache) and the +# snapshot tag prefix below stays precise to this file's snapshots. +_ENV_NAME = "branch-composed-proof" +# Tag pattern from DockerSandbox.snapshot(): +# _sanitize_docker_image_name(f"bf-snap-{environment_name}-{suffix}") +_SNAP_IMAGE_PREFIX = f"bf-snap-{_ENV_NAME}-" + +# sqlite3 CLI is required by ManifestEnvironment's `.backup`-based +# snapshot/restore; busybox provides find/sort/xargs/stat/sha256sum. +_DOCKERFILE = """\ +FROM alpine:3.20 +RUN apk add --no-cache sqlite +WORKDIR /app +""" + +# Marker prefix for value-bearing exec output. `_run_docker_compose_command` +# merges stderr into stdout, so compose warnings (e.g. "Found orphan +# containers") can pollute the stream; values are extracted by prefix match +# instead of trusting raw stdout. +_MARKER = "BFPROOF:" + + +def _docker_unavailable_reason() -> str | None: + """Return why real-Docker tests cannot run here, or None if they can.""" + if shutil.which("docker") is None: + return "docker CLI not installed" + try: + result = subprocess.run( + ["docker", "info"], capture_output=True, timeout=10, check=False + ) + except (subprocess.TimeoutExpired, OSError) as exc: + return f"docker daemon unreachable: {exc}" + if result.returncode != 0: + return "docker daemon not running (`docker info` failed)" + return None + + +@pytest.fixture(scope="module") +def docker_prereqs() -> None: + """Skip (never error) when the Docker daemon is absent. + + Deferred to fixture time — like ``test_smoke.smoke_prereqs`` — so the + subprocess only fires when a live test is actually selected. + """ + reason = _docker_unavailable_reason() + if reason: + pytest.skip(reason) + + +@pytest.fixture(scope="module") +def environment_dir(tmp_path_factory: pytest.TempPathFactory) -> Path: + """A minimal sandbox environment definition: just a Dockerfile.""" + env_dir = tmp_path_factory.mktemp("branch-composed-env") + (env_dir / "Dockerfile").write_text(_DOCKERFILE) + return env_dir + + +def _sweep_project_containers(session_id: str) -> None: + """Best-effort removal of any container left carrying our project label. + + ``DockerSandbox.restore`` re-creates ``main`` outside compose (raw + ``docker run`` with compose labels); if ``compose down`` misses it, this + sweep — the same label filter ``_force_kill_project`` uses — catches it. + """ + project = _sanitize_docker_compose_project_name(session_id) + listed = subprocess.run( + [ + "docker", + "ps", + "-aq", + "--filter", + f"label=com.docker.compose.project={project}", + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + container_ids = listed.stdout.split() + if container_ids: + subprocess.run( + ["docker", "rm", "-f", "-v", *container_ids], + capture_output=True, + timeout=60, + check=False, + ) + subprocess.run( + ["docker", "network", "rm", f"{project}_default"], + capture_output=True, + timeout=30, + check=False, + ) + + +def _remove_snapshot_images() -> None: + """``docker rmi`` every snapshot image this file's tests created. + + Matches on the exact ``bf-snap--`` prefix so unrelated images + (including other tests' ``bf-snap-*``) are never touched. + """ + listed = subprocess.run( + ["docker", "images", "--format", "{{.Repository}}"], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + for repository in listed.stdout.split(): + if repository.startswith(_SNAP_IMAGE_PREFIX): + subprocess.run( + ["docker", "rmi", "-f", repository], + capture_output=True, + timeout=60, + check=False, + ) + + +@asynccontextmanager +async def _live_sandbox( + environment_dir: Path, rollout_dir: Path +) -> AsyncIterator[DockerSandbox]: + """Start a real DockerSandbox; guarantee teardown of containers + snapshots. + + Mirrors the ``sandbox_type == "docker"`` branch of + ``benchflow.sandbox.setup._create_sandbox_environment``. One sandbox per + test keeps every round trip starting from a clean compose-managed + ``main`` container. + """ + session_id = f"bf-branch-composed-{uuid.uuid4().hex[:8]}" + rollout_paths = RolloutPaths(rollout_dir=rollout_dir) + rollout_paths.mkdir() + sandbox = DockerSandbox( + environment_dir=environment_dir, + environment_name=_ENV_NAME, + session_id=session_id, + rollout_paths=rollout_paths, + task_env_config=SandboxConfig(), + ) + try: + await sandbox.start(force_build=False) + yield sandbox + finally: + try: + # delete=False: plain `compose down` — containers and network go, + # the built bf__ image stays cached for the next run. + await sandbox.stop(delete=False) + finally: + _sweep_project_containers(session_id) + _remove_snapshot_images() + + +async def _exec_ok(sandbox: DockerSandbox, command: str) -> None: + """Exec a state-mutating command and assert it succeeded.""" + result = await sandbox.exec(command, timeout_sec=120) + assert result.return_code == 0, ( + f"in-sandbox command failed (rc={result.return_code}): {command!r}\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + + +async def _exec_value(sandbox: DockerSandbox, command: str) -> str: + """Exec a value-producing command; return its single-line output. + + The value is echoed behind a unique marker prefix and extracted by prefix + match, so merged compose warnings cannot corrupt digest comparisons. + """ + wrapped = f'__bf_out="$({command})" && echo "{_MARKER}${{__bf_out}}"' + result = await sandbox.exec(wrapped, timeout_sec=120) + assert result.return_code == 0, ( + f"in-sandbox command failed (rc={result.return_code}): {command!r}\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + values = [ + line[len(_MARKER) :].strip() + for line in (result.stdout or "").splitlines() + if line.startswith(_MARKER) + ] + assert len(values) == 1, ( + f"expected exactly one {_MARKER} line for {command!r}, " + f"got {values!r}\nstdout: {result.stdout}" + ) + return values[0] + + +def _digest_command(exclude_glob: str | None = None) -> str: + """One deterministic line summarizing /app: file contents + tree + modes. + + Content hashes plus a ``stat`` listing of names and permission bits, all + digested by a final ``sha256sum``. ``exclude_glob`` drops files by + basename (used for the sqlite DB, whose ``.backup``-restored copy is + logically — not necessarily byte — identical; the DB is asserted through + a SQL query instead). + + Delegates to the src helper so the proofs measure with the same oracle + production records — the reviewer's P1-A repro (PR #1046 second review) + drove *this* helper's original ``find | sort | xargs`` pipeline, which + word-split legal filenames and emitted a successful wrong digest; the + null-safe fail-closed form now lives in exactly one place. + """ + from benchflow.sandbox.workspace_digest import workspace_digest_command + + return workspace_digest_command("/app", exclude_basename=exclude_glob) + + +async def test_container_layer_zero_delta_round_trip( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """Sandbox.snapshot → destructive mutation → Sandbox.restore is lossless. + + Zero delta ⇒ identical state digest; and the digest is a real oracle — + a post-restore mutation (the negative control) moves it again. + """ + async with _live_sandbox(environment_dir, tmp_path) as sandbox: + await _exec_ok( + sandbox, + "mkdir -p /app/t1/nested" + " && echo alpha > /app/t1/state.txt" + " && echo beta > /app/t1/nested/deep.txt" + " && chmod 640 /app/t1/state.txt" + " && chmod 700 /app/t1/nested", + ) + digest_before = await _exec_value(sandbox, _digest_command()) + + image = await sandbox.snapshot() + assert image.provider == "docker" + assert image.ref.startswith(_SNAP_IMAGE_PREFIX) + + await _exec_ok( + sandbox, + "echo corrupted > /app/t1/state.txt" + " && rm -rf /app/t1/nested" + " && echo extra > /app/t1/extra.txt" + " && chmod 600 /app/t1/state.txt", + ) + digest_mutated = await _exec_value(sandbox, _digest_command()) + assert digest_mutated != digest_before, ( + "mutation must move the digest — otherwise the digest is no oracle" + ) + + await sandbox.restore(image) + + digest_restored = await _exec_value(sandbox, _digest_command()) + assert digest_restored == digest_before, ( + "zero-delta restore must reproduce the exact pre-snapshot state" + ) + + # Negative control: a known-bad delta after restore is detectable. + await _exec_ok(sandbox, "echo tampered >> /app/t1/state.txt") + digest_tampered = await _exec_value(sandbox, _digest_command()) + assert digest_tampered != digest_before, ( + "post-restore tampering must be visible in the digest" + ) + + +async def test_composed_two_layer_round_trip( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """checkpoint_composed/restore_composed round-trip both layers (RFC §3.1). + + A real ManifestEnvironment (``[environment.state] kind=sqlite``) over the + live sandbox: after restore, the file-plane digest AND a sqlite query + both read their pre-checkpoint values, and the node's recorded + StageSnapshot carries both layer refs. + """ + db_path = "/app/t2/env.db" + db_query = f"sqlite3 {db_path} \"SELECT v FROM kv WHERE k='phase';\"" + file_digest = _digest_command(exclude_glob="env.db*") + + async with _live_sandbox(environment_dir, tmp_path) as sandbox: + await _exec_ok( + sandbox, "mkdir -p /app/t2 && echo file-plane > /app/t2/file.txt" + ) + await _exec_ok( + sandbox, + f"sqlite3 {db_path} " + '"CREATE TABLE kv(k TEXT PRIMARY KEY, v TEXT); ' + "INSERT INTO kv VALUES('phase','before');\"", + ) + + # Real env-plane construction: manifest with declared sqlite state + # over the injected live sandbox (the test_manifest_env shape, no + # fakes). No [[services]]: the state plane is what T2 exercises. + manifest = EnvironmentManifest( + name=_ENV_NAME, + image=f"bf__{_ENV_NAME}", + state=StateSpec(kind="sqlite", paths=[db_path]), + ) + environment = ManifestEnvironment(manifest, sandbox=sandbox) + await environment.provision(None) + probe = await environment.readiness() + assert probe.ready, f"environment readiness failed: {probe.error}" + + digest_before = await _exec_value(sandbox, file_digest) + assert await _exec_value(sandbox, db_query) == "before" + + tree = RolloutTree() + node = tree.root + snap = await checkpoint_composed( + node, environment=environment, sandbox=sandbox, stage="pre-verify" + ) + # The recorded StageSnapshot must carry both layer refs. + assert isinstance(snap, StageSnapshot) + assert node.state["snapshot"] is snap + assert snap.stage == "pre-verify" + assert isinstance(snap.environment_ref, StateSnapshot) + assert snap.environment_ref.path.startswith("/tmp/benchflow-snapshots/") + assert snap.sandbox_ref is not None + assert snap.sandbox_ref.provider == "docker" + assert snap.sandbox_ref.ref.startswith(_SNAP_IMAGE_PREFIX) + + # Destructive mutation on both planes. + await _exec_ok( + sandbox, f"sqlite3 {db_path} \"UPDATE kv SET v='after' WHERE k='phase';\"" + ) + await _exec_ok( + sandbox, + "echo scribbled > /app/t2/file.txt && echo junk > /app/t2/junk.txt", + ) + assert await _exec_value(sandbox, db_query) == "after" + digest_mutated = await _exec_value(sandbox, file_digest) + assert digest_mutated != digest_before + + await restore_composed(node, environment=environment, sandbox=sandbox) + + assert await _exec_value(sandbox, file_digest) == digest_before, ( + "composed restore must reproduce the pre-checkpoint file plane" + ) + assert await _exec_value(sandbox, db_query) == "before", ( + "composed restore must roll the declared sqlite state back" + ) + + +async def test_zero_delta_reward_proxy_invariant( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """RFC §5 (a) in miniature: zero-delta restore preserves the reward. + + A deterministic in-sandbox verifier — the reward function's stand-in — + reads PASS on the checkpointed state, FAIL after a destructive mutation + (the detectable known-bad delta), and PASS again after restore. Uses the + composed ops in their sandbox-only shape (``require_layers={"sandbox"}`` + per RFC §3.1: a stateless env + snapshot-capable sandbox can branch). + """ + verifier = ( + 'if [ "$(cat /app/t3/answer.txt 2>/dev/null)" = "42" ]' + " && [ -f /app/t3/nested/marker ];" + " then echo PASS; else echo FAIL; fi" + ) + + async with _live_sandbox(environment_dir, tmp_path) as sandbox: + await _exec_ok( + sandbox, + "mkdir -p /app/t3/nested" + " && echo 42 > /app/t3/answer.txt" + " && touch /app/t3/nested/marker", + ) + assert await _exec_value(sandbox, verifier) == "PASS" + + tree = RolloutTree() + node = tree.root + snap = await checkpoint_composed(node, sandbox=sandbox, stage="pre-verify") + assert snap.environment_ref is None + assert snap.sandbox_ref is not None + + await _exec_ok( + sandbox, "echo 43 > /app/t3/answer.txt && rm -f /app/t3/nested/marker" + ) + assert await _exec_value(sandbox, verifier) == "FAIL", ( + "a known-bad delta must flip the verifier — otherwise PASS-after-" + "restore would prove nothing" + ) + + await restore_composed(node, sandbox=sandbox) + + assert await _exec_value(sandbox, verifier) == "PASS", ( + "zero-delta restore must return the verifier to its pre-snapshot " + "verdict — the reward-equality invariant" + ) + + +async def test_restore_keeps_the_rollout_bind_mounts_and_the_host_sees_writes( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """T2 proof for the live regression: a restored container is still mounted. + + ``DockerSandbox`` bind-mounts the rollout's ``verifier`` / ``agent`` / + ``artifacts`` directories into ``main``. ``restore()`` re-creates that + container outside compose, and a replacement created without those mounts + is silently broken: the verifier's ``reward.txt`` is written inside the + container, the host reads nothing, and the branch child is scored from a + file that does not exist (the run that reported two arms at ``0.00``). + + Two assertions, in the order that matters: the mounts are *declared* on + the restored container (``docker inspect``), and they are *effective* — a + file written to the mounted path inside the container appears on the host, + which is the property the verifier actually depends on. The pre-restore + write proves the mount worked before, so a post-restore failure is + attributable to the restore. + """ + rollout_dir = tmp_path / "mounted-run" + async with _live_sandbox(environment_dir, rollout_dir) as sandbox: + host_verifier_dir = (rollout_dir / "verifier").resolve() + container_verifier_dir = str(SandboxPaths.verifier_dir) + + await _exec_ok(sandbox, f"echo pre > {container_verifier_dir}/reward.txt") + assert (host_verifier_dir / "reward.txt").read_text().strip() == "pre", ( + "the compose-created container must bind-mount the verifier dir — " + "otherwise this test proves nothing about restore" + ) + + image = await sandbox.snapshot() + await sandbox.restore(image) + + mounts = subprocess.run( + [ + "docker", + "inspect", + "-f", + "{{json .Mounts}}", + await sandbox._main_container_id() or "", + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert mounts.returncode == 0, mounts.stderr + destinations = { + mount["Destination"] for mount in json.loads(mounts.stdout or "[]") + } + assert container_verifier_dir in destinations, ( + "the restored container dropped the verifier bind mount: " + f"{sorted(destinations)}" + ) + assert str(SandboxPaths.agent_dir) in destinations + assert str(SandboxPaths.artifacts_dir) in destinations + + # Effective, not just declared: the host must see a post-restore write. + await _exec_ok(sandbox, f"echo 1.0 > {container_verifier_dir}/reward.txt") + assert (host_verifier_dir / "reward.txt").read_text().strip() == "1.0", ( + "a file written to the mounted verifier path after restore must " + "appear on the host — this is the read the verifier depends on" + ) + + # And the live mount check agrees with reality. + assert await sandbox.has_host_mount( + host_dir=host_verifier_dir, container_dir=container_verifier_dir + ) + + +async def test_a_childs_verifier_run_cannot_destroy_the_parents_evidence( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """T2 proof for the shared-mount clobbering, on a real bind mount. + + Because ``restore()`` replays the rollout's bind mounts (the test above), + every branch child writes ``/logs/verifier`` straight into the *parent's* + host directory — and a child whose own rollout paths differ sees + ``has_host_mount() == False``, so it runs ``clear_verifier_output_dir`` + first, whose ``find /logs/verifier -mindepth 1 -exec rm -rf {} +`` empties + the parent's directory before writing anything of its own. That is the + exact command replayed here, in the container, against a live mount. + + The three assertions are the three moves of + :class:`~benchflow.branch_artifacts.MountedArtifacts`: after ``hold`` the + parent's files are out of the blast radius (and the mount really did carry + the deletion through to the host, so this is not a vacuous test); after + ``hand_off`` the child's own output is under its node directory; after + ``release`` the parent's evidence is back at its canonical path with the + transient hold directory gone. + """ + from benchflow.branch_artifacts import ( + MountedArtifacts, + child_mount_dir, + parent_hold_dir, + ) + + rollout_dir = tmp_path / "branch-evidence-run" + async with _live_sandbox(environment_dir, rollout_dir) as sandbox: + host_verifier_dir = (rollout_dir / "verifier").resolve() + container_verifier_dir = str(SandboxPaths.verifier_dir) + + # The parent's own verifier run, written through the mount. + await _exec_ok(sandbox, f"echo 1.0 > {container_verifier_dir}/reward.txt") + await _exec_ok( + sandbox, f"echo 'parent tests passed' > {container_verifier_dir}/stdout.txt" + ) + assert (host_verifier_dir / "reward.txt").read_text().strip() == "1.0" + + holder = MountedArtifacts.hold(run_dir=rollout_dir, parent_id="root") + + # A child's verifier: clear the shared output dir, then score itself. + await _exec_ok( + sandbox, + f"find {container_verifier_dir} -mindepth 1 -exec rm -rf -- {{}} +", + ) + assert not (host_verifier_dir / "reward.txt").exists(), ( + "the child's clear must reach the host through the mount — " + "otherwise this test proves nothing about the clobbering" + ) + assert ( + parent_hold_dir(rollout_dir, "root") / "verifier" / "reward.txt" + ).read_text().strip() == "1.0" + await _exec_ok(sandbox, f"echo 0.0 > {container_verifier_dir}/reward.txt") + + child_dir = child_mount_dir(rollout_dir, "root", "n1") + holder.hand_off(child_dir) + assert (child_dir / "verifier" / "reward.txt").read_text().strip() == "0.0" + + holder.release() + assert (host_verifier_dir / "reward.txt").read_text().strip() == "1.0" + assert ( + host_verifier_dir / "stdout.txt" + ).read_text().strip() == "parent tests passed" + assert not parent_hold_dir(rollout_dir, "root").exists() + + # The mount still works after all that moving — the mount root was + # emptied, never removed. + await _exec_ok(sandbox, f"echo post > {container_verifier_dir}/after.txt") + assert (host_verifier_dir / "after.txt").read_text().strip() == "post" + + +async def test_export_import_round_trip_restores_the_recorded_image_id( + docker_prereqs: None, environment_dir: Path, tmp_path: Path +) -> None: + """--keep-snapshots end-to-end (RFC §3.6): commit → export → destroy → + import restores the exact recorded image. + + Guards "feat(rollout): stage snapshots record their lifetime; + --keep-snapshots on bench eval run with a tested import path" at the T2 + tier: the committed stage image is ``docker save``d through the shared + export machinery, removed (what cleanup's ``compose down --rmi all`` does + to every ``bf-snap-*`` image), then re-imported through the recorded + ``stage_snapshots.json`` — and the loaded image id equals both what the + export recorded from the tar's manifest and what ``docker image inspect`` + reports afterwards. + """ + from benchflow.branch_policy import export_stage_snapshot + from benchflow.snapshot_import import import_stage_snapshots + + async with _live_sandbox(environment_dir, tmp_path / "rollout") as sandbox: + await _exec_ok(sandbox, "echo retained > /app/retained.txt") + image = await sandbox.snapshot() + run_dir = tmp_path / "run" + run_dir.mkdir() + exported = await export_stage_snapshot( + sandbox, sandbox_ref=image.ref, out_dir=run_dir + ) + assert exported["image_id"] is not None, ( + "a real `docker save` tar must yield its config digest" + ) + (run_dir / "stage_snapshots.json").write_text( + json.dumps( + { + "schema_version": 1, + "stages": { + "pre-verify": { + "environment_ref": None, + "sandbox_ref": image.ref, + "layers": ["sandbox"], + "exchanges_completed": None, + "ephemeral": False, + "exported": exported, + } + }, + }, + sort_keys=True, + indent=2, + ) + + "\n" + ) + # The run is over as far as this image is concerned. + subprocess.run( + ["docker", "rmi", "-f", image.ref], + capture_output=True, + timeout=60, + check=True, + ) + + [restored] = import_stage_snapshots(run_dir) + + assert restored.sandbox_ref == image.ref + assert restored.image_id == exported["image_id"] + inspected = subprocess.run( + ["docker", "image", "inspect", "--format", "{{.Id}}", image.ref], + capture_output=True, + text=True, + timeout=30, + check=True, + ) + assert inspected.stdout.strip() == exported["image_id"] diff --git a/tests/test_branch_config_delta.py b/tests/test_branch_config_delta.py new file mode 100644 index 000000000..cbc7e3070 --- /dev/null +++ b/tests/test_branch_config_delta.py @@ -0,0 +1,396 @@ +"""Regression tests for executing config_override branch deltas. + +Guards "feat(branch): execute config_override deltas as fresh rollouts from +env-ready" (docs/rollout-branching-rfc.md §3.3; FrontierPhysics#73). PR number +to be added on submission. + +``config_override`` was schema- and provenance-stable but fail-closed: the +C-axis overlay is deep-merged into the task's resolved config by ``setup()``, +which only a fresh child rollout re-runs. A child forked from the ``env-ready`` +snapshot therefore executes the delta through the same fresh-rollout path the +skills ablation uses (``use_prebuilt_env``): its RolloutConfig is the parent's +with the overlay deep-merged through the EXISTING allowlisted machinery +(``benchflow._utils.config_override``, #790) — same allowlist, same fail-closed +rejection of scorer-touching keys, same content addressing. + +These tests pin that the child really runs under the merged config (the +effective agent timeout its execution receives, not just a recorded label), +that a non-allowlisted key and every non-env-ready branch point fail closed +before any child runs, that the delta composes with ``skill_mode`` and +``injected_prompt`` on one child, and that provenance carries the overlay's +sha256 and its allowlisted keys. + +Unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from benchflow._utils.config_override import overlay_hash +from benchflow.branch_delta import BranchDelta +from benchflow.branch_skill import child_skill_config +from benchflow.rollout import RolloutConfig +from benchflow.rollout_branch import BranchDeltaNotSupported +from tests.test_branch_skill_delta import ( + AGENT_EVENTS, + FakePlanes, + _capture_env_ready, + _fake_verifier, + _parent, + _task_dir, +) + +#: The one controlled change the tests fork under: a 7-second agent budget. +#: Small and distinctive — no TaskConfig default is 7, so an assertion on the +#: child's effective timeout cannot pass by coincidence. +OVERLAY = {"agent": {"timeout_sec": 7}} + + +class TimeoutRecordingPlanes(FakePlanes): + """FakePlanes that records the effective agent timeout each child ran under. + + The overlay's observable effect: ``setup()`` resolves ``self._timeout`` + from the (merged) task config's ``agent.timeout_sec``, and ``execute()`` + hands exactly that to ``execute_prompts``. Recording it here measures the + world the child ran in, not the label its config carries. + """ + + def __init__(self) -> None: + super().__init__() + self.timeouts: list[int] = [] + self.prompts: list[list[str]] = [] + + async def execute_prompts(self, client, session, prompts, timeout, **kwargs): + self.timeouts.append(int(timeout)) + self.prompts.append(list(prompts)) + return list(AGENT_EVENTS), 1 + + +# 1. The delta executes: the child runs fresh under the merged config + + +async def test_config_delta_child_runs_fresh_with_the_merged_config( + tmp_path: Path, monkeypatch +): + """The config ablation, executed for real. + + The control child runs under the task's own agent budget; the delta child + runs under the overlay's — asserted on the timeout its execution actually + received, which only differs if the child's own setup() re-resolved the + merged config. The parent's config is untouched afterwards. + """ + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + assert value == 1.0 + control_timeout, child_timeout = planes.timeouts + assert child_timeout == 7 + assert control_timeout != 7 # the task's own default, not the overlay + assert rollout._config.config_override is None # the parent is untouched + # each child re-installed for itself — the fresh-rollout path, not in-place + assert len(planes.deployments) == 2 + + +async def test_config_delta_child_config_json_records_the_merged_overlay( + tmp_path: Path, monkeypatch +): + """The child is a first-class rollout, so its own config.json carries the + #790 overlay record — keys, sha256, and the patch itself.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + control = json.loads((children_dir / "n2" / "config.json").read_text()) + child = json.loads((children_dir / "n3" / "config.json").read_text()) + assert "config_override" not in control + assert child["config_override"] == { + "keys": ["agent"], + "sha256": overlay_hash(OVERLAY), + "patch": OVERLAY, + } + + +async def test_config_delta_merges_over_the_parents_own_overlay( + tmp_path: Path, monkeypatch +): + """A parent that ran under an overlay of its own hands the child both: + dropping the parent's overlay would silently vary two things and label the + comparison as one delta.""" + planes = TimeoutRecordingPlanes() + rollout = _parent( + _task_dir(tmp_path), + tmp_path, + planes=planes, + config_override={"metadata": {"tags": ["parent-run"]}}, + ) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + control = json.loads((children_dir / "n2" / "config.json").read_text()) + child = json.loads((children_dir / "n3" / "config.json").read_text()) + # the control child inherits the parent's overlay unchanged (zero delta) + assert control["config_override"]["patch"] == {"metadata": {"tags": ["parent-run"]}} + # the delta child runs under parent's overlay + the one recorded change + assert child["config_override"]["patch"] == { + "metadata": {"tags": ["parent-run"]}, + "agent": {"timeout_sec": 7}, + } + assert sorted(child["config_override"]["keys"]) == ["agent", "metadata"] + assert planes.timeouts[1] == 7 + + +# 2. Combinations: one child, one config_override plus another executable field + + +async def test_config_delta_composes_with_skill_mode_on_one_child( + tmp_path: Path, monkeypatch +): + """config_override + skill_mode are both fields on the child's config, so + one child can carry both: it deploys the switched pack AND runs under the + merged budget.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="no-skill"), + BranchDelta(skill_mode="with-skill", config_override=OVERLAY), + ], + ) + + no_skill, with_skill = planes.deployments + assert no_skill["skills_dir"] is None + assert with_skill["skill_files"] == ["demo"] + assert planes.timeouts[0] != 7 # the pure-skill arm keeps the task budget + assert planes.timeouts[1] == 7 # the combined arm runs under the overlay + + +async def test_config_delta_composes_with_an_injected_prompt_on_one_child( + tmp_path: Path, monkeypatch +): + """config_override + injected_prompt: the injection rides the fresh-child + prompt path (the child's continuation prompt) while the overlay rides the + config — neither swallows the other.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + None, + BranchDelta(config_override=OVERLAY, injected_prompt="Follow PLAN.md."), + ], + ) + + assert planes.prompts == [["Solve the task."], ["Follow PLAN.md."]] + assert planes.timeouts == [planes.timeouts[0], 7] + + +# 3. Fail closed before any child runs + + +async def test_a_non_allowlisted_key_fails_before_any_child_runs(tmp_path: Path): + """The same fail-closed allowlist as the run-level overlay (#790): a + scorer-touching patch dies at delta validation, with nothing quiesced, + restored, or forked — never at child setup after a snapshot was consumed.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(ValueError, match="verifier") as excinfo: + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + None, + BranchDelta(config_override={"verifier": {"timeout_sec": 1}}), + ], + ) + + assert "deltas[1].config_override" in str(excinfo.value) + assert rollout._env.restored == [] + assert rollout._environment.restored == [] + assert planes.deployments == [] + assert [node for node in rollout.tree.nodes() if "delta" in node.state] == [] + + +async def test_config_delta_at_a_cursor_branch_names_env_ready(tmp_path: Path): + """A cursor branch forks after the parent's setup() consumed its config — + fail closed naming the boundary that works.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + env = rollout._environment + + with pytest.raises(BranchDeltaNotSupported, match="config_override") as excinfo: + await rollout.branch( + 2, + snapshot_layers={"environment", "sandbox"}, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + assert "env-ready" in str(excinfo.value) + assert "setup()" in str(excinfo.value) + assert env.snapshots == [] + assert rollout._cursor.children == [] + + +async def test_config_delta_at_another_stage_names_env_ready(tmp_path: Path): + """pre-verify is a recorded boundary too — and still the wrong one: the + state the config governs was already consumed by then.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await rollout.mark_stage("pre-verify", snapshot_layers={"environment", "sandbox"}) + + with pytest.raises(BranchDeltaNotSupported, match="'env-ready'") as excinfo: + await rollout.branch_at_stage( + "pre-verify", + 2, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + assert "'pre-verify'" in str(excinfo.value) + assert planes.deployments == [] + + +async def test_config_delta_with_an_explicit_run_child_is_rejected(tmp_path: Path): + """A caller-supplied runner owns the child's execution, so the engine + cannot run setup() under the merged config.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout) + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="run_child"): + await rollout.branch_at_stage( + "env-ready", + 2, + run_child=run_child, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + assert planes.deployments == [] + + +# 4. Provenance: the sha and the allowlisted keys flow into the lineage + + +async def test_provenance_carries_the_overlay_sha_and_keys(tmp_path: Path, monkeypatch): + """Per-child provenance.json and tree.json record the overlay as its + sha256 (#790's exact hash) plus the allowlisted keys it patched — never + the raw patch, which lives only in the child's own config.json.""" + planes = TimeoutRecordingPlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(config_override=OVERLAY)], + ) + + run_dir = rollout._rollout_dir + children_dir = run_dir / "branches" / "root" / "children" + control = json.loads((children_dir / "n2" / "provenance.json").read_text()) + child = json.loads((children_dir / "n3" / "provenance.json").read_text()) + assert control["delta"]["config_override_sha256"] is None + assert "config_override_keys" not in control["delta"] # zero-delta shape + assert child["delta"]["config_override_sha256"] == overlay_hash(OVERLAY) + assert child["delta"]["config_override_keys"] == ["agent"] + assert child["delta_execution"] == "fresh-rollout" + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["n3"]["delta"]["config_override_sha256"] == overlay_hash(OVERLAY) + assert nodes["n3"]["delta"]["config_override_keys"] == ["agent"] + + +# 5. The derived child config + + +def test_child_config_merges_the_delta_overlay_over_the_parents(tmp_path: Path): + """The merge lives in child_skill_config so every fresh child derives its + config the same way: delta keys win, untouched parent keys survive.""" + task = _task_dir(tmp_path) + parent = RolloutConfig.from_legacy( + task_path=task, + agent="oracle", + jobs_dir=tmp_path / "jobs", + config_override={"agent": {"timeout_sec": 30}, "metadata": {"x": 1}}, + ) + + child = child_skill_config( + parent, + skill_mode="no-skill", + jobs_dir=tmp_path / "jobs", + job_name="children", + rollout_name="n2", + config_override={"agent": {"timeout_sec": 7}}, + ) + + assert child.config_override == { + "agent": {"timeout_sec": 7}, + "metadata": {"x": 1}, + } + # the parent's dict is not mutated, and no-delta children inherit it as-is + assert parent.config_override == { + "agent": {"timeout_sec": 30}, + "metadata": {"x": 1}, + } + inherited = child_skill_config( + parent, + skill_mode="no-skill", + jobs_dir=tmp_path / "jobs", + job_name="children", + rollout_name="n3", + ) + assert inherited.config_override == parent.config_override + + +def test_provenance_dict_records_keys_only_when_the_field_is_set() -> None: + """``config_override_keys`` joins the provenance block only for a set + overlay, so every other delta keeps the exact dict shape the RFC pinned.""" + unset: dict[str, Any] = BranchDelta(skill_mode="no-skill").provenance_dict() + assert "config_override_keys" not in unset + provenance = BranchDelta( + config_override={"metadata": {}, "agent": {"timeout_sec": 7}} + ).provenance_dict() + assert provenance["config_override_keys"] == ["agent", "metadata"] diff --git a/tests/test_branch_deltas_lineage.py b/tests/test_branch_deltas_lineage.py new file mode 100644 index 000000000..474ff4dcf --- /dev/null +++ b/tests/test_branch_deltas_lineage.py @@ -0,0 +1,606 @@ +"""Regression tests for branch deltas + lineage artifacts. + +Guards the deltas + lineage layer ("feat(branch): per-child deltas + branch +lineage artifacts"; docs/rollout-branching-rfc.md WS-2; FrontierPhysics#73). +PR number to be added on submission. + +A branch child's delta (RFC §3.3) is the recorded exactly-one-controlled-change +it runs under: v1 executes ``injected_prompt`` (the child's user-visible first +message), while ``environment_ref`` / ``config_override`` / ``skill_mode`` are +schema-and-provenance-stable but fail closed until the child-as-fresh-rollout +follow-on. Lineage (RFC §3.4) makes a branch leave evidence: a deterministic +``tree.json`` (each node carrying the delta provenance the engine attached at +fork time) plus per-child +``branches//children//provenance.json`` / +``reward.json``, with artifact-write failures isolated from the branch result. + +These are unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from datetime import datetime +from pathlib import Path + +import pytest + +from benchflow._utils.content_address import sha256_prefixed +from benchflow.branch import StageSnapshot +from benchflow.branch_delta import BranchDelta +from benchflow.branch_lineage import child_provenance, serialize_tree +from benchflow.environment.protocol import StateSnapshot +from benchflow.rollout import _TERMINAL_PHASES, Rollout, RolloutConfig, Scene +from benchflow.rollout_branch import ( + _UNSUPPORTED_DELTA_FIELDS, + BranchDeltaNotSupported, +) +from benchflow.sandbox.protocol import SandboxImage +from benchflow.trajectories.tree import RolloutTree, Step + + +class FakeEnvironment: + """Environment-plane stand-in recording snapshot/restore calls.""" + + def __init__(self) -> None: + self.snapshots: list[StateSnapshot] = [] + self.restored: list[StateSnapshot] = [] + + async def snapshot(self) -> StateSnapshot: + snap = StateSnapshot(id=f"env-snap-{len(self.snapshots) + 1}", path="/tmp/x") + self.snapshots.append(snap) + return snap + + async def restore(self, snap: StateSnapshot) -> None: + self.restored.append(snap) + + +class FakeSnapSandbox: + """Snapshot-capable Sandbox stand-in.""" + + supports_snapshot = True + + def __init__(self) -> None: + self.snapshots: list[SandboxImage] = [] + self.restored: list[SandboxImage] = [] + + async def snapshot(self, name: str | None = None) -> SandboxImage: + img = SandboxImage(provider="fake", ref=f"bf-snap-{len(self.snapshots) + 1}") + self.snapshots.append(img) + return img + + async def restore(self, image: SandboxImage) -> None: + self.restored.append(image) + + +def _rollout(tmp_path: Path) -> Rollout: + return Rollout( + RolloutConfig(task_path=tmp_path / "task", scenes=[Scene.single(agent="dummy")]) + ) + + +def _fake_agent_boundary(monkeypatch, received_prompts: list[list[str] | None]): + """Fake connect/execute/verify at the class boundary, recording prompts.""" + + async def fake_connect(self): + self._acp_client = object() + + async def fake_disconnect(self): + self._acp_client = None + + async def fake_execute(self, prompts=None, *, node=None): + received_prompts.append(prompts) + return [], 0 + + async def fake_verify(self): + return {"reward": 1.0} + + monkeypatch.setattr(Rollout, "connect", fake_connect) + monkeypatch.setattr(Rollout, "disconnect", fake_disconnect) + monkeypatch.setattr(Rollout, "execute", fake_execute) + monkeypatch.setattr(Rollout, "verify", fake_verify) + + +# 1. BranchDelta schema: validation and content-addressed provenance + + +def test_branch_delta_rejects_an_unknown_skill_mode(): + """skill_mode outside {no-skill, with-skill} fails at construction.""" + with pytest.raises(ValueError, match="skill_mode"): + BranchDelta(skill_mode="self-gen") + with pytest.raises(ValueError, match="skill_mode"): + BranchDelta(skill_mode="skills-on") + + +def test_branch_delta_is_empty_only_when_every_field_is_unset(): + assert BranchDelta().is_empty + assert not BranchDelta(injected_prompt="plan").is_empty + assert not BranchDelta(skill_mode="no-skill").is_empty + + +def test_provenance_hashes_are_content_addressed_and_key_order_stable(): + """Same config_override content -> same sha; different content -> different. + + The hash is over canonical JSON (sort_keys=True), so key order must not + change it — the run-level overlay's exact hashing (#790). + """ + a = BranchDelta(config_override={"agent": {"timeout_sec": 60}, "metadata": {}}) + b = BranchDelta(config_override={"metadata": {}, "agent": {"timeout_sec": 60}}) + c = BranchDelta(config_override={"agent": {"timeout_sec": 61}}) + + sha_a = a.provenance_dict()["config_override_sha256"] + sha_b = b.provenance_dict()["config_override_sha256"] + sha_c = c.provenance_dict()["config_override_sha256"] + + assert sha_a == sha_b + assert sha_a != sha_c + assert sha_a.startswith("sha256:") + + +def test_provenance_records_the_prompt_hash_never_the_prompt_text(): + """No raw prompt content in provenance — only its sha256 digest.""" + prompt = "SECRET ORACLE PLAN: mine the sqlite db first." + delta = BranchDelta(injected_prompt=prompt, skill_mode="with-skill") + + prov = delta.provenance_dict() + + assert prov["injected_prompt_sha256"] == sha256_prefixed(prompt.encode()) + assert prov["skill_mode"] == "with-skill" + assert prov["environment_ref"] is None + assert prov["config_override_sha256"] is None + assert prompt not in json.dumps(prov) + + +# 2. Engine: delta validation fails closed before anything runs + + +async def test_deltas_length_must_match_n(tmp_path: Path): + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="one entry per child"): + await rollout.branch(2, run_child=run_child, deltas=[None]) + + +@pytest.mark.parametrize( + "delta", + [ + BranchDelta(environment_ref="env0@outage"), + BranchDelta(config_override={"agent": {"timeout_sec": 60}}), + BranchDelta(skill_mode="with-skill"), + ], + ids=["environment_ref", "config_override", "skill_mode"], +) +async def test_unsupported_delta_fields_fail_closed_before_any_child_runs( + tmp_path: Path, delta: BranchDelta +): + """environment_ref/config_override/skill_mode raise BranchDeltaNotSupported + at a *cursor* branch. + + Every one of them executes only from the ``env-ready`` stage snapshot — + the fresh-child path — so at the cursor each gate names that boundary in + a typed error that also names the field, and fires before + quiesce/checkpoint: no snapshot taken, no child forked or run. + """ + rollout = _rollout(tmp_path) + env = FakeEnvironment() + rollout._environment = env + parent = rollout._cursor + ran: list[str] = [] + + async def run_child(child): + ran.append(child.id) + return 1.0 + + field_name = next( + name + for name in ("environment_ref", "config_override", "skill_mode") + if getattr(delta, name) is not None + ) + with pytest.raises(BranchDeltaNotSupported, match=field_name) as excinfo: + await rollout.branch(2, run_child=run_child, deltas=[None, delta]) + + assert "use_prebuilt_env" in str(excinfo.value) + assert isinstance(excinfo.value, NotImplementedError) + # fail closed BEFORE any child ran: nothing snapshotted, nothing forked + assert ran == [] + assert env.snapshots == [] + assert parent.children == [] + + +async def test_injected_prompt_with_an_explicit_run_child_is_rejected( + tmp_path: Path, +): + """An injected_prompt delta needs the default runner to deliver it — a + caller-supplied run_child owns the child's prompts, so the combination + fails closed instead of silently not delivering the injection.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="injected_prompt"): + await rollout.branch( + 2, + run_child=run_child, + deltas=[BranchDelta(injected_prompt="go"), None], + ) + + +# 3. Engine: injected_prompt executes and is provenance-recorded + + +async def test_injected_prompt_becomes_the_childs_continuation_prompt( + tmp_path: Path, monkeypatch +): + """The child runs with the injected prompt as its user-visible message; + the zero-delta sibling keeps the rollout's resolved prompts (None here).""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + received: list[list[str] | None] = [] + _fake_agent_boundary(monkeypatch, received) + + prompt = "Execute PLAN.md verbatim; do not re-research." + value = await rollout.branch(2, deltas=[None, BranchDelta(injected_prompt=prompt)]) + + assert value == 1.0 + assert received == [None, [prompt]] + + +async def test_injected_prompt_provenance_records_the_hash_not_the_text( + tmp_path: Path, monkeypatch +): + """branches//children//provenance.json carries the sha only.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + received: list[list[str] | None] = [] + _fake_agent_boundary(monkeypatch, received) + + prompt = "Execute PLAN.md verbatim; do not re-research." + await rollout.branch(2, deltas=[None, BranchDelta(injected_prompt=prompt)]) + + children_dir = run_dir / "branches" / "root" / "children" + prov = json.loads((children_dir / "n2" / "provenance.json").read_text()) + assert prov["delta"]["injected_prompt_sha256"] == sha256_prefixed(prompt.encode()) + assert prompt not in (children_dir / "n2" / "provenance.json").read_text() + zero = json.loads((children_dir / "n1" / "provenance.json").read_text()) + assert zero["delta"]["injected_prompt_sha256"] is None + + +# 4. Lineage: tree.json structure, determinism, and failure isolation + + +async def test_tree_json_records_nodes_snapshot_refs_rewards_and_deltas( + tmp_path: Path, monkeypatch +): + """tree.json after a composed branch: schema_version 1, one entry per + node, the branch point carrying both layers' snapshot refs and V(parent), + each child carrying its reward and delta provenance.""" + rollout = _rollout(tmp_path) + env, sandbox = FakeEnvironment(), FakeSnapSandbox() + rollout._environment = env + rollout._env = sandbox + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + received: list[list[str] | None] = [] + _fake_agent_boundary(monkeypatch, received) + + prompt = "Use the oracle plan." + await rollout.branch( + 2, + snapshot_layers={"environment", "sandbox"}, + deltas=[None, BranchDelta(injected_prompt=prompt)], + ) + + payload = json.loads((run_dir / "tree.json").read_text()) + assert payload["schema_version"] == 1 + assert payload["cut_point"] is None + + nodes = {node["id"]: node for node in payload["nodes"]} + root = nodes["root"] + assert root["parent"] is None + assert root["snapshot"] == {"environment": "env-snap-1", "sandbox": "bf-snap-1"} + assert root["value"] == 1.0 + + children = [node for node in payload["nodes"] if node["parent"] == "root"] + assert len(children) == 2 + assert [child["reward"] for child in children] == [1.0, 1.0] + assert children[0]["delta"]["injected_prompt_sha256"] is None + assert children[1]["delta"]["injected_prompt_sha256"] == sha256_prefixed( + prompt.encode() + ) + assert prompt not in (run_dir / "tree.json").read_text() + + +def test_serialize_tree_is_deterministic_across_calls(tmp_path: Path): + """Two serializations of the same tree are byte-identical — sorted keys, + trailing newline, no wall-clock timestamps. Node entries come from each + node's own recorded state (snapshot / reward / delta), never from + positional guessing.""" + tree = RolloutTree() + snap = StageSnapshot( + environment_ref=StateSnapshot(id="env-snap-1", path="/tmp/x"), + sandbox_ref=SandboxImage(provider="fake", ref="bf-snap-1"), + stage="env-ready", + ) + tree.root.state["snapshot"] = snap + tree.root.state["value"] = 0.5 + deltas = [None, BranchDelta(injected_prompt="plan")] + for reward, delta in zip((0.0, 1.0), deltas, strict=True): + child = tree.attach(tree.root) + child.state["reward"] = reward + child.state["delta"] = ( + delta if delta is not None else BranchDelta() + ).provenance_dict() + + first_dir, second_dir = tmp_path / "a", tmp_path / "b" + first_dir.mkdir() + second_dir.mkdir() + first = serialize_tree(tree, run_dir=first_dir) + second = serialize_tree(tree, run_dir=second_dir) + + text = first.read_text() + assert text == second.read_text() + assert text.endswith("\n") + # the stage tag comes from the StageSnapshot itself + root_entry = next( + node for node in json.loads(text)["nodes"] if node["id"] == "root" + ) + assert root_entry["stage"] == "env-ready" + + +async def test_artifact_write_failure_never_corrupts_the_branch_result( + tmp_path: Path, +): + """An unwritable run dir (here: a plain file) is logged and swallowed — + the branch still returns V(parent) and the tree still grew.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + not_a_dir = tmp_path / "run" + not_a_dir.write_text("occupied") + rollout._rollout_dir = not_a_dir + parent = rollout._cursor + + returns = iter([0.0, 1.0]) + + async def run_child(child): + return next(returns) + + value = await rollout.branch(2, run_child=run_child) + + assert value == 0.5 + assert len(parent.children) == 2 + assert not_a_dir.read_text() == "occupied" + + +async def test_children_artifacts_are_written_and_parseable(tmp_path: Path): + """branches//children// gets provenance.json + (kind benchflow-branch, legacy env-only snapshot ref, cursor-tagged stage) + and reward.json.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + + returns = iter([0.25, 0.75]) + + async def run_child(child): + return next(returns) + + await rollout.branch(2, run_child=run_child) + + for child_id, expected in [("n1", 0.25), ("n2", 0.75)]: + child_dir = run_dir / "branches" / "root" / "children" / child_id + prov = json.loads((child_dir / "provenance.json").read_text()) + assert prov["kind"] == "benchflow-branch" + assert prov["parent_rollout"] == str(run_dir) + assert prov["branch_stage"] == "cursor:root" + assert prov["snapshot_ref"] == {"environment": "env-snap-1", "sandbox": None} + assert prov["cut_point"] is None + assert prov["delta"] == BranchDelta().provenance_dict() + assert json.loads((child_dir / "reward.json").read_text()) == { + "reward": expected + } + + +async def test_second_branch_at_the_same_parent_never_misattributes( + tmp_path: Path, monkeypatch +): + """Two sequential branch() calls at one parent keep both events' evidence. + + Deltas are attached to each child node at fork time, so the second event + cannot claim the first event's children; artifacts are namespaced by node + id, so nothing is overwritten. + """ + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + received: list[list[str] | None] = [] + _fake_agent_boundary(monkeypatch, received) + + first, second = "first-event plan", "second-event plan" + await rollout.branch(2, deltas=[None, BranchDelta(injected_prompt=first)]) + await rollout.branch(2, deltas=[BranchDelta(injected_prompt=second), None]) + + children_dir = run_dir / "branches" / "root" / "children" + shas = { + child_id: json.loads((children_dir / child_id / "provenance.json").read_text())[ + "delta" + ]["injected_prompt_sha256"] + for child_id in ("n1", "n2", "n3", "n4") + } + # first event: n1 zero-delta, n2 first prompt; second event: n3 second + # prompt, n4 zero-delta — no overwrites, no misattribution. + assert shas == { + "n1": None, + "n2": sha256_prefixed(first.encode()), + "n3": sha256_prefixed(second.encode()), + "n4": None, + } + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["n2"]["delta"]["injected_prompt_sha256"] == sha256_prefixed( + first.encode() + ) + assert nodes["n3"]["delta"]["injected_prompt_sha256"] == sha256_prefixed( + second.encode() + ) + assert nodes["n1"]["delta"]["injected_prompt_sha256"] is None + assert nodes["n4"]["delta"]["injected_prompt_sha256"] is None + + +async def test_two_branch_points_keep_per_node_delta_provenance( + tmp_path: Path, monkeypatch +): + """A tree with two branch points records every child's delta on its node + — the old unique-branch-point fallback dropped all delta provenance + here.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + received: list[list[str] | None] = [] + _fake_agent_boundary(monkeypatch, received) + + first, second = "root-event plan", "deep-event plan" + await rollout.branch(2, deltas=[None, BranchDelta(injected_prompt=first)]) + # continue linearly past the branch, then branch again at the new cursor + rollout._cursor = rollout._tree.advance(rollout._cursor, Step(id="s1")) + await rollout.branch(2, deltas=[BranchDelta(injected_prompt=second), None]) + + # first event at root (children n1, n2); linear node n3; second event at + # n3 (children n4, n5) + root_children = run_dir / "branches" / "root" / "children" + deep_children = run_dir / "branches" / "n3" / "children" + assert json.loads((root_children / "n2" / "provenance.json").read_text())["delta"][ + "injected_prompt_sha256" + ] == sha256_prefixed(first.encode()) + assert json.loads((deep_children / "n4" / "provenance.json").read_text())["delta"][ + "injected_prompt_sha256" + ] == sha256_prefixed(second.encode()) + + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["n2"]["delta"]["injected_prompt_sha256"] == sha256_prefixed( + first.encode() + ) + assert nodes["n4"]["delta"]["injected_prompt_sha256"] == sha256_prefixed( + second.encode() + ) + assert nodes["n1"]["delta"]["injected_prompt_sha256"] is None + assert nodes["n5"]["delta"]["injected_prompt_sha256"] is None + assert "delta" not in nodes["n3"] # the linear node carries no delta + + +def test_child_provenance_shape_matches_the_rfc(): + """child_provenance builds the RFC §3.4 dict for a composed snapshot.""" + snap = StageSnapshot( + environment_ref=StateSnapshot(id="env-snap-9", path="/tmp/x"), + sandbox_ref=SandboxImage(provider="fake", ref="bf-snap-9"), + stage="env-ready", + ) + prov = child_provenance( + "/runs/parent-1", + branch_stage="env-ready", + snapshot=snap, + delta=BranchDelta(skill_mode="no-skill"), + ) + assert prov == { + "kind": "benchflow-branch", + "parent_rollout": "/runs/parent-1", + "branch_stage": "env-ready", + "snapshot_ref": {"sandbox": "bf-snap-9", "environment": "env-snap-9"}, + "cut_point": None, + "delta": { + "environment_ref": None, + "config_override_sha256": None, + "skill_mode": "no-skill", + "injected_prompt_sha256": None, + }, + } + + +# 5. Engine: the derived unsupported-field set fails closed by construction + + +def test_unsupported_delta_fields_are_derived_from_the_schema(): + """The blocklist is derived at import (all BranchDelta fields minus the + executable set), so a future BranchDelta field is unsupported-by-default + — the engine fails closed on it instead of silently ignoring it. Every + current field has since gained an execution path (a fresh child rollout + from the env-ready snapshot) and is gated per branch point instead — see + tests/test_branch_skill_delta.py, tests/test_branch_config_delta.py and + tests/test_branch_environment_delta.py — leaving the derived set empty + but still load-bearing for the next schema field.""" + assert set(_UNSUPPORTED_DELTA_FIELDS) == set() + + +# 6. "branched" is a terminal phase and the result is real + + +async def test_branch_first_result_is_none_without_setup(tmp_path: Path): + """A branch-first rollout that never ran setup() reads result as None. + + "branched" is terminal, but with no run directory there is nothing to + build result artifacts in — the property stays graceful (the pre-branch + contract) instead of raising RuntimeError from _require_rollout_dir. + """ + assert "branched" in _TERMINAL_PHASES + + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + + async def run_child(child): + return 1.0 + + assert rollout.result is None # pre-branch: no terminal phase yet + await rollout.branch(2, run_child=run_child) + + assert rollout._phase == "branched" + assert rollout.result is None # no setup() -> no run dir -> graceful None + + +async def test_branched_result_surfaces_the_branch_aggregate(tmp_path: Path): + """A set-up rollout's branched result carries V(cursor), not rewards=None. + + branch() restores the parent's linear state, so ``_rewards`` rolls back + to its pre-branch value even though the children verified; the built + result must surface the aggregate the engine recorded on the branch + point as ``rewards={"reward": , "source": "branch_aggregate"}`` — + exercised through the real result/_build_result path (no monkeypatch). + """ + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + rollout._started_at = datetime.now() + + returns = iter([0.25, 0.75]) + + async def run_child(child): + return next(returns) + + await rollout.branch(2, run_child=run_child) + + result = rollout.result + assert result is not None + assert result.rewards == {"reward": 0.5, "source": "branch_aggregate"} + persisted = json.loads((run_dir / "result.json").read_text()) + assert persisted["rewards"] == {"reward": 0.5, "source": "branch_aggregate"} diff --git a/tests/test_branch_environment_delta.py b/tests/test_branch_environment_delta.py new file mode 100644 index 000000000..00a30a4f4 --- /dev/null +++ b/tests/test_branch_environment_delta.py @@ -0,0 +1,480 @@ +"""Regression tests for executing service-level environment_ref branch deltas. + +Guards "feat(branch): execute service-level environment_ref deltas from +env-ready" (docs/rollout-branching-rfc.md §3.3 — the ``env0@prod`` vs +``env0@outage`` tool-outage perturbation; FrontierPhysics#73). PR number to be +added on submission. + +``environment_ref`` was schema- and provenance-stable but fail-closed. The +executable slice is *service topology*: the env-ready snapshot commits the +**parent's** container, and restoring it kills every framework-started service +with it — so for a manifest pair sharing the same image with +``owns_lifecycle = false``, what the fresh child provisions over the restore +IS the child's service set, and swapping the manifest executes the delta +soundly (the RFC's "service stop/start bracketing around state restore", +§3.1/§3.3). Everything outside that slice fails closed before anything is +quiesced: an image-changing manifest breaks the restore-the-parent-container +premise (it would need a rebuild path, which contradicts branching from a +snapshot), and an entrypoint-owned lifecycle starts whatever the image bakes +in, so a recorded service delta would not be enforced. + +The provisioning step also covers the control arm: a zero-delta child of a +manifest-bound parent re-provisions the *parent's* manifest, so the baseline +arm of a framework-started environment no longer scores a world whose services +all died with the container restore. + +Unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from benchflow.branch_delta import BranchDelta +from benchflow.branch_skill import ( + BranchEnvironmentImageConflict, + resolve_environment_ref_delta, +) +from benchflow.environment.manifest import load_manifest +from benchflow.rollout_branch import BranchDeltaNotSupported +from tests.test_branch_skill_delta import ( + FakePlanes, + _capture_env_ready, + _fake_verifier, + _parent, + _task_dir, +) + +IMAGE = "bf-env0:2026.08" + +_SERVICE_BLOCKS = { + "gmail": '[[environment.services]]\nname = "gmail"\ncommand = "claw-gmail --port 8081"\nport = 8081\n', + "gcal": '[[environment.services]]\nname = "gcal"\ncommand = "claw-gcal --port 8082"\nport = 8082\n', +} + + +def _manifest_toml( + *, image: str = IMAGE, services: tuple[str, ...], owns_lifecycle: bool = False +) -> str: + blocks = "\n".join(_SERVICE_BLOCKS[name] for name in services) + return ( + "[environment]\n" + 'name = "env0"\n' + f'image = "{image}"\n' + f"owns_lifecycle = {str(owns_lifecycle).lower()}\n\n" + blocks + ) + + +def _registry(tmp_path: Path, monkeypatch) -> Path: + """A local env registry holding the documented prod/outage pair plus the + two unbranchable variants the gates exist for.""" + registry = tmp_path / "registry" + registry.mkdir() + (registry / "env0@prod.toml").write_text(_manifest_toml(services=("gmail", "gcal"))) + (registry / "env0@outage.toml").write_text(_manifest_toml(services=("gmail",))) + (registry / "env0@rebuilt.toml").write_text( + _manifest_toml(image="bf-env0:rebuilt", services=("gmail",)) + ) + (registry / "env0@baked.toml").write_text( + _manifest_toml(services=(), owns_lifecycle=True) + ) + monkeypatch.setenv("BENCHFLOW_ENV_REGISTRY", str(registry)) + return registry + + +class FakeManifestEnvironment: + """Environment-plane stand-in recording what was provisioned, and when.""" + + ready = True + + def __init__(self, manifest, calls: list[str]) -> None: + self.manifest = manifest + self.calls = calls + self.provision_ctx: Any = None + self.torn_down = False + + async def provision(self, ctx: Any) -> Any: + self.provision_ctx = ctx + self.calls.append( + f"provision:{','.join(s.name for s in self.manifest.services)}" + ) + return SimpleNamespace(name=self.manifest.name, endpoints={}) + + async def readiness(self) -> Any: + self.calls.append("readiness") + return SimpleNamespace( + ready=self.ready, + checked=[f"http://localhost:{s.port}" for s in self.manifest.services], + error=None if self.ready else "gmail never responded", + ) + + async def teardown(self) -> None: + self.torn_down = True + + +class EnvPlanes(FakePlanes): + """FakePlanes with the manifest-environment factory the child provisions + through, recording every plane it builds and the order of the calls.""" + + environment_cls = FakeManifestEnvironment + + def __init__(self) -> None: + super().__init__() + self.environments: list[FakeManifestEnvironment] = [] + self.calls: list[str] = [] + + def manifest_environment(self, manifest, *, sandbox): + environment = self.environment_cls(manifest, self.calls) + self.environments.append(environment) + return environment + + async def deploy_skills(self, env, task_path, skills_dir, *args, **kwargs): + self.calls.append("deploy_skills") + await super().deploy_skills(env, task_path, skills_dir, *args, **kwargs) + + +def _manifest_parent(tmp_path: Path, monkeypatch, *, planes: EnvPlanes): + """A manifest-bound parent positioned as if it had run past env-ready.""" + _registry(tmp_path, monkeypatch) + return _parent( + _task_dir(tmp_path), + tmp_path, + planes=planes, + environment_manifest=load_manifest("env0@prod"), + ) + + +# 1. The delta executes: the child provisions the child manifest's service set + + +async def test_environment_delta_child_provisions_the_child_manifests_services( + tmp_path: Path, monkeypatch +): + """The tool-outage ablation, executed for real. + + The control child re-provisions the parent's own manifest (both services); + the delta child provisions the outage manifest's subset over the same + restored container — asserted on the provisioned service sets, which is + the world each arm actually ran in, not a recorded label. + """ + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + assert value == 1.0 + control_env, delta_env = planes.environments + assert [s.name for s in control_env.manifest.services] == ["gmail", "gcal"] + assert [s.name for s in delta_env.manifest.services] == ["gmail"] + assert control_env.provision_ctx == {"task_id": "task"} + assert delta_env.provision_ctx == {"task_id": "task"} + # both children's planes were torn down by their own cleanup() + assert control_env.torn_down and delta_env.torn_down + # ... and the parent's own config still binds the prod manifest + assert [s.name for s in rollout._config.environment_manifest.services] == [ + "gmail", + "gcal", + ] + + +async def test_environment_delta_child_config_json_records_the_swapped_manifest( + tmp_path: Path, monkeypatch +): + """The child is a first-class rollout, so its own config.json records the + manifest it actually provisioned — name, image, and the differing service + set — while the control child records the parent's.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + control = json.loads((children_dir / "n2" / "config.json").read_text()) + child = json.loads((children_dir / "n3" / "config.json").read_text()) + assert control["environment_manifest"]["services"] == ["gmail", "gcal"] + assert child["environment_manifest"]["services"] == ["gmail"] + assert child["environment_manifest"]["image"] == IMAGE + assert child["environment_manifest"]["owns_lifecycle"] is False + + +async def test_child_environment_is_provisioned_after_restore_before_install( + tmp_path: Path, monkeypatch +): + """The RFC §3.1 bracket: services come up after the container roll-back + (which killed the previous child's) and before install_agent(), mirroring + the linear lifecycle's start() -> install_agent() order.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + sandbox = rollout._env + original_restore = sandbox.restore + + async def recording_restore(image): + planes.calls.append("sandbox.restore") + await original_restore(image) + + sandbox.restore = recording_restore + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + assert planes.calls == [ + "sandbox.restore", + "provision:gmail,gcal", + "readiness", + "deploy_skills", + "sandbox.restore", + "provision:gmail", + "readiness", + "deploy_skills", + ] + + +async def test_zero_delta_children_reprovision_the_parents_manifest( + tmp_path: Path, monkeypatch +): + """The control arm's honesty: restoring the container kills every + framework-started service, so a zero-delta child of a manifest-bound + parent re-provisions the parent's own service set — without this, the + baseline arm of an outage comparison scores a dead world and calls it + prod.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage("env-ready", 2) + + assert [[s.name for s in env.manifest.services] for env in planes.environments] == [ + ["gmail", "gcal"], + ["gmail", "gcal"], + ] + + +async def test_a_failed_readiness_gate_fails_the_child_loudly( + tmp_path: Path, monkeypatch +): + """A child whose swapped environment never became ready must not score: + the gate raises (the same contract as start()), the fork ends, and the + child's own cleanup still tears its plane down.""" + + class NotReadyEnvironment(FakeManifestEnvironment): + ready = False + + class NotReadyPlanes(EnvPlanes): + environment_cls = NotReadyEnvironment + + planes = NotReadyPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + with pytest.raises(RuntimeError, match="not ready"): + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + assert planes.environments[0].torn_down + assert planes.deployments == [] # readiness gates before install_agent() + + +# 2. Provenance: the ref is recorded verbatim on the child + + +async def test_provenance_records_the_environment_ref_verbatim( + tmp_path: Path, monkeypatch +): + """Per-child provenance.json and tree.json carry the registry ref exactly + as requested (it is already content-addressed by the registry), plus the + fresh-rollout execution marker.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + control = json.loads((children_dir / "n2" / "provenance.json").read_text()) + child = json.loads((children_dir / "n3" / "provenance.json").read_text()) + assert control["delta"]["environment_ref"] is None + assert child["delta"]["environment_ref"] == "env0@outage" + assert child["delta_execution"] == "fresh-rollout" + nodes = { + node["id"]: node + for node in json.loads((rollout._rollout_dir / "tree.json").read_text())[ + "nodes" + ] + } + assert nodes["n3"]["delta"]["environment_ref"] == "env0@outage" + + +# 3. The boundary: everything outside the services slice fails closed + + +async def test_an_image_changing_manifest_fails_closed_with_the_typed_error( + tmp_path: Path, monkeypatch +): + """The restore-the-parent-container premise: the snapshot commits the + parent's image, so a manifest naming a different image needs a rebuild + path — which contradicts branching from a snapshot. The typed error names + both images and fires before anything is restored or forked.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(BranchEnvironmentImageConflict) as excinfo: + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@rebuilt")], + ) + + message = str(excinfo.value) + assert IMAGE in message and "bf-env0:rebuilt" in message + assert "rebuild" in message + assert "contradicts branching from a snapshot" in message + assert isinstance(excinfo.value, BranchDeltaNotSupported) + assert planes.environments == [] + assert rollout._env.restored == [] + assert [node for node in rollout.tree.nodes() if "delta" in node.state] == [] + + +async def test_an_entrypoint_owned_lifecycle_fails_closed(tmp_path: Path, monkeypatch): + """Same image, but the child manifest hands the lifecycle to the image + entrypoint: the restored container starts whatever the image bakes in, so + the framework cannot enforce a manifest-declared service difference — the + delta would be recorded but not executed. Fail closed instead.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(BranchDeltaNotSupported, match="owns_lifecycle"): + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@baked")], + ) + + assert planes.environments == [] + assert rollout._env.restored == [] + + +async def test_an_unresolvable_ref_fails_closed_naming_the_registry_problem( + tmp_path: Path, monkeypatch +): + """A ref the registry cannot resolve is a delta that cannot be recorded + honestly, let alone executed — fail closed with the resolution error.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(BranchDeltaNotSupported, match="does not resolve"): + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(environment_ref="env0@nonexistent")], + ) + + assert planes.environments == [] + + +async def test_environment_delta_at_a_cursor_branch_names_env_ready( + tmp_path: Path, monkeypatch +): + """A cursor branch keeps the parent's provisioned services alive across + the fork, so the swap could only be recorded, never enforced — fail closed + naming the boundary that works.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + env = rollout._environment + + with pytest.raises(BranchDeltaNotSupported, match="environment_ref") as excinfo: + await rollout.branch( + 2, + snapshot_layers={"environment", "sandbox"}, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + assert "env-ready" in str(excinfo.value) + assert env.snapshots == [] + assert rollout._cursor.children == [] + + +async def test_environment_delta_with_an_explicit_run_child_is_rejected( + tmp_path: Path, monkeypatch +): + """A caller-supplied runner owns the child's execution, so the engine + cannot provision the child manifest's services.""" + planes = EnvPlanes() + rollout = _manifest_parent(tmp_path, monkeypatch, planes=planes) + await _capture_env_ready(rollout) + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="run_child"): + await rollout.branch_at_stage( + "env-ready", + 2, + run_child=run_child, + deltas=[None, BranchDelta(environment_ref="env0@outage")], + ) + + assert planes.environments == [] + + +# 4. The resolution helper is the single source of the boundary + + +def test_resolve_environment_ref_delta_returns_the_outage_manifest( + tmp_path: Path, monkeypatch +): + """The happy path of the shared gate (engine validation, the child runner, + and the ablate pre-flight all call this one function).""" + _registry(tmp_path, monkeypatch) + parent = load_manifest("env0@prod") + + child = resolve_environment_ref_delta(parent, "env0@outage") + + assert child.image == parent.image + assert [s.name for s in child.services] == ["gmail"] + assert child.owns_lifecycle is False + + +def test_resolve_environment_ref_delta_requires_a_manifest_bound_parent( + tmp_path: Path, monkeypatch +): + """No parent manifest means no Environment plane to swap: the child would + run a world its snapshot never contained.""" + _registry(tmp_path, monkeypatch) + + with pytest.raises(BranchDeltaNotSupported, match="no environment manifest"): + resolve_environment_ref_delta(None, "env0@outage", subject="arm 'env:...'") diff --git a/tests/test_branch_skill_delta.py b/tests/test_branch_skill_delta.py new file mode 100644 index 000000000..195678ed8 --- /dev/null +++ b/tests/test_branch_skill_delta.py @@ -0,0 +1,1147 @@ +"""Regression tests for executing branch children at the env-ready boundary. + +Guards the skill-delta execution path ("feat(branch): execute skill_mode deltas +as fresh rollouts from the env-ready snapshot"; docs/rollout-branching-rfc.md +§3.3 / WS-4b; FrontierPhysics#73) and its generalization to every child of that +boundary. PR number to be added on submission. + +``skill_mode`` was schema-stable but fail-closed: skills are deployed by +``install_agent()``, so a branch at the cursor forks a world that has already +resolved the question. The ``env-ready`` stage snapshot (WS-4a) is taken before +``install_agent()``, so a child restored from it can re-run installation under +the switched mode — as a *fresh* Rollout over the restored sandbox +(``use_prebuilt_env``, #388). These tests pin that the child really varies +skills (the staged build context and the deploy call, not just a recorded +label), that every other branch point still fails closed, and that lineage +records the effective mode and the fresh-rollout execution. + +Section 5 pins the same property for children that carry *no* skill delta. The +seam is the stage, not the delta: ``env-ready`` precedes ``install_agent()``, +so a child restored there has no agent binary, no sandbox user, no seeded +verifier workspace, no lockdown and no skill pack. An in-place child forked +there — the shape an ``inject:`` ablation arm used to take — either dies +connecting to an agent the restore deleted or scores a world missing everything +installation deploys, and reports it as an ordinary single-delta child. Every +engine-run child of that boundary therefore runs as a fresh rollout. + +These are unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest + +from benchflow.branch import UNSCORED_KEY +from benchflow.branch_delta import BranchDelta +from benchflow.branch_skill import child_skill_config, resolve_child_skill_policy +from benchflow.environment.protocol import StateSnapshot +from benchflow.rollout import Rollout, RolloutConfig, Scene +from benchflow.rollout_branch import ( + _EXECUTABLE_DELTA_FIELDS, + _UNSUPPORTED_DELTA_FIELDS, + BranchChildExecutionNotSupported, + BranchDeltaNotSupported, + BranchParentSkillModeConflict, +) +from benchflow.sandbox.protocol import SandboxImage +from benchflow.task.paths import RolloutPaths +from benchflow.trajectories.tree import Step + +SKILL_BODY = "---\nname: demo\n---\n\nUse the demo skill.\n" +DOCKERFILE = "FROM python:3.12-slim\nCOPY skills /skills\nCOPY . /app\n" +# One tool call: a rollout that ends with zero tokens AND zero tool calls is +# classified as a suspected silent API failure and has its reward nulled, so a +# fake agent has to show some activity to be scoreable. +AGENT_EVENTS: list[dict[str, Any]] = [{"type": "tool_call", "tool_name": "bash"}] + + +def _task_dir(tmp_path: Path, *, bundled_skills: bool = True) -> Path: + """A minimal real task — the skill policy resolves against this on disk.""" + task = tmp_path / "task" + (task / "environment").mkdir(parents=True) + (task / "task.toml").write_text('version = "1.0"\n') + (task / "instruction.md").write_text("Solve the task.") + (task / "environment" / "Dockerfile").write_text(DOCKERFILE) + if bundled_skills: + skill = task / "environment" / "skills" / "demo" + skill.mkdir(parents=True) + (skill / "SKILL.md").write_text(SKILL_BODY) + return task + + +class FakeEnv: + """Environment-plane stand-in recording snapshot/restore calls.""" + + def __init__(self) -> None: + self.snapshots: list[StateSnapshot] = [] + self.restored: list[StateSnapshot] = [] + + async def snapshot(self) -> StateSnapshot: + snap = StateSnapshot(id=f"env-snap-{len(self.snapshots) + 1}", path="/tmp/x") + self.snapshots.append(snap) + return snap + + async def restore(self, snap: StateSnapshot) -> None: + self.restored.append(snap) + + +class FakeSandbox: + """Snapshot-capable Sandbox stand-in; ``exec`` answers the workspace probe.""" + + supports_snapshot = True + + def __init__(self) -> None: + self.snapshots: list[SandboxImage] = [] + self.restored: list[SandboxImage] = [] + self.stopped = False + + async def snapshot(self, name: str | None = None) -> SandboxImage: + img = SandboxImage(provider="fake", ref=f"bf-snap-{len(self.snapshots) + 1}") + self.snapshots.append(img) + return img + + async def restore(self, image: SandboxImage) -> None: + self.restored.append(image) + + async def exec(self, command: str, **_kw: Any) -> Any: + return SimpleNamespace(return_code=0, stdout="/app", stderr="") + + async def stop(self, delete: bool = False) -> None: + self.stopped = True + + +class FakePlanes: + """Recording stand-in for the concrete plane bundle (contracts.RolloutPlanes). + + ``deploy_skills`` is the call that actually puts (or does not put) the pack + into the sandbox, so it records what the rollout resolved *and* what the + staged build context looked like at that moment — the temp task copy is + deleted by cleanup(), so the evidence has to be taken while it exists. + """ + + def __init__(self) -> None: + self.deployments: list[dict[str, Any]] = [] + self.dockerfile_injections: list[Path] = [] + + # --- host-side setup ------------------------------------------------- + def install_docker_compat(self) -> None: + return None + + def resolve_locked_paths(self, sandbox_user, locked_paths): + return [] + + def resolve_agent_env(self, agent, model, agent_env): + return dict(agent_env or {}) + + def agent_launch(self, agent, *, disallow_web_tools): + return "oracle" + + def inject_skills_into_dockerfile( + self, task_path, skills_dir, *, sandbox_dir="/skills" + ): + self.dockerfile_injections.append(Path(skills_dir)) + + def extract_usage(self, runtime): + return {"usage_source": "unavailable"} + + # --- install_agent (oracle path) ------------------------------------- + async def setup_sandbox_user( + self, env, sandbox_user, *, workspace, timeout_sec=120 + ): + return workspace + + async def snapshot_build_config(self, env, *, workspace): + return None + + async def seed_verifier_workspace(self, env, *, workspace, sandbox_user): + return None + + async def deploy_skills(self, env, task_path, skills_dir, *args, **kwargs): + staged = Path(task_path) + dockerfile = staged / "environment" / "Dockerfile" + self.deployments.append( + { + "skills_dir": Path(skills_dir) if skills_dir is not None else None, + "skill_files": sorted( + p.parent.name for p in Path(skills_dir).glob("*/SKILL.md") + ) + if skills_dir is not None + else [], + "staged_skills_dir_exists": ( + staged / "environment" / "skills" + ).is_dir(), + "dockerfile": dockerfile.read_text() if dockerfile.exists() else "", + } + ) + + async def lockdown_paths(self, env, locked_paths): + return None + + # --- connect / execute ------------------------------------------------ + async def ensure_litellm_runtime(self, *args, **kwargs): + return kwargs.get("agent_env", {}), None + + async def connect_acp(self, *args, **kwargs): + async def close() -> None: + return None + + return SimpleNamespace(close=close), SimpleNamespace(), None, "oracle" + + async def execute_prompts(self, client, session, prompts, timeout, **kwargs): + return list(AGENT_EVENTS), 1 + + +def _fake_verifier(monkeypatch, reward: float = 1.0) -> None: + """Fake the child's verifier I/O; everything before it stays real.""" + + async def fake_publish(*_a, **_kw): + return None + + async def fake_verify_rollout(*_a, **_kw): + return {"reward": reward}, None, None + + monkeypatch.setattr( + "benchflow.rollout._publish_trajectory_for_verifier", fake_publish + ) + monkeypatch.setattr("benchflow.rollout._verify_rollout", fake_verify_rollout) + + +def _fake_verifier_without_reward(monkeypatch, *, error: str) -> None: + """Fake a verifier that ran but left no reward — the shape of the live bug. + + ``_verify_rollout`` returns ``(rewards, verifier_error, timeout_diag)``; + when the reward file never reaches the host, ``rewards`` is ``None`` and + the error explains why. ``Rollout.verify()`` then returns ``None``. + """ + + async def fake_publish(*_a, **_kw): + return None + + async def fake_verify_rollout(*_a, **_kw): + return None, error, None + + monkeypatch.setattr( + "benchflow.rollout._publish_trajectory_for_verifier", fake_publish + ) + monkeypatch.setattr("benchflow.rollout._verify_rollout", fake_verify_rollout) + + +def _parent( + task: Path, + tmp_path: Path, + *, + planes: FakePlanes, + skill_mode: str = "no-skill", + **config: Any, +) -> Rollout: + """A parent rollout positioned as if it had run past env-ready.""" + rollout = Rollout( + RolloutConfig( + task_path=task, + scenes=[Scene.single(agent="oracle")], + jobs_dir=tmp_path / "jobs", + skill_mode=skill_mode, + planes=planes, + **config, + ) + ) + rollout._environment, rollout._env = FakeEnv(), FakeSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir(exist_ok=True) + rollout._rollout_dir = run_dir + # What setup() would have left behind. Without it an in-place branch child + # dies in verify() on a None _rollout_paths, which would let a test "pass" + # for an incidental reason instead of on the property under test. + rollout._rollout_paths = RolloutPaths(rollout_dir=run_dir) + rollout._rollout_paths.mkdir() + return rollout + + +async def _capture_env_ready( + rollout: Rollout, *, layers: set[str] | None = None +) -> None: + """Record env-ready at the root, then run past it linearly.""" + await rollout.mark_stage( + "env-ready", + snapshot_layers={"environment", "sandbox"} if layers is None else layers, + ) + rollout._cursor = rollout._tree.advance(rollout._cursor, Step(id="s1")) + + +# 1. The ablation: a skill_mode child runs a fresh rollout with the switched mode + + +async def test_skill_delta_children_deploy_different_skills_at_env_ready( + tmp_path: Path, monkeypatch +): + """The with-skill/no-skill ablation, executed for real. + + Each child re-runs install_agent() as its own Rollout over the restored + env-ready sandbox, so the evidence is the deploy call itself: the with-skill + child deploys the task's bundled pack, the no-skill child deploys nothing + AND its staged build context no longer contains the pack or the COPY line + that would smuggle it in. + """ + planes = FakePlanes() + task = _task_dir(tmp_path) + rollout = _parent(task, tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert value == 1.0 + with_skill, no_skill = planes.deployments + assert with_skill["skill_files"] == ["demo"] + assert with_skill["staged_skills_dir_exists"] is True + assert "COPY skills /skills" in with_skill["dockerfile"] + assert no_skill["skills_dir"] is None + assert no_skill["skill_files"] == [] + assert no_skill["staged_skills_dir_exists"] is False + assert "COPY skills /skills" not in no_skill["dockerfile"] + # the pack is stripped from the child's staged copy, never from the task + assert (task / "environment" / "skills" / "demo" / "SKILL.md").exists() + + +async def test_skill_child_never_injects_skills_into_the_staged_dockerfile( + tmp_path: Path, monkeypatch +): + """The child adopts the restored sandbox, so its Dockerfile is never built. + + Real-docker repro: the with-skill child staged its skills into the task copy + and appended ``COPY _deps/skills /skills/`` to a Dockerfile nothing rebuilds. + ``deploy_skills`` then read that line as "already baked into the image", + skipped the runtime upload, and the link step failed closed with + ``experiment_fidelity/skill_deployment_missing``. Injection must not happen + for a caller-owned sandbox — the pack has to arrive by runtime upload. + """ + planes = FakePlanes() + task = _task_dir(tmp_path) + rollout = _parent(task, tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert planes.dockerfile_injections == [] + with_skill, no_skill = planes.deployments + # The pack still reaches the with-skill child — by upload, not by build. + assert with_skill["skill_files"] == ["demo"] + assert no_skill["skills_dir"] is None + + +async def test_skill_delta_child_is_a_separate_rollout_over_the_parent_sandbox( + tmp_path: Path, monkeypatch +): + """Fresh rollout, same (restored) sandbox: the child owns its own config, + run directory and result, and must not stop the caller's container (#388).""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + sandbox = rollout._env + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + parent_task_path = rollout._effective_task_path + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + configs = { + child.name: json.loads((child / "config.json").read_text()) + for child in sorted(children_dir.iterdir()) + } + assert [cfg["skill_mode"] for cfg in configs.values()] == [ + "with-skill", + "no-skill", + ] + assert [cfg["include_task_skills"] for cfg in configs.values()] == [True, False] + assert all((child / "result.json").exists() for child in children_dir.iterdir()) + # the parent instance was never re-entered: its own setup state is untouched + assert rollout._effective_task_path == parent_task_path + assert rollout._config.skill_mode == "no-skill" + assert sandbox.stopped is False + assert len(sandbox.restored) == 2 # one container roll-back per child + + +async def test_skill_delta_child_restores_the_snapshot_before_installing( + tmp_path: Path, monkeypatch +): + """Order matters: a child that installed before the roll-back would deploy + into the previous child's container.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + calls: list[str] = [] + sandbox = rollout._env + original_restore = sandbox.restore + + async def recording_restore(image): + calls.append("sandbox.restore") + await original_restore(image) + + sandbox.restore = recording_restore + original_deploy = planes.deploy_skills + + async def recording_deploy(*args, **kwargs): + calls.append("deploy_skills") + await original_deploy(*args, **kwargs) + + planes.deploy_skills = recording_deploy + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="no-skill"), + BranchDelta(skill_mode="with-skill"), + ], + ) + + assert calls == [ + "sandbox.restore", + "deploy_skills", + "sandbox.restore", + "deploy_skills", + ] + + +async def test_skill_delta_child_delivers_an_injected_prompt( + tmp_path: Path, monkeypatch +): + """A skill delta that also injects a prompt: the fresh rollout runs the + injection as its continuation prompt, not the task's base prompt.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + prompts: list[list[str]] = [] + + async def recording_execute(client, session, sent, timeout, **kwargs): + prompts.append(list(sent)) + return list(AGENT_EVENTS), 1 + + planes.execute_prompts = recording_execute + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill", injected_prompt="Follow PLAN.md."), + ], + ) + + assert prompts == [["Solve the task."], ["Follow PLAN.md."]] + + +async def test_a_child_whose_verifier_produced_no_reward_is_never_scored_zero( + tmp_path: Path, monkeypatch +): + """The real-world failure this guard exists for, in miniature. + + A live ablation lost both children's rewards (the restored container had no + host-mounted ``/logs``, so the verifier's ``reward.txt`` was never + downloaded and ``verify()`` returned ``None``) and the run reported two + confident ``0.00``s. A child that produced no reward is *unscored*, not a + zero: the node carries the reason and no ``reward``, no ``reward.json`` is + fabricated for it, and V(parent) is undefined rather than averaged from + nothing. + """ + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier_without_reward( + monkeypatch, + error=("No reward file found at /run/verifier/reward.txt or reward.json"), + ) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + children = [node for node in rollout.tree.nodes() if "delta" in node.state] + assert len(children) == 2 + assert all("reward" not in child.state for child in children) + assert all( + "No reward file found" in child.state[UNSCORED_KEY] for child in children + ) + assert value is None + assert "value" not in rollout.tree.root.state + + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + assert not any( + (child / "reward.json").exists() for child in sorted(children_dir.iterdir()) + ) + tree = json.loads((rollout._rollout_dir / "tree.json").read_text()) + scored = [node for node in tree["nodes"] if "reward" in node] + assert scored == [] + + +# 2. Lineage: the effective mode and the fresh-rollout execution are recorded + + +async def test_provenance_records_the_effective_mode_and_fresh_rollout( + tmp_path: Path, monkeypatch +): + """Per-child provenance.json carries the delta's skill_mode verbatim and + delta_execution=fresh-rollout; tree.json carries the same on the node.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + run_dir = rollout._rollout_dir + children_dir = run_dir / "branches" / "root" / "children" + provenance = [ + json.loads((children_dir / child_id / "provenance.json").read_text()) + for child_id in ("n2", "n3") + ] + assert [p["delta"]["skill_mode"] for p in provenance] == [ + "with-skill", + "no-skill", + ] + assert [p["delta_execution"] for p in provenance] == [ + "fresh-rollout", + "fresh-rollout", + ] + assert [p["branch_stage"] for p in provenance] == ["env-ready", "env-ready"] + assert provenance[0]["snapshot_ref"] == { + "environment": "env-snap-1", + "sandbox": "bf-snap-1", + } + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["n2"]["delta_execution"] == "fresh-rollout" + assert "delta_execution" not in nodes["root"] + + +async def test_the_child_rollout_carries_the_branch_source_provenance( + tmp_path: Path, monkeypatch +): + """The child is a first-class rollout, so its own config.json/result.json + record which rollout it forked from — the seam a continued run uses.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch, reward=0.5) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + child_dir = rollout._rollout_dir / "branches" / "root" / "children" / "n2" + source = json.loads((child_dir / "config.json").read_text())["source"] + assert source["kind"] == "benchflow-branch" + assert source["branch_stage"] == "env-ready" + assert source["delta_execution"] == "fresh-rollout" + assert json.loads((child_dir / "result.json").read_text())["rewards"] == { + "reward": 0.5 + } + assert json.loads((child_dir / "reward.json").read_text()) == {"reward": 0.5} + assert value == 0.5 + + +# 3. Everywhere else still fails closed + + +async def test_skill_delta_at_a_cursor_branch_names_env_ready(tmp_path: Path): + """A cursor branch forks after install_agent(), so the delta could only be + recorded, never executed — fail closed naming the boundary that works.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + env = rollout._environment + + with pytest.raises(BranchDeltaNotSupported, match="skill_mode") as excinfo: + await rollout.branch( + 2, + snapshot_layers={"environment", "sandbox"}, + deltas=[None, BranchDelta(skill_mode="with-skill")], + ) + + assert "env-ready" in str(excinfo.value) + assert "install_agent" in str(excinfo.value) + assert env.snapshots == [] # nothing quiesced, snapshotted, or forked + assert rollout._cursor.children == [] + + +async def test_skill_delta_at_another_stage_names_env_ready( + tmp_path: Path, monkeypatch +): + """pre-verify is a recorded boundary too — and still the wrong one.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await rollout.mark_stage("pre-verify", snapshot_layers={"environment", "sandbox"}) + + with pytest.raises(BranchDeltaNotSupported, match="'env-ready'") as excinfo: + await rollout.branch_at_stage( + "pre-verify", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert "'pre-verify'" in str(excinfo.value) + assert planes.deployments == [] + + +async def test_skill_delta_needs_the_container_layer_in_the_snapshot( + tmp_path: Path, +): + """An environment-state-only env-ready snapshot cannot roll the container + back, so a no-skill child would re-install on top of the parent's pack — + fail closed instead of measuring nothing.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout, layers={"environment"}) + + with pytest.raises(BranchDeltaNotSupported, match="sandbox") as excinfo: + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert "container filesystem" in str(excinfo.value) + assert rollout._environment.restored == [] + assert planes.deployments == [] + + +async def test_skill_delta_cannot_fork_a_parent_that_baked_its_pack_in( + tmp_path: Path, monkeypatch +): + """The mislabeled-arm hole: a ``no-skill`` arm of a ``with-skill`` parent. + + Guards the fix from "fix(branch): gate skill deltas on the parent's own + skill mode". The gates below it check the stage, the layer and the runner — + none of them look at what the *parent* installed. A ``with-skill`` parent's + ``setup()`` calls ``inject_skills_into_dockerfile``, so the pack is in the + image, and the ``env-ready`` snapshot is a commit of that image. The + ``no-skill`` child then restored the pack, resolved no host directory, + deployed nothing on top of it — and ran *with* the skills while its + ``config.json``, its provenance and the ablation row all said ``no-skill``. + + The first assertion is the mechanism, taken from the parent's own real + ``setup()``: the pack is baked into the image the snapshot commits. The + rest is the gate. + """ + planes = FakePlanes() + task = _task_dir(tmp_path) + rollout = _parent(task, tmp_path, planes=planes, skill_mode="with-skill") + await rollout.setup() + assert planes.dockerfile_injections != [], ( + "precondition: a with-skill parent bakes the pack into the image its " + "env-ready snapshot commits" + ) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + with pytest.raises(BranchParentSkillModeConflict) as excinfo: + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert "'with-skill'" in str(excinfo.value) # names the parent's own mode + assert "no-skill" in str(excinfo.value) # and the mode it must be run in + assert isinstance(excinfo.value, BranchDeltaNotSupported) + # fail closed before anything is restored, installed or forked + assert rollout._env.restored == [] + assert rollout._environment.restored == [] + assert planes.deployments == [] + assert [node for node in rollout.tree.nodes() if "delta" in node.state] == [] + + +@pytest.mark.parametrize("parent_mode", ["with-skill", "self-gen"]) +@pytest.mark.parametrize( + "deltas", + [ + ["with-skill", "no-skill"], + [None, "no-skill"], + ["no-skill", None], + ["no-skill", "no-skill"], + ], + ids=[ + "both-arms", + "control-then-no-skill", + "no-skill-then-control", + "both-no-skill", + ], +) +async def test_no_arm_forked_from_a_skilled_parent_can_claim_no_skill( + tmp_path: Path, parent_mode: str, deltas: list[str | None] +): + """A mislabeled arm is unobtainable, not merely unlikely. + + Guards the fix from "fix(branch): gate skill deltas on the parent's own + skill mode". ``bench eval ablate`` only ever built a ``no-skill`` parent, + so the hole was reachable solely through the Python API — which is exactly + the caller that has no CLI default protecting it. Every shape of that call + that could produce a ``no-skill``-labelled arm over a parent that may have + baked a pack in is refused here, and none of them leaves a forked child + behind to be reported. + """ + planes = FakePlanes() + rollout = _parent( + _task_dir(tmp_path), tmp_path, planes=planes, skill_mode=parent_mode + ) + await _capture_env_ready(rollout) + + with pytest.raises(BranchParentSkillModeConflict, match=parent_mode): + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + None if mode is None else BranchDelta(skill_mode=mode) + for mode in deltas + ], + ) + + assert planes.deployments == [] + assert [node for node in rollout.tree.nodes() if "delta" in node.state] == [] + + +async def test_a_no_skill_parent_is_the_one_the_ablation_forks_from( + tmp_path: Path, monkeypatch +): + """The other half of the gate: the parent it *does* accept bakes nothing. + + Guards the fix from "fix(branch): gate skill deltas on the parent's own + skill mode" against being tightened into a gate nothing can pass. The + accepted parent's own ``setup()`` injects no skills into the Dockerfile, so + the image the arms restore carries no pack and each arm's own + ``deploy_skills`` is the only thing that puts one there. + """ + planes = FakePlanes() + rollout = _parent( + _task_dir(tmp_path), tmp_path, planes=planes, skill_mode="no-skill" + ) + await rollout.setup() + assert planes.dockerfile_injections == [] # nothing baked into the image + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert value == 1.0 + with_skill, no_skill = planes.deployments + assert with_skill["skill_files"] == ["demo"] + assert no_skill["skills_dir"] is None + assert planes.dockerfile_injections == [] # still nothing built with a pack + + +async def test_skill_delta_with_an_explicit_run_child_is_rejected(tmp_path: Path): + """A caller-supplied runner owns the child's execution, so the engine + cannot re-run install_agent() under the switched mode.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout) + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="run_child"): + await rollout.branch_at_stage( + "env-ready", + 2, + run_child=run_child, + deltas=[None, BranchDelta(skill_mode="no-skill")], + ) + + assert planes.deployments == [] + + +async def test_with_skill_against_a_task_without_skills_fails_before_any_child( + tmp_path: Path, +): + """The ablation's precondition is resolved up front: a with-skill delta on + a task shipping no pack raises the setup() error before the branch restores + anything, not halfway through the fork.""" + planes = FakePlanes() + task = _task_dir(tmp_path, bundled_skills=False) + rollout = _parent(task, tmp_path, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(FileNotFoundError, match="no bundled skills"): + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[ + BranchDelta(skill_mode="with-skill"), + BranchDelta(skill_mode="no-skill"), + ], + ) + + assert rollout._env.restored == [] + assert planes.deployments == [] + + +async def test_an_environment_ref_delta_needs_a_manifest_bound_parent( + tmp_path: Path, +): + """``environment_ref`` executes only as a variation of the parent's own + bound manifest ("feat(branch): execute service-level environment_ref + deltas from env-ready") — a parent with no Environment plane fails closed + at the env-ready boundary too, before anything is restored. The executed + slice is pinned in tests/test_branch_environment_delta.py.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout) + + with pytest.raises(BranchDeltaNotSupported, match="no environment manifest"): + await rollout.branch_at_stage( + "env-ready", 2, deltas=[None, BranchDelta(environment_ref="env0@outage")] + ) + + assert planes.deployments == [] + assert rollout._env.restored == [] + + +def test_every_delta_field_is_executable_and_the_blocklist_is_empty(): + """The blocklist is derived from the schema minus the executable set, so + graduating an axis is one edit and a future field stays + unsupported-by-default. ``config_override`` and ``environment_ref`` + graduated last ("feat(branch): execute config_override deltas as fresh + rollouts from env-ready" / "feat(branch): execute service-level + environment_ref deltas from env-ready"), leaving the derived blocklist + empty — but still load-bearing the day BranchDelta grows a field.""" + assert set(_EXECUTABLE_DELTA_FIELDS) == { + "injected_prompt", + "skill_mode", + "config_override", + "environment_ref", + } + assert set(_UNSUPPORTED_DELTA_FIELDS) == set() + + +# 4. The derived child config + + +def test_child_config_replaces_both_recorded_skill_modes(tmp_path: Path): + """The skill policy resolves from recorded_skill_mode (artifact_skill_mode + or skill_mode), so writing only one of them would leave a from_legacy-built + parent's mode in place and the delta would silently do nothing.""" + task = _task_dir(tmp_path) + parent = RolloutConfig.from_legacy( + task_path=task, + agent="oracle", + skill_mode="with-skill", + jobs_dir=tmp_path / "jobs", + ) + assert parent.recorded_skill_mode == "with-skill" + + child = child_skill_config( + parent, + skill_mode="no-skill", + jobs_dir=tmp_path / "jobs", + job_name="children", + rollout_name="n2", + ) + + assert child.recorded_skill_mode == "no-skill" + assert child.skills_dir is None # a no-skill config cannot carry skills_dir + assert resolve_child_skill_policy(child, child.skill_mode).host_dir is None + assert child.task_path == parent.task_path # nothing else moved + assert child.snapshot_stages == frozenset() + assert parent.skill_mode == "with-skill" # the parent config is untouched + + +def test_child_config_keeps_a_strategy_user_materializable(tmp_path: Path): + """A loop strategy materializes a user in __post_init__; carrying both into + the child config is rejected by RolloutConfig, so the child re-materializes + the identical user from the same spec.""" + parent = RolloutConfig( + task_path=_task_dir(tmp_path), + scenes=[Scene.single(agent="oracle")], + jobs_dir=tmp_path / "jobs", + loop_strategy="verify-retry:k=3", + ) + assert parent.user is not None + + child = child_skill_config( + parent, + skill_mode="with-skill", + jobs_dir=tmp_path / "jobs", + job_name="children", + rollout_name="n2", + ) + + assert child.user is not None + assert type(child.user) is type(parent.user) + assert child.max_user_rounds == parent.max_user_rounds + + +# 5. Every env-ready child is a fresh rollout — the delta does not decide it + + +async def test_injected_prompt_child_at_env_ready_reinstalls_the_agent( + tmp_path: Path, monkeypatch +): + """The WS-4c hole: an ``inject:`` arm forked from ``env-ready``. + + ``env-ready`` is captured before ``install_agent()``, so the restored world + has no agent binary, no sandbox user, no seeded verifier workspace, no + lockdown and no skill pack. Routing that child through the in-place default + runner called ``connect()`` on a world where nothing had been installed: + against a real sandbox it dies launching a deleted binary, and against one + where the agent survives in the base image it scores a skill-less, + lockdown-less world and the report calls it "the parent's world plus one + injected prompt". + + The evidence that it is fixed is the deploy call itself — one per child, + the same call the skills arms are measured by. Under the old routing + ``planes.deployments`` is empty because no child ever re-installed. + """ + planes = FakePlanes() + rollout = _parent( + _task_dir(tmp_path), tmp_path, planes=planes, skill_mode="with-skill" + ) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + value = await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(injected_prompt="Follow PLAN.md verbatim.")], + ) + + assert value == 1.0 + assert len(planes.deployments) == 2, ( + "each env-ready child must re-run install_agent() for itself — an " + "in-place child connects to an agent the restore deleted" + ) + # ... and re-installs the *parent's* mode: an inject arm is a zero delta on + # skills, so both children must carry the pack the parent ran with. + assert [deployment["skill_files"] for deployment in planes.deployments] == [ + ["demo"], + ["demo"], + ] + assert len(rollout._env.restored) == 2 # one container roll-back per child + + +async def test_zero_delta_child_at_env_ready_is_a_fresh_rollout_too( + tmp_path: Path, monkeypatch +): + """No deltas at all is the same hole: nothing about a delta made the + in-place child unsound, the boundary did. Both children install.""" + planes = FakePlanes() + rollout = _parent( + _task_dir(tmp_path), tmp_path, planes=planes, skill_mode="no-skill" + ) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage("env-ready", 2) + + assert len(planes.deployments) == 2 + assert [deployment["skills_dir"] for deployment in planes.deployments] == [ + None, + None, + ] # the parent's own recorded mode, re-installed unchanged + children_dir = rollout._rollout_dir / "branches" / "root" / "children" + configs = [ + json.loads((child / "config.json").read_text()) + for child in sorted(children_dir.iterdir()) + ] + assert [cfg["skill_mode"] for cfg in configs] == ["no-skill", "no-skill"] + assert all((child / "result.json").exists() for child in children_dir.iterdir()) + + +async def test_injected_prompt_at_env_ready_reaches_the_fresh_child( + tmp_path: Path, monkeypatch +): + """The injection is still the child's user-visible first message — routing + it through a fresh rollout must not swallow it.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + prompts: list[list[str]] = [] + + async def recording_execute(client, session, sent, timeout, **kwargs): + prompts.append(list(sent)) + return list(AGENT_EVENTS), 1 + + planes.execute_prompts = recording_execute + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", + 2, + deltas=[None, BranchDelta(injected_prompt="Follow PLAN.md verbatim.")], + ) + + assert prompts == [["Solve the task."], ["Follow PLAN.md verbatim."]] + + +async def test_env_ready_lineage_marks_every_child_fresh_rollout( + tmp_path: Path, monkeypatch +): + """``delta_execution`` is a property of how the engine ran the child, so a + zero-delta env-ready child records it too — a reader cannot infer it from + the delta, which is empty.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + _fake_verifier(monkeypatch) + await _capture_env_ready(rollout) + + await rollout.branch_at_stage( + "env-ready", 2, deltas=[None, BranchDelta(injected_prompt="plan")] + ) + + run_dir = rollout._rollout_dir + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["n2"]["delta_execution"] == "fresh-rollout" + assert nodes["n3"]["delta_execution"] == "fresh-rollout" + children_dir = run_dir / "branches" / "root" / "children" + assert [ + json.loads((children_dir / cid / "provenance.json").read_text())[ + "delta_execution" + ] + for cid in ("n2", "n3") + ] == ["fresh-rollout", "fresh-rollout"] + + +async def test_env_ready_children_need_the_container_layer_without_any_delta( + tmp_path: Path, +): + """An environment-state-only env-ready snapshot cannot undo one child's + installation before the next child installs, so the fork fails closed — + even with no deltas at all, where no skill-delta gate would fire.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout, layers={"environment"}) + + with pytest.raises(BranchChildExecutionNotSupported, match="sandbox") as excinfo: + await rollout.branch_at_stage("env-ready", 2) + + assert "install_agent" in str(excinfo.value) + assert isinstance(excinfo.value, NotImplementedError) + # fail closed before anything is restored or installed + assert rollout._environment.restored == [] + assert rollout._env.restored == [] + assert planes.deployments == [] + assert [node for node in rollout.tree.nodes() if "delta" in node.state] == [] + + +async def test_a_caller_supplied_runner_at_env_ready_still_owns_execution( + tmp_path: Path, +): + """The engine never second-guesses an explicit ``run_child`` — including + the layer gate, which exists only because *the engine* would install.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + await _capture_env_ready(rollout, layers={"environment"}) + ran: list[str] = [] + + async def run_child(child): + ran.append(child.id) + return 1.0 + + value = await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + assert value == 1.0 + assert ran == ["n2", "n3"] + assert planes.deployments == [] + + +async def test_an_injected_prompt_after_install_still_runs_in_place( + tmp_path: Path, monkeypatch +): + """The counterpart guard: ``pre-verify`` is captured *after* + ``install_agent()``, so its children continue the installed world in place + and must NOT re-install — re-installing there would be the second-order + version of the same bug, changing the world an inject arm measures.""" + planes = FakePlanes() + rollout = _parent(_task_dir(tmp_path), tmp_path, planes=planes) + # The in-place runner replays the parent's own resolved prompts; the parent + # here never ran setup(), so seed them the way setup() would. + rollout._resolved_prompts = ["Solve the task."] + _fake_verifier(monkeypatch) + prompts: list[list[str] | None] = [] + + async def recording_execute(client, session, sent, timeout, **kwargs): + prompts.append(list(sent) if sent is not None else None) + return list(AGENT_EVENTS), 1 + + planes.execute_prompts = recording_execute + await rollout.mark_stage("pre-verify", snapshot_layers={"environment", "sandbox"}) + + value = await rollout.branch_at_stage( + "pre-verify", 2, deltas=[None, BranchDelta(injected_prompt="Retry the fix.")] + ) + + assert value == 1.0 + assert planes.deployments == [] # nothing re-installed + assert prompts == [["Solve the task."], ["Retry the fix."]] + nodes = { + node["id"]: node + for node in json.loads((rollout._rollout_dir / "tree.json").read_text())[ + "nodes" + ] + } + assert "delta_execution" not in nodes["n1"] diff --git a/tests/test_branch_stage_policy.py b/tests/test_branch_stage_policy.py new file mode 100644 index 000000000..cb6b6255e --- /dev/null +++ b/tests/test_branch_stage_policy.py @@ -0,0 +1,1027 @@ +"""Regression tests for the stage-boundary snapshot policy. + +Guards the stage-boundary policy ("feat(branch): stage-boundary snapshot policy +for rollout branching"; docs/rollout-branching-rfc.md §3.2 / WS-4a; +FrontierPhysics#73). PR number to be added on submission. + +The four cascade stages pin to existing lifecycle transitions: ``env-ready`` +at the end of ``start()`` (before ``install_agent()``), ``post-research`` at an +explicit ``mark_stage()`` inside ``execute()``, ``pre-verify`` immediately +before ``planes.harden_before_verify``, and ``post-verify`` after ``verify()``. +Capture is opt-in via ``RolloutConfig.snapshot_stages`` — the default set is +empty and must keep today's behavior exactly, with no snapshot call anywhere. +``branch_at_stage()`` then forks from a recorded stage instead of +checkpointing at the cursor, and records the stage name as ``branch_stage``. + +These are unit tests against fakes — no Docker, Daytona, or API keys. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from benchflow.branch import StageSnapshot +from benchflow.branch_delta import BranchDelta +from benchflow.branch_stage import ( + AUTO_STAGES, + BRANCH_STAGES, + MARKED_STAGES, + BranchStageNotCaptured, +) +from benchflow.environment.manifest import EnvironmentManifest +from benchflow.environment.manifest_env import ManifestEnvironment +from benchflow.environment.protocol import StateSnapshot +from benchflow.rollout import Rollout, RolloutConfig, Scene +from benchflow.rollout_branch import BranchDeltaNotSupported +from benchflow.sandbox.protocol import SandboxImage, SandboxSnapshotNotSupported +from benchflow.trajectories.tree import Step + + +class FakeEnv: + """Stateful in-memory Environment — snapshot copies state, restore rolls back. + + ``calls`` may be shared with the fake sandbox and the faked lifecycle + boundaries to record ordering across all of them. + """ + + def __init__(self, calls: list[str] | None = None) -> None: + self.state: dict[str, str] = {"db": "initial"} + self._saved: dict[str, dict[str, str]] = {} + self.calls = calls if calls is not None else [] + self.snapshots: list[StateSnapshot] = [] + self.restored: list[StateSnapshot] = [] + + async def snapshot(self) -> StateSnapshot: + snap = StateSnapshot(id=f"env-snap-{len(self.snapshots) + 1}", path="/tmp/x") + self._saved[snap.id] = dict(self.state) + self.snapshots.append(snap) + self.calls.append("env.snapshot") + return snap + + async def restore(self, snap: StateSnapshot) -> None: + self.state = dict(self._saved[snap.id]) + self.restored.append(snap) + self.calls.append("env.restore") + + +class FakeSnapSandbox: + """Snapshot-capable Sandbox stand-in with in-memory filesystem state.""" + + supports_snapshot = True + + def __init__(self, calls: list[str] | None = None) -> None: + self.fs: dict[str, str] = {"/workspace": "clean"} + self._images: dict[str, dict[str, str]] = {} + self.calls = calls if calls is not None else [] + self.snapshots: list[SandboxImage] = [] + self.restored: list[SandboxImage] = [] + + async def snapshot(self, name: str | None = None) -> SandboxImage: + img = SandboxImage(provider="fake", ref=f"bf-snap-{len(self.snapshots) + 1}") + self._images[img.ref] = dict(self.fs) + self.snapshots.append(img) + self.calls.append("sandbox.snapshot") + return img + + async def restore(self, image: SandboxImage) -> None: + self.fs = dict(self._images[image.ref]) + self.restored.append(image) + self.calls.append("sandbox.restore") + + +class NoSnapshotSandbox: + """Sandbox without container snapshot — the capability gate must fail closed.""" + + supports_snapshot = False + + def __init__(self, calls: list[str] | None = None) -> None: + self.calls = calls if calls is not None else [] + + async def snapshot(self, name: str | None = None) -> SandboxImage: + self.calls.append("sandbox.snapshot") + raise SandboxSnapshotNotSupported("NoSnapshotSandbox does not support snapshot") + + async def restore(self, image: SandboxImage) -> None: + self.calls.append("sandbox.restore") + raise SandboxSnapshotNotSupported("NoSnapshotSandbox does not support restore") + + +# A manifest with no [environment.state] — a *stateless* environment whose +# snapshot() raises (the manifest_env fail-closed precedent, #387). +_STATELESS_MANIFEST = EnvironmentManifest.model_validate_toml( + """ +[environment] +name = "chi-bench" +image = "chi-bench:latest" +ports = [8020] +owns_lifecycle = true +""" +) + + +def _rollout(tmp_path: Path, **config) -> Rollout: + return Rollout( + RolloutConfig( + task_path=tmp_path / "task", + scenes=[Scene.single(agent="dummy")], + **config, + ) + ) + + +def _fake_start_boundary(monkeypatch, calls: list[str]) -> None: + """Fake everything ``start()`` does before the env-ready boundary.""" + + async def fake_start_env_and_upload(*_a, **_kw): + calls.append("start-env") + + async def fake_healthcheck(*_a, **_kw): + calls.append("healthcheck") + + async def fake_setup_commands(*_a, **_kw): + calls.append("env-setup-commands") + + monkeypatch.setattr( + "benchflow.rollout._start_env_and_upload", fake_start_env_and_upload + ) + monkeypatch.setattr( + "benchflow.rollout._run_environment_healthcheck", fake_healthcheck + ) + monkeypatch.setattr( + "benchflow.rollout._run_environment_setup_commands", fake_setup_commands + ) + + async def fake_install_agent(self): + calls.append("install_agent") + + monkeypatch.setattr(Rollout, "install_agent", fake_install_agent) + + +def _fake_verify_boundary(monkeypatch, calls: list[str], tmp_path: Path) -> None: + """Fake ``verify()``'s I/O; ``_verify_rollout`` hardens first, as in prod.""" + + async def fake_publish(*_a, **_kw): + calls.append("publish-trajectory") + + async def fake_verify_rollout(*_a, **_kw): + calls.append("harden_before_verify") + calls.append("verifier") + return {"reward": 1.0}, None, None + + monkeypatch.setattr( + "benchflow.rollout._publish_trajectory_for_verifier", fake_publish + ) + monkeypatch.setattr("benchflow.rollout._verify_rollout", fake_verify_rollout) + + +def _ready_to_verify(rollout: Rollout, tmp_path: Path) -> None: + """Minimal state so ``verify()`` reaches its boundaries with fakes in place.""" + rollout._trajectory = [{"role": "agent", "text": "done"}] + rollout._rollout_paths = SimpleNamespace(agent_dir=tmp_path) + + +# 1. Taxonomy + config validation + + +def test_taxonomy_splits_auto_detectable_from_marked_stages(): + """post-research is the one boundary no lifecycle transition can detect.""" + assert BRANCH_STAGES == ( + "env-ready", + "post-research", + "pre-verify", + "post-verify", + ) + assert sorted(AUTO_STAGES) == ["env-ready", "post-verify", "pre-verify"] + assert sorted(MARKED_STAGES) == ["post-research"] + + +def test_default_config_requests_no_stage_snapshots(tmp_path: Path): + """The default is empty — the policy is strictly opt-in.""" + rollout = _rollout(tmp_path) + assert rollout._config.snapshot_stages == frozenset() + assert rollout._config.snapshot_layers == frozenset({"environment"}) + assert rollout.stage_snapshots == {} + + +def test_unknown_stage_name_is_rejected_at_config_validation(tmp_path: Path): + """A typo fails closed at construction, not at the boundary it would miss.""" + with pytest.raises(ValueError, match="unknown snapshot_stages entry"): + _rollout(tmp_path, snapshot_stages={"env-ready", "post-executuion"}) + + +def test_snapshot_stages_rejects_a_bare_string(tmp_path: Path): + """A string would silently iterate into characters — every one unknown.""" + with pytest.raises(ValueError, match="collection of stage names"): + _rollout(tmp_path, snapshot_stages="env-ready") + + +def test_snapshot_stages_normalizes_any_collection(tmp_path: Path): + rollout = _rollout(tmp_path, snapshot_stages=["env-ready", "pre-verify"]) + assert rollout._config.snapshot_stages == frozenset({"env-ready", "pre-verify"}) + + +async def test_mark_stage_rejects_a_name_outside_the_taxonomy(tmp_path: Path): + rollout = _rollout(tmp_path) + rollout._environment = FakeEnv() + + with pytest.raises(ValueError, match="unknown stage"): + await rollout.mark_stage("post-planning") + + +async def test_branch_at_stage_rejects_a_name_outside_the_taxonomy(tmp_path: Path): + rollout = _rollout(tmp_path) + rollout._environment = FakeEnv() + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="unknown stage"): + await rollout.branch_at_stage("post-planning", 2, run_child=run_child) + + +# 2. Opt-in capture: the default lifecycle takes zero snapshots + + +async def test_default_config_takes_zero_snapshots_across_the_lifecycle( + tmp_path: Path, monkeypatch +): + """Zero overhead by default: not one snapshot call, not one artifact.""" + calls: list[str] = [] + rollout = _rollout(tmp_path) + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + rollout._environment, rollout._env = env, sandbox + rollout._rollout_dir = tmp_path + _fake_start_boundary(monkeypatch, calls) + _fake_verify_boundary(monkeypatch, calls, tmp_path) + _ready_to_verify(rollout, tmp_path) + + await rollout.start() + await rollout.install_agent() + await rollout.verify() + + assert env.snapshots == [] + assert sandbox.snapshots == [] + assert rollout.stage_snapshots == {} + assert not (tmp_path / "stage_snapshots.json").exists() + + +# 3. Each auto stage fires at its lifecycle boundary + + +async def test_env_ready_is_captured_at_the_end_of_start_before_install_agent( + tmp_path: Path, monkeypatch +): + """env-ready precedes install_agent() so a child re-runs skill deployment + from a skill-free world (RFC §3.2).""" + calls: list[str] = [] + rollout = _rollout( + tmp_path, + snapshot_stages={"env-ready"}, + snapshot_layers={"environment", "sandbox"}, + ) + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + rollout._environment, rollout._env = env, sandbox + _fake_start_boundary(monkeypatch, calls) + + await rollout.start() + await rollout.install_agent() + + assert calls == [ + "start-env", + "healthcheck", + "env-setup-commands", + "env.snapshot", + "sandbox.snapshot", + "install_agent", + ] + snap = rollout.stage_snapshots["env-ready"] + assert isinstance(snap, StageSnapshot) + assert snap.stage == "env-ready" + assert snap.environment_ref is env.snapshots[0] + assert snap.sandbox_ref is sandbox.snapshots[0] + # the stage tag rides on the node's own recorded checkpoint + assert rollout._cursor.state["snapshot"] is snap + + +async def test_pre_verify_is_captured_before_hardening_and_post_verify_after( + tmp_path: Path, monkeypatch +): + """pre-verify lands immediately before planes.harden_before_verify (which + _verify_rollout runs first), post-verify after the verifier scored.""" + calls: list[str] = [] + rollout = _rollout(tmp_path, snapshot_stages={"pre-verify", "post-verify"}) + env = FakeEnv(calls) + rollout._environment, rollout._env = env, FakeSnapSandbox(calls) + _fake_verify_boundary(monkeypatch, calls, tmp_path) + _ready_to_verify(rollout, tmp_path) + + rewards = await rollout.verify() + + assert rewards == {"reward": 1.0} + assert calls == [ + "publish-trajectory", + "env.snapshot", + "harden_before_verify", + "verifier", + "env.snapshot", + ] + assert [s.stage for s in rollout.stage_snapshots.values()] == [ + "pre-verify", + "post-verify", + ] + + +async def test_only_the_requested_stages_are_captured(tmp_path: Path, monkeypatch): + """A stage nobody asked for costs a membership test and nothing else.""" + calls: list[str] = [] + rollout = _rollout(tmp_path, snapshot_stages={"pre-verify"}) + env = FakeEnv(calls) + rollout._environment, rollout._env = env, FakeSnapSandbox(calls) + _fake_start_boundary(monkeypatch, calls) + _fake_verify_boundary(monkeypatch, calls, tmp_path) + _ready_to_verify(rollout, tmp_path) + + await rollout.start() + await rollout.verify() + + assert list(rollout.stage_snapshots) == ["pre-verify"] + assert len(env.snapshots) == 1 + + +# 4. Capability gaps fail closed, and the diagnostic names the stage + + +async def test_stage_capture_names_the_stage_when_the_sandbox_cannot_snapshot( + tmp_path: Path, monkeypatch +): + """The #384 container-snapshot gate, attributed to the boundary that + needed it — and start() fails closed instead of continuing uncheckpointed.""" + calls: list[str] = [] + rollout = _rollout( + tmp_path, + snapshot_stages={"env-ready"}, + snapshot_layers={"environment", "sandbox"}, + ) + env, sandbox = FakeEnv(calls), NoSnapshotSandbox() + rollout._environment, rollout._env = env, sandbox + _fake_start_boundary(monkeypatch, calls) + + with pytest.raises(RuntimeError, match="container-level snapshot/restore") as exc: + await rollout.start() + + assert "env-ready" in str(exc.value) + assert env.snapshots == [] # the gate fires before any layer is snapshotted + assert sandbox.calls == [] + assert rollout.stage_snapshots == {} + + +async def test_stage_capture_names_the_stage_on_a_stateless_environment( + tmp_path: Path, monkeypatch +): + """A stateless ManifestEnvironment's typed error propagates, annotated with + the stage that asked for the environment layer.""" + calls: list[str] = [] + rollout = _rollout(tmp_path, snapshot_stages={"pre-verify"}) + rollout._environment = ManifestEnvironment( + _STATELESS_MANIFEST, sandbox=FakeSnapSandbox() + ) + rollout._env = FakeSnapSandbox(calls) + _fake_verify_boundary(monkeypatch, calls, tmp_path) + _ready_to_verify(rollout, tmp_path) + + with pytest.raises(RuntimeError, match="stateless") as exc: + await rollout.verify() + + assert any("pre-verify" in note for note in exc.value.__notes__) + assert "harden_before_verify" not in calls + assert rollout.stage_snapshots == {} + + +async def test_stage_capture_names_the_stage_without_an_environment_plane( + tmp_path: Path, monkeypatch +): + calls: list[str] = [] + rollout = _rollout(tmp_path, snapshot_stages={"env-ready"}) + rollout._env = FakeSnapSandbox(calls) + assert rollout._environment is None + _fake_start_boundary(monkeypatch, calls) + + with pytest.raises(RuntimeError, match="needs the Environment plane") as exc: + await rollout.start() + + assert "env-ready" in str(exc.value) + + +# 5. Manual marking — the post-research boundary + + +async def test_mark_stage_registers_post_research_at_the_current_point( + tmp_path: Path, +): + """The harness marks the mid-execute() boundary it alone can see.""" + rollout = _rollout(tmp_path) + env = FakeEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + rollout._cursor = rollout._tree.advance(rollout._cursor, Step(id="s1")) + research_node = rollout._cursor + + snap = await rollout.mark_stage("post-research") + + assert rollout.stage_snapshots == {"post-research": snap} + assert snap.stage == "post-research" + assert snap.environment_ref is env.snapshots[0] + assert research_node.state["snapshot"] is snap + + +async def test_mark_stage_takes_the_layers_it_is_given(tmp_path: Path): + rollout = _rollout(tmp_path) + env, sandbox = FakeEnv(), FakeSnapSandbox() + rollout._environment, rollout._env = env, sandbox + + snap = await rollout.mark_stage( + "post-research", snapshot_layers={"environment", "sandbox"} + ) + + assert snap.sandbox_ref is sandbox.snapshots[0] + + +class _FakeUsageGateway: + """A usage-gateway server double exposing the RFC §3.5 stage-marker seam.""" + + def __init__(self, count: int | None) -> None: + self._count = count + self.calls = 0 + + async def live_exchange_count(self) -> int | None: + self.calls += 1 + return self._count + + +async def test_mark_stage_records_the_completed_exchange_index(tmp_path: Path): + """Guards "feat(branch): record stage markers with trajectory exchange + indices": marking a stage records which LLM exchange had completed at that + moment — read from the usage gateway's drained live capture — into the + snapshot meta and into ``stage_snapshots.json``, so a stage-named replay + cut (``bench eval continue --cut-stage``) can resolve the prefix later. + """ + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + gateway = _FakeUsageGateway(7) + rollout._usage_runtime = SimpleNamespace(server=gateway) + + snap = await rollout.mark_stage("post-research") + + assert gateway.calls == 1 + assert snap.meta["exchanges_completed"] == 7 + recorded = json.loads((run_dir / "stage_snapshots.json").read_text()) + assert recorded["stages"]["post-research"]["exchanges_completed"] == 7 + + +async def test_stage_capture_records_null_when_the_gateway_cannot_count( + tmp_path: Path, +): + """An unavailable count is recorded as null, never a fabricated number. + + Both absence shapes degrade the same way: a gateway whose drained tail + answers ``None`` (it could not catch up), and a runtime whose read raises. + """ + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + + class _RaisingGateway: + async def live_exchange_count(self) -> int | None: + raise RuntimeError("tail read failed") + + rollout._usage_runtime = SimpleNamespace(server=_FakeUsageGateway(None)) + unknown = await rollout.mark_stage("post-research") + rollout._usage_runtime = SimpleNamespace(server=_RaisingGateway()) + raising = await rollout.mark_stage("pre-verify") + + assert unknown.meta["exchanges_completed"] is None + assert raising.meta["exchanges_completed"] is None + recorded = json.loads((run_dir / "stage_snapshots.json").read_text()) + assert recorded["stages"]["post-research"]["exchanges_completed"] is None + assert recorded["stages"]["pre-verify"]["exchanges_completed"] is None + + +# 6. Branching from a recorded stage + + +async def _capture_env_ready(rollout: Rollout) -> StageSnapshot: + """Record env-ready at the root, then run past it linearly.""" + snap = await rollout.mark_stage("env-ready") + rollout._cursor = rollout._tree.advance(rollout._cursor, Step(id="s1")) + return snap + + +async def test_branch_at_stage_restores_the_recorded_snapshot(tmp_path: Path): + """The children restore env-ready's world — no new checkpoint is taken, and + they fork at the node that stage was captured on, not at the cursor.""" + rollout = _rollout(tmp_path) + env = FakeEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + root = rollout._cursor + snap = await _capture_env_ready(rollout) + env.state["db"] = "post-agent" + cursor_before = rollout._cursor + + returns = iter([1.0, 0.0]) + + async def run_child(child): + return next(returns) + + value = await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + assert value == 0.5 + assert len(env.snapshots) == 1 # the stage snapshot, not one per branch + assert env.restored == [snap.environment_ref, snap.environment_ref] + assert env.state == {"db": "initial"} # rolled back to the env-ready world + assert [child.parent for child in root.children[1:]] == [root, root] + assert rollout._cursor is cursor_before + + +async def test_branch_at_stage_records_the_stage_name_in_provenance(tmp_path: Path): + """branch_stage is the stage name, never the cursor: fallback.""" + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + await _capture_env_ready(rollout) + + async def run_child(child): + return 1.0 + + await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + children_dir = run_dir / "branches" / "root" / "children" + provenance = json.loads((children_dir / "n2" / "provenance.json").read_text()) + assert provenance["branch_stage"] == "env-ready" + assert provenance["snapshot_ref"]["environment"] == "env-snap-1" + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert nodes["root"]["stage"] == "env-ready" + assert nodes["n2"]["parent"] == "root" + + +async def test_branch_at_stage_value_covers_only_this_fork(tmp_path: Path): + """The stage node already carries the linear continuation, which has no + reward — averaging it in would silently drag V toward zero.""" + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + root = rollout._cursor + await _capture_env_ready(rollout) + + async def run_child(child): + return 1.0 + + value = await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + assert len(root.children) == 3 # linear child + two branch children + assert value == 1.0 + assert root.state["value"] == 1.0 + + +async def test_branch_at_stage_on_a_stage_that_was_never_captured(tmp_path: Path): + """Fail closed with the captured stages named — never degrade to a + checkpoint at the cursor, which would fork a different world.""" + rollout = _rollout(tmp_path) + env = FakeEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + await rollout.mark_stage("env-ready") + + async def run_child(child): + return 1.0 + + with pytest.raises(BranchStageNotCaptured, match="post-research") as exc: + await rollout.branch_at_stage("post-research", 2, run_child=run_child) + + assert "['env-ready']" in str(exc.value) + assert len(env.snapshots) == 1 + assert env.restored == [] + + +async def test_branch_at_stage_rejects_layers_the_stage_did_not_capture( + tmp_path: Path, +): + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + await rollout.mark_stage("env-ready") + + async def run_child(child): + return 1.0 + + with pytest.raises(ValueError, match="disagrees with the layers"): + await rollout.branch_at_stage( + "env-ready", + 2, + run_child=run_child, + snapshot_layers={"environment", "sandbox"}, + ) + + +class RefLessEnv(FakeEnv): + """An Environment plane whose ``snapshot()`` yields no ref. + + A third-party plane, not the shipped ``ManifestEnvironment``. Nothing + between the protocol and ``checkpoint_composed`` requires the returned + handle to be non-``None``, so a plane shaped like this produces a + ``StageSnapshot`` with no layer refs at all. + """ + + async def snapshot(self) -> StateSnapshot: # type: ignore[override] + self.calls.append("env.snapshot") + return None # type: ignore[return-value] + + +async def test_branch_at_stage_rejects_a_stage_that_captured_no_layer( + tmp_path: Path, +): + """A stage branch enforces the non-empty layer check the cursor path does. + + Guards the fix from "fix(branch): stage branches enforce the non-empty + layer check". ``_resolve_layers`` — which rejects an empty layer set — ran + only on the cursor arm. On the ``at_stage`` arm the layers were *derived* + from the recorded snapshot's refs, and a snapshot carrying neither ref + derived the empty set: the capability gate then had nothing to check, + ``restore_composed(environment=None, sandbox=None)`` rolled nothing back + before each child, and provenance recorded a clean stage fork. Every child + would have run in the world the previous one left, and the fork would have + published a V computed across them. + + ``pre-verify`` rather than ``env-ready``: at ``env-ready`` the + fresh-children gate already demands the container layer, so the hole is + only reachable at a boundary whose children run in place. + """ + rollout = _rollout(tmp_path) + env = RefLessEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + await rollout.mark_stage("pre-verify") + recorded = rollout.stage_snapshots["pre-verify"] + assert recorded.environment_ref is None and recorded.sandbox_ref is None + + ran: list[str] = [] + + async def run_child(child): + ran.append(child.id) + return 1.0 + + with pytest.raises(ValueError, match="at least one layer") as excinfo: + await rollout.branch_at_stage("pre-verify", 2, run_child=run_child) + + assert "pre-verify" in str(excinfo.value) # names the fork that has no point + # fail closed: nothing restored, no child run, no fork recorded + assert env.restored == [] + assert ran == [] + assert rollout.tree.root.children == [] + + +async def test_branch_at_stage_composes_both_layers_when_the_stage_did( + tmp_path: Path, +): + """A stage captured with both layers restores sandbox-then-env per child.""" + calls: list[str] = [] + rollout = _rollout(tmp_path) + env, sandbox = FakeEnv(calls), FakeSnapSandbox(calls) + rollout._environment, rollout._env = env, sandbox + await rollout.mark_stage("env-ready", snapshot_layers={"environment", "sandbox"}) + calls.clear() + + async def run_child(child): + return 1.0 + + await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + assert calls == [ + "sandbox.restore", + "env.restore", + "sandbox.restore", + "env.restore", + ] + + +async def test_branch_at_stage_reuses_the_engine_delta_validation(tmp_path: Path): + """Stage branching is the ordinary engine path — a not-yet-executable + delta still fails closed before any child runs.""" + rollout = _rollout(tmp_path) + env = FakeEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + await rollout.mark_stage("env-ready") + + async def run_child(child): + return 1.0 + + with pytest.raises(BranchDeltaNotSupported, match="skill_mode"): + await rollout.branch_at_stage( + "env-ready", + 2, + run_child=run_child, + deltas=[None, BranchDelta(skill_mode="with-skill")], + ) + + assert env.restored == [] + + +async def test_a_childs_own_stage_capture_stays_scoped_to_the_child( + tmp_path: Path, +): + """Branch children are isolated sub-rollouts: a child running through its + own pre-verify boundary must not become the *parent's* pre-verify — a + later branch there would fork the child's world, not the parent's.""" + rollout = _rollout(tmp_path, snapshot_stages={"pre-verify"}) + env = FakeEnv() + rollout._environment, rollout._env = env, FakeSnapSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + await _capture_env_ready(rollout) + + async def run_child(child): + await rollout._capture_stage("pre-verify") + return 1.0 + + await rollout.branch_at_stage("env-ready", 2, run_child=run_child) + + assert list(rollout.stage_snapshots) == ["env-ready"] + recorded = json.loads((run_dir / "stage_snapshots.json").read_text()) + assert list(recorded["stages"]) == ["env-ready"] + + +# 7. The stage registry artifact + + +async def test_stage_snapshots_json_is_deterministic(tmp_path: Path): + """Sorted keys, per-layer refs, the layers each stage captured — and + rewritten in full on every capture so a partial run still has evidence.""" + rollout = _rollout(tmp_path) + rollout._environment, rollout._env = FakeEnv(), FakeSnapSandbox() + run_dir = tmp_path / "run" + run_dir.mkdir() + rollout._rollout_dir = run_dir + + await rollout.mark_stage("pre-verify") + partial = (run_dir / "stage_snapshots.json").read_text() + await rollout.mark_stage("env-ready", snapshot_layers={"environment", "sandbox"}) + + assert json.loads(partial)["stages"] == { + "pre-verify": { + "environment_ref": "env-snap-1", + "sandbox_ref": None, + "layers": ["environment"], + # No usage gateway is attached to this fake rollout, so the + # exchange index of the capture is recorded as an honest null + # ("feat(branch): record stage markers with trajectory exchange + # indices"). + "exchanges_completed": None, + } + } + assert (run_dir / "stage_snapshots.json").read_text() == ( + "{\n" + ' "schema_version": 1,\n' + ' "stages": {\n' + ' "env-ready": {\n' + ' "environment_ref": "env-snap-2",\n' + ' "exchanges_completed": null,\n' + ' "layers": [\n' + ' "environment",\n' + ' "sandbox"\n' + " ],\n" + ' "sandbox_ref": "bf-snap-1"\n' + " },\n" + ' "pre-verify": {\n' + ' "environment_ref": "env-snap-1",\n' + ' "exchanges_completed": null,\n' + ' "layers": [\n' + ' "environment"\n' + " ],\n" + ' "sandbox_ref": null\n' + " }\n" + " }\n" + "}\n" + ) + + +# 8. Snapshot lifetime at cleanup (RFC §3.6): stage_snapshots.json must say +# whether each recorded ref outlived the run — the P1-B finding of PR +# #1046's second review was three valid-looking bf-snap-… refs whose +# images `docker image inspect` could no longer resolve. + + +class _StoppableSandbox: + """Cleanup-facing sandbox fake: records stop/export ordering.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + + async def stop(self, delete: bool) -> None: + self.calls.append(f"stop:{delete}") + + +class _ExportableSandbox(_StoppableSandbox): + """A sandbox that can docker-save its committed snapshot images.""" + + def __init__(self, tar_bytes: bytes = b"fake-docker-save-tar") -> None: + super().__init__() + self.tar_bytes = tar_bytes + + async def export_image(self, ref: str, target_path) -> None: + self.calls.append(f"export:{ref}") + Path(target_path).write_bytes(self.tar_bytes) + + +def _docker_save_tar_bytes(config_entry: str) -> bytes: + """Minimal bytes in `docker save` layout: a manifest.json naming Config.""" + import io + import tarfile + + manifest = json.dumps( + [{"Config": config_entry, "RepoTags": ["bf-snap-x:latest"], "Layers": []}] + ).encode() + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w") as tar: + info = tarfile.TarInfo("manifest.json") + info.size = len(manifest) + tar.addfile(info, io.BytesIO(manifest)) + return buffer.getvalue() + + +def _cleanup_ready_rollout(tmp_path: Path, *, env, keep_snapshots: bool = False): + """A minimally-populated Rollout whose ``cleanup()`` runs against fakes. + + Mirrors the ``Rollout.__new__`` pattern of ``test_usage_required``: only + the attributes cleanup() touches are set, plus a recorded ``pre-verify`` + stage snapshot and its capture-time ``stage_snapshots.json``. + """ + from benchflow.branch_lineage import write_stage_snapshots + + rollout = Rollout.__new__(Rollout) + rollout._config = RolloutConfig( + task_path=tmp_path / "task", + scenes=[Scene.single(agent="dummy")], + keep_snapshots=keep_snapshots, + ) + rollout._error = None + rollout._trajectory = [] + rollout._acp_client = None + rollout._agent_launch = "" + rollout._env = env + rollout._environment = None + rollout._rollout_dir = tmp_path + rollout._stage_snapshots = { + "pre-verify": StageSnapshot( + environment_ref=None, + sandbox_ref=SandboxImage(provider="fake", ref="bf-snap-77"), + stage="pre-verify", + ) + } + write_stage_snapshots(run_dir=tmp_path, snapshots=rollout._stage_snapshots) + return rollout + + +def _recorded_stage(tmp_path: Path, stage: str = "pre-verify") -> dict: + return json.loads((tmp_path / "stage_snapshots.json").read_text())["stages"][stage] + + +async def test_cleanup_marks_destroyed_stage_refs_ephemeral(tmp_path: Path): + """Guards "feat(rollout): stage snapshots record their lifetime; + --keep-snapshots on bench eval run with a tested import path" (the + truthful half): a plain run's cleanup destroys the committed images, so + the artifact entry must say ``ephemeral: true, exported: null`` instead + of leaving a bare ref that no longer resolves.""" + sandbox = _StoppableSandbox() + rollout = _cleanup_ready_rollout(tmp_path, env=sandbox) + assert "ephemeral" not in _recorded_stage(tmp_path) # capture-time entry + + await rollout.cleanup() + + entry = _recorded_stage(tmp_path) + assert entry["ephemeral"] is True + assert entry["exported"] is None + assert "export_error" not in entry + # The recorded refs themselves are untouched — the lifetime is an + # annotation, not a rewrite of what was captured. + assert entry["sandbox_ref"] == "bf-snap-77" + assert entry["layers"] == ["sandbox"] + assert sandbox.calls == ["stop:True"] + + +async def test_cleanup_with_keep_snapshots_exports_before_the_images_die( + tmp_path: Path, +): + """The durable half: ``RolloutConfig.keep_snapshots`` docker-saves each + captured stage image into /snapshots/ *before* the sandbox stop + destroys it, recording the tar's path, sha256 and image id with + ``ephemeral: false``.""" + import hashlib + + config_hex = "ab" * 32 + tar_bytes = _docker_save_tar_bytes(f"blobs/sha256/{config_hex}") + sandbox = _ExportableSandbox(tar_bytes) + rollout = _cleanup_ready_rollout(tmp_path, env=sandbox, keep_snapshots=True) + + await rollout.cleanup() + + assert sandbox.calls == ["export:bf-snap-77", "stop:True"] + tar_path = tmp_path / "snapshots" / "bf-snap-77.tar" + assert tar_path.read_bytes() == tar_bytes + entry = _recorded_stage(tmp_path) + assert entry["ephemeral"] is False + assert entry["exported"] == { + "path": str(tar_path), + "sha256": "sha256:" + hashlib.sha256(tar_bytes).hexdigest(), + "image_id": f"sha256:{config_hex}", + } + + +async def test_cleanup_records_a_failed_export_and_still_stops(tmp_path: Path): + """A backend that cannot export (no export_image) records export_error, + keeps the entry ephemeral, and never blocks teardown.""" + sandbox = _StoppableSandbox() + rollout = _cleanup_ready_rollout(tmp_path, env=sandbox, keep_snapshots=True) + + await rollout.cleanup() + + entry = _recorded_stage(tmp_path) + assert entry["ephemeral"] is True + assert entry["exported"] is None + assert "--keep-snapshots" in entry["export_error"] + assert sandbox.calls == ["stop:True"] + assert not (tmp_path / "snapshots").exists() + + +async def test_cleanup_preserves_an_entry_an_earlier_writer_annotated( + tmp_path: Path, +): + """`bench eval ablate --keep-snapshots` exports the branched stage into + its own out-dir and annotates the parent's stage_snapshots.json before + cleanup runs; cleanup must preserve that record instead of re-marking the + exported snapshot ephemeral.""" + from benchflow.branch_lineage import annotate_stage_snapshots_file + + sandbox = _StoppableSandbox() + rollout = _cleanup_ready_rollout(tmp_path, env=sandbox) + exported_entry = { + **_recorded_stage(tmp_path), + "ephemeral": False, + "exported": { + "path": str(tmp_path / "out" / "snapshots" / "bf-snap-77.tar"), + "sha256": "sha256:" + "cd" * 32, + "image_id": "sha256:" + "ab" * 32, + }, + } + annotate_stage_snapshots_file( + run_dir=tmp_path, annotations={"pre-verify": exported_entry} + ) + + await rollout.cleanup() + + assert _recorded_stage(tmp_path) == exported_entry + + +async def test_cleanup_leaves_an_externally_owned_sandbox_and_its_file_alone( + tmp_path: Path, +): + """An externally-owned sandbox (use_prebuilt_env, #388) is not stopped, + so its snapshot images survive cleanup — marking them ephemeral would be + the opposite lie. The capture-time file stays exactly as written.""" + sandbox = _StoppableSandbox() + rollout = _cleanup_ready_rollout(tmp_path, env=sandbox) + rollout._env_externally_owned = True + before = (tmp_path / "stage_snapshots.json").read_text() + + await rollout.cleanup() + + assert (tmp_path / "stage_snapshots.json").read_text() == before + assert sandbox.calls == [] + + +def test_image_id_is_read_from_the_save_tars_manifest(tmp_path: Path): + """Both `docker save` manifest Config spellings resolve to the image id + (`sha256:` — what `docker image inspect` reports after + `docker load`); unreadable bytes are the honest null, never a guess.""" + from benchflow.branch_policy import image_id_from_export_tar + + hex_id = "12" * 32 + current = tmp_path / "current.tar" + current.write_bytes(_docker_save_tar_bytes(f"blobs/sha256/{hex_id}")) + assert image_id_from_export_tar(current) == f"sha256:{hex_id}" + + legacy = tmp_path / "legacy.tar" + legacy.write_bytes(_docker_save_tar_bytes(f"{hex_id}.json")) + assert image_id_from_export_tar(legacy) == f"sha256:{hex_id}" + + garbage = tmp_path / "garbage.tar" + garbage.write_bytes(b"fake-docker-save-tar") + assert image_id_from_export_tar(garbage) is None diff --git a/tests/test_cli_arg_validation.py b/tests/test_cli_arg_validation.py index e8d32128b..751c07149 100644 --- a/tests/test_cli_arg_validation.py +++ b/tests/test_cli_arg_validation.py @@ -96,6 +96,33 @@ def test_reasoning_effort_bogus_clean_error(tmp_path: Path): assert "Traceback (most recent call last)" not in result.output +def test_reasoning_effort_unsupported_by_agent_clean_error(tmp_path: Path): + """Guards P2-A of the PR #1046 second review: gemini declares no ACP + effort config option, so a requested effort is decidable at planning — + it must be rejected before any provisioning, not after the sandbox was + built and the agent installed.""" + task = _task_dir(tmp_path) + with patch.object(Evaluation, "run", new=_fake_run_pass): + result = CliRunner().invoke( + app, + [ + "eval", + "create", + "--tasks-dir", + str(task), + "--agent", + "gemini", + "--sandbox", + "docker", + "--reasoning-effort", + "high", + ], + ) + assert result.exit_code == 1 + assert "not supported by agent 'gemini'" in result.stderr + assert "Traceback (most recent call last)" not in result.output + + def test_tasks_dir_missing_clean_error(tmp_path: Path): with patch.object(Evaluation, "run", new=_fake_run_pass): result = CliRunner().invoke( diff --git a/tests/test_cli_docs_drift.py b/tests/test_cli_docs_drift.py index 5d27d001a..15646e389 100644 --- a/tests/test_cli_docs_drift.py +++ b/tests/test_cli_docs_drift.py @@ -179,6 +179,7 @@ def test_documented_subcommands_exist() -> None: ["eval", "list"], ["eval", "metrics"], ["eval", "view"], + ["eval", "import-snapshots"], ["train", "convert"], ["train", "validate"], ["train", "run"], diff --git a/tests/test_cli_live_progress.py b/tests/test_cli_live_progress.py index 9f72a3e43..c82bccf92 100644 --- a/tests/test_cli_live_progress.py +++ b/tests/test_cli_live_progress.py @@ -1249,3 +1249,57 @@ def boom(*args): Evaluation._fire_progress(boom, "task-x") # must NOT raise Evaluation._fire_progress(None) # None callback is a no-op assert seen == [("task-x",)] + + +def test_ctrf_outcome_map_reports_every_test_not_only_the_failures(tmp_path): + """Guards "feat(ablate): per-test attribution so a scalar tie cannot hide a + behavioral difference": the per-test reading of the same CTRF report the + failure line mines. Passes are evidence too — an ablation arm that flips a + test *to* passing is exactly the difference a 0.00/0.00 scalar tie hides — + and CTRF statuses pass through verbatim rather than being re-classified. + """ + from benchflow.cli._failure_evidence import artifact_test_outcomes + + rollout = tmp_path / "lake__ab12cd34" + _write_ctrf_tests( + rollout / "verifier", + [ + {"name": "::TestLake::test_trend_result", "status": "passed"}, + {"name": "::TestLake::test_dominant_factor", "status": "failed"}, + {"name": "::TestLake::test_optional", "status": "skipped"}, + ], + ) + + assert artifact_test_outcomes(rollout) == { + "test_dominant_factor": "failed", + "test_optional": "skipped", + "test_trend_result": "passed", + } + # No CTRF report at all is "not observed", never an empty map: the caller + # has to be able to say "scalar-only attribution" instead of "all tie". + assert artifact_test_outcomes(tmp_path / "missing__ab12cd34") is None + + +def test_ctrf_outcome_map_keeps_raw_node_ids_when_short_names_collide(tmp_path): + """Same-named tests in two files must not collapse into one row. + + Guards "feat(ablate): per-test attribution …": the display shortening that + keeps the failure line inside its char budget is only safe while it stays + injective — collapsing ``a.py::test_x`` and ``b.py::test_x`` would hide one + arm's outcome behind the other's. + """ + from benchflow.cli._failure_evidence import artifact_test_outcomes + + rollout = tmp_path / "collide__ab12cd34" + _write_ctrf_tests( + rollout / "verifier", + [ + {"name": "tests/test_a.py::test_x", "status": "passed"}, + {"name": "tests/test_b.py::test_x", "status": "failed"}, + ], + ) + + assert artifact_test_outcomes(rollout) == { + "tests/test_a.py::test_x": "passed", + "tests/test_b.py::test_x": "failed", + } diff --git a/tests/test_eval_worker_retry.py b/tests/test_eval_worker_retry.py index eca1c78e2..919ab43aa 100644 --- a/tests/test_eval_worker_retry.py +++ b/tests/test_eval_worker_retry.py @@ -74,3 +74,21 @@ def test_evaluation_config_retry_excludes_provider_auth_by_default(): eval_cfg = _evaluation_config({"tasks_dir": "/tmp/tasks", "agent": "oracle"}) assert not eval_cfg.retry.should_retry(_PROVIDER_AUTH_ERROR) assert not eval_cfg.retry.should_retry(_PROVIDER_RATE_LIMIT_ERROR) + + +def test_request_global_rejections_are_never_retried(): + """Guards P2-A of the PR #1046 second review: a request-global setting + rejection (unsupported --reasoning-effort/model for the agent) is + deterministic — retrying re-runs the identical doomed configuration. The + category is excluded by default, the #917/provider_auth pattern.""" + from benchflow._utils.scoring import REQUEST_GLOBAL, classify_error + + error = ( + "ACPRequestGlobalError: request-global setting rejected: " + "reasoning_effort='high' was requested for agent 'gemini', but that " + "agent does not declare an ACP effort config option" + ) + assert classify_error(error) == REQUEST_GLOBAL + cfg = RetryConfig.from_mapping(None) + assert REQUEST_GLOBAL in cfg.exclude_categories + assert not cfg.should_retry(error) diff --git a/tests/test_live_llm_trajectory.py b/tests/test_live_llm_trajectory.py index 7ae3419cd..3dadfd331 100644 --- a/tests/test_live_llm_trajectory.py +++ b/tests/test_live_llm_trajectory.py @@ -659,3 +659,80 @@ async def test_partial_read_of_a_growing_log_is_recorded_as_measured_lag( assert process._live_capture_lag_bytes == len(record) - 128 assert process._live_capture_stall_ticks == 0 # advancing, so not a stall assert process._live_capture_stall_warned is False + + +# ── live_exchange_count: the stage-marker read (RFC §3.5) ───────────────── + + +@pytest.mark.asyncio +async def test_live_exchange_count_drains_the_tail_before_counting(tmp_path): + """Guards "feat(branch): record stage markers with trajectory exchange + indices": the stage-marker read must not report a lagging mirror's count. + + Three records sit in the gateway log but the poller never ran (the armed + capture's task is cancelled), so the mirror holds zero exchanges when the + read starts. ``live_exchange_count`` drains to EOF itself and returns 3 — + the index a stage-named replay cut can trust. + """ + log = ( + _callback_line(content="one") + + _callback_line(content="two") + + _callback_line(content="three") + ).encode() + process = _sandbox_process(_SandboxWithCallbackLog(log)) + await _arm_capture(process, tmp_path / "trajectory" / "llm_trajectory.jsonl") + + assert await process.live_exchange_count() == 3 + + +@pytest.mark.asyncio +async def test_live_exchange_count_is_none_before_capture_starts(): + """No live capture, no count — never a fabricated zero.""" + process = _sandbox_process(_SandboxWithCallbackLog(b"")) + + assert await process.live_exchange_count() is None + + +@pytest.mark.asyncio +async def test_live_exchange_count_refuses_a_stale_lower_bound(tmp_path, monkeypatch): + """A tail that cannot catch up within the drain budget answers None. + + A stage tag that undercounts would make ``--cut-stage`` silently replay a + shorter prefix than the stage boundary saw — worse than no tag. Shrink the + per-poll read budget below one record and the drain budget to one poll, so + the tail is guaranteed to still be behind when the budget runs out. + """ + monkeypatch.setattr(runtime_mod, "_LIVE_CAPTURE_CHUNK_BYTES", 64) + monkeypatch.setattr(runtime_mod, "_LIVE_CAPTURE_MAX_READS_PER_TICK", 1) + monkeypatch.setattr(runtime_mod, "_LIVE_EXCHANGE_DRAIN_ATTEMPTS", 1) + log = (_callback_line(content="x" * 400) * 3).encode() + process = _sandbox_process(_SandboxWithCallbackLog(log)) + await _arm_capture(process, tmp_path / "trajectory" / "llm_trajectory.jsonl") + + assert await process.live_exchange_count() is None + + +@pytest.mark.asyncio +async def test_live_exchange_count_counts_failure_records_too(tmp_path): + """Failure records are exchanges in llm_trajectory.jsonl, so the index + must count them — the replay prefix is counted on the same per-line basis.""" + failure = json.dumps( + { + "event": "failure", + "request": { + "method": "POST", + "path": "/v1/chat/completions", + "body": {"messages": [{"role": "user", "content": "hello"}]}, + }, + "error": {"status_code": 429, "type": "rate_limit"}, + "start_time": "2026-07-11T00:00:00Z", + "end_time": "2026-07-11T00:00:01Z", + "duration_ms": 1000, + }, + separators=(",", ":"), + ) + log = (_callback_line(content="ok") + failure + "\n").encode() + process = _sandbox_process(_SandboxWithCallbackLog(log)) + await _arm_capture(process, tmp_path / "trajectory" / "llm_trajectory.jsonl") + + assert await process.live_exchange_count() == 2 diff --git a/tests/test_loop_strategies.py b/tests/test_loop_strategies.py index bbdfb2dfd..63c13648c 100644 --- a/tests/test_loop_strategies.py +++ b/tests/test_loop_strategies.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest from benchflow.contracts import RoundResult @@ -392,7 +394,12 @@ def test_from_legacy_overrides_document_user_with_warning( monkeypatch.setattr( rollout_config, "_task_document_user_runtime", - lambda *a, **kw: (PassthroughUser(), 4), + # The compiled-runtime shape _task_document_user_runtime returns + # since the forked-snapshot stage-request wiring (user, round + # budget, and the task's snapshot_stages request). + lambda *a, **kw: SimpleNamespace( + user=PassthroughUser(), max_rounds=4, snapshot_stages=frozenset() + ), ) with caplog.at_level("WARNING"): config = RolloutConfig.from_legacy( @@ -416,7 +423,12 @@ def test_single_shot_suppresses_document_user_with_warning_on_both_paths( monkeypatch.setattr( rollout_config, "_task_document_user_runtime", - lambda *a, **kw: (PassthroughUser(), 4), + # The compiled-runtime shape _task_document_user_runtime returns + # since the forked-snapshot stage-request wiring (user, round + # budget, and the task's snapshot_stages request). + lambda *a, **kw: SimpleNamespace( + user=PassthroughUser(), max_rounds=4, snapshot_stages=frozenset() + ), ) with caplog.at_level("WARNING"): direct = RolloutConfig(task_path=tmp_path, loop_strategy="single-shot") @@ -443,7 +455,9 @@ def test_document_user_fallback_identical_on_both_paths( monkeypatch.setattr( rollout_config, "_task_document_user_runtime", - lambda *a, **kw: (doc_user, 4), + lambda *a, **kw: SimpleNamespace( + user=doc_user, max_rounds=4, snapshot_stages=frozenset() + ), ) direct = RolloutConfig(task_path=tmp_path) legacy = RolloutConfig.from_legacy(task_path=tmp_path) diff --git a/tests/test_rollout_branch.py b/tests/test_rollout_branch.py index 79029cce3..1609e5220 100644 --- a/tests/test_rollout_branch.py +++ b/tests/test_rollout_branch.py @@ -11,10 +11,12 @@ from __future__ import annotations +import json from pathlib import Path import pytest +from benchflow.branch import UNSCORED_KEY from benchflow.environment.manifest import EnvironmentManifest from benchflow.environment.manifest_env import ManifestEnvironment from benchflow.environment.protocol import StateSnapshot @@ -227,6 +229,86 @@ async def fake_verify_inner(self): assert rollout._rewards == rewards_before +async def test_branch_restores_the_parents_result_bearing_state( + tmp_path: Path, monkeypatch +): + """The parent's *reported* state survives its in-place children too. + + Guards the fix from "fix(branch): restore result-bearing rollout state + after in-place children", found live by @Galius5136 on a two-arm + ``pre-verify`` ablation: ``_rewards`` was scoped, but ``_timing``, + ``_verifier_error``, ``_diagnostics`` and the native-usage counters were + not — so the parent's ``timing.json`` carried its own ``agent_execution`` + plus both arms', and its ``result.json`` carried the last child's verifier + error and diagnostics as if they were the parent's own. + + The children here mutate those fields the way the real phases do: in + place, on the shared dicts (``execute()`` accumulates ``agent_execution``, + ``_verify_rollout`` writes into the same ``_timing``), which is also what + makes the deep copy load-bearing — a captured reference would be corrupted + by the first child and restore nothing for the second. + """ + from benchflow.diagnostics import TransportClosedDiagnostic + + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + + async def fake_execute_prompts(*_a, **_kw): + return [{"role": "agent", "text": "parent-step"}], 1 + + monkeypatch.setattr(rollout._planes, "execute_prompts", fake_execute_prompts) + rollout._acp_client = object() + await rollout.execute(["p1"]) + + # The parent's own result-bearing state, as a linear run leaves it. + rollout._verifier_error = "parent verifier: 1 test failed" + rollout._error = "parent agent: idle for 600s" + rollout._diagnostics.set( + TransportClosedDiagnostic(raw_message="parent", transport_diagnosis="unknown") + ) + rollout._native_usage_metrics = {"total_tokens": 100} + timing_before = dict(rollout._timing) + assert "agent_execution" in timing_before + diagnostics_before = rollout._diagnostics.to_result_fields() + usage_before = dict(rollout._native_usage_metrics) + parent = rollout._cursor + + async def fake_connect_inner(self): + self._acp_client = object() + + async def fake_disconnect_inner(self): + self._acp_client = None + + async def fake_verify_inner(self): + # Exactly the writes the real phases make: in place, on the parent's + # own dicts, plus the two error channels verify() assigns outright. + self._timing["verification"] = 99.0 + self._native_usage_metrics["total_tokens"] += 1_000 + self._verifier_error = "child verifier: exit 137" + self._error = "child agent: transport closed" + self._diagnostics.set( + TransportClosedDiagnostic(raw_message="child", transport_diagnosis="child") + ) + self._rewards = {"reward": 1.0} + return self._rewards + + monkeypatch.setattr(Rollout, "connect", fake_connect_inner) + monkeypatch.setattr(Rollout, "disconnect", fake_disconnect_inner) + monkeypatch.setattr(Rollout, "verify", fake_verify_inner) + + assert await rollout.branch(2) == 1.0 + + # No accumulation, no inherited verdicts: every field is the parent's own. + assert rollout._timing == timing_before + assert rollout._verifier_error == "parent verifier: 1 test failed" + assert rollout._error == "parent agent: idle for 600s" + assert rollout._diagnostics.to_result_fields() == diagnostics_before + assert rollout._native_usage_metrics == usage_before + # ... and the children did run through those fields (two children, each + # writing once), so the assertions above are not passing on an empty fork. + assert [child.state["reward"] for child in parent.children] == [1.0, 1.0] + + async def test_post_branch_execute_grows_off_the_parent_cursor( tmp_path: Path, monkeypatch ): @@ -362,13 +444,16 @@ async def fake_verify(self): assert value == 1.0 -async def test_default_runner_empty_reward_falls_back_to_zero( +async def test_default_runner_empty_reward_is_unscored_not_zero( tmp_path: Path, monkeypatch ): - """SHOULD-FIX 9: the default runner returns 0.0 when verify() yields nothing. + """A child whose verify() yielded nothing is unscored — never a 0.0. - verify() can return None or an empty dict; the per-child runner must - treat both as a 0.0 return rather than crashing. + Supersedes the original SHOULD-FIX 9 expectation (both shapes fell back to + a 0.0 return). verify() returning ``None`` or ``{}`` still must not crash + the branch, but a fabricated zero is indistinguishable from a real failing + score — a lost reward would read as evidence. The child records why it is + unscored, carries no ``reward``, and V(parent) is undefined. """ rollout = _rollout(tmp_path) rollout._environment = FakeEnvironment() @@ -394,8 +479,84 @@ async def fake_verify(self): value = await rollout.branch(2) - # both children scored 0.0 (None and {} both fall back), V(parent) = 0.0 - assert value == 0.0 + # neither child was scored (None and {} are both "no reward"), so there is + # nothing to average: V(parent) is undefined, not 0.0. + assert value is None + children = list(rollout.tree.root.children) + assert len(children) == 2 + assert all("reward" not in child.state for child in children) + assert all( + "produced no verifier reward" in child.state[UNSCORED_KEY] for child in children + ) + + +async def test_a_second_fork_with_an_unscored_child_clears_the_first_forks_value( + tmp_path: Path, +): + """An undefined V must not leave the previous fork's V behind. + + Guards the fix from "fix(branch): clear a stale branch value when V is + undefined". A node can be branched more than once. The unscored path set + the local ``value = None`` and skipped *writing* ``parent.state["value"]`` + — but never removed the one the earlier fork wrote, so ``branch()`` + returned ``None`` while ``tree.json`` still published the first fork's + number as the value of a fork that has no value. The stale number is the + dangerous half: the API's ``None`` is at least honest. + """ + from benchflow.branch import UnscoredChildError + from benchflow.branch_lineage import serialize_tree + + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + parent = rollout._cursor + + async def scored(child): + return 1.0 + + async def unscored(child): + raise UnscoredChildError("the verifier never reported a reward") + + assert await rollout.branch(2, run_child=scored) == 1.0 + assert parent.state["value"] == 1.0 + + assert await rollout.branch(2, run_child=unscored) is None + + assert "value" not in parent.state + run_dir = tmp_path / "lineage" + run_dir.mkdir() + serialize_tree(rollout.tree, run_dir=run_dir) + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + assert "value" not in nodes[parent.id] + + +async def test_a_rescored_fork_still_records_its_own_value(tmp_path: Path): + """The counterpart: clearing an undefined V must not clear a defined one. + + Guards the fix from "fix(branch): clear a stale branch value when V is + undefined" against being written as an unconditional ``pop`` — a second + fork whose children *were* scored publishes its own V, overwriting the + first fork's. + """ + from benchflow.branch import UnscoredChildError + + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + parent = rollout._cursor + + async def unscored(child): + raise UnscoredChildError("the verifier never reported a reward") + + async def scored(child): + return 0.5 + + assert await rollout.branch(2, run_child=unscored) is None + assert "value" not in parent.state + + assert await rollout.branch(2, run_child=scored) == 0.5 + assert parent.state["value"] == 0.5 async def test_linear_rollout_run_never_branches(tmp_path: Path, monkeypatch): @@ -449,3 +610,445 @@ async def run_child(child): # the real restore path ran once per child: cp from the snapshot dir restore_cmds = [c for c in sandbox.exec_calls if c.startswith("cp ")] assert len(restore_cmds) == 2 + + +# The parent's on-disk evidence across a fork +# +# Children share the parent's sandbox, and the sandbox bind-mounts the parent +# rollout's agent/artifacts/verifier directories into the container +# (``DockerSandbox.__init__``: host_*_path -> /logs/*), with ``restore()`` +# deliberately replaying those mounts. So a child's writes to /logs/verifier +# land in the *parent's* verifier directory on the host. These tests stand in +# for that at the unit layer: the child runner writes to the same host paths a +# mounted container would. + + +def _mounted_rollout(tmp_path: Path) -> Rollout: + """A rollout whose run directory exists, with the parent's own evidence.""" + from benchflow.task.paths import RolloutPaths + + rollout = _rollout(tmp_path) + run_dir = tmp_path / "run" + rollout._rollout_dir = run_dir + paths = RolloutPaths(rollout_dir=run_dir) + paths.mkdir() + (paths.verifier_dir / "reward.txt").write_text("parent-1.0\n") + (paths.verifier_dir / "test-stdout.txt").write_text("parent verifier output\n") + (paths.agent_dir / "acp_trajectory.jsonl").write_text('{"parent": true}\n') + (paths.artifacts_dir / "answer.md").write_text("the parent's answer\n") + return rollout + + +def _mounted_writer(rollout: Rollout, marks: list[str]): + """A child runner that writes through the shared mounts, as a child does.""" + from benchflow.task.paths import RolloutPaths + + paths = RolloutPaths(rollout_dir=rollout._rollout_dir) + order = iter(marks) + + async def run_child(child): + mark = next(order) + # What a real child does first: clear_verifier_output_dir wipes + # /logs/verifier — which is the parent's directory — then the verifier + # writes the child's own files there. + for entry in sorted(paths.verifier_dir.iterdir()): + entry.unlink() + (paths.verifier_dir / "reward.txt").write_text(f"{mark}\n") + (paths.verifier_dir / f"{mark}-only.txt").write_text("child-scoped\n") + (paths.agent_dir / "acp_trajectory.jsonl").write_text( + f'{{"child": "{mark}"}}\n' + ) + (paths.artifacts_dir / "answer.md").write_text(f"{mark} answer\n") + return 1.0 + + return run_child + + +async def test_branch_children_do_not_clobber_the_parents_verifier_evidence( + tmp_path: Path, +): + """The parent's own scoring evidence must survive its children. + + Without the fix the run ends with the last child's ``reward.txt`` sitting + where the parent's belongs — and, because a child clears /logs/verifier + before writing, the parent's other files are gone entirely. No reported + number is wrong, but the parent's row in the ablation table can no longer + be audited against anything on disk. + """ + from benchflow.task.paths import RolloutPaths + + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + paths = RolloutPaths(rollout_dir=rollout._rollout_dir) + + value = await rollout.branch( + 2, run_child=_mounted_writer(rollout, ["child-a", "child-b"]) + ) + + assert value == 1.0 + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" + assert ( + paths.verifier_dir / "test-stdout.txt" + ).read_text() == "parent verifier output\n" + assert ( + paths.agent_dir / "acp_trajectory.jsonl" + ).read_text() == '{"parent": true}\n' + assert (paths.artifacts_dir / "answer.md").read_text() == "the parent's answer\n" + # nothing of the children's is left mixed into the parent's directories + assert not (paths.verifier_dir / "child-a-only.txt").exists() + assert not (paths.verifier_dir / "child-b-only.txt").exists() + + +async def test_each_child_keeps_the_output_it_wrote_to_the_shared_mounts( + tmp_path: Path, +): + """Preserving the parent must not throw the children's evidence away. + + Each child's mounted output is moved into its own artifact directory under + ``mounted/``, keyed by node id — so the two arms of an ablation each have + their verifier output, and neither is the other's. + """ + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + + await rollout.branch(2, run_child=_mounted_writer(rollout, ["child-a", "child-b"])) + + children = rollout._rollout_dir / "branches" / "root" / "children" + first = children / "n1" / "mounted" + second = children / "n2" / "mounted" + assert (first / "verifier" / "reward.txt").read_text() == "child-a\n" + assert (second / "verifier" / "reward.txt").read_text() == "child-b\n" + assert ( + first / "agent" / "acp_trajectory.jsonl" + ).read_text() == '{"child": "child-a"}\n' + assert (second / "artifacts" / "answer.md").read_text() == "child-b answer\n" + # each child's directory holds only its own files + assert not (first / "verifier" / "child-b-only.txt").exists() + assert not (second / "verifier" / "child-a-only.txt").exists() + # and the parent's entries are back at their canonical paths, with the + # transient hold directory gone (its presence would mean the fork died) + assert not (rollout._rollout_dir / "branches" / "root" / "parent").exists() + + +async def test_a_child_that_raises_still_leaves_its_mounted_output(tmp_path: Path): + """A crashed child's verifier output is often the only evidence of why — + so the hand-off happens in the same finally that records its wall clock, + and the parent's evidence is restored even though branch() propagates.""" + from benchflow.task.paths import RolloutPaths + + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + paths = RolloutPaths(rollout_dir=rollout._rollout_dir) + + async def run_child(child): + (paths.verifier_dir / "test-stdout.txt").write_text("boom: exit 137\n") + raise RuntimeError("verifier crashed") + + with pytest.raises(RuntimeError, match="verifier crashed"): + await rollout.branch(2, run_child=run_child) + + child_mounted = ( + rollout._rollout_dir / "branches" / "root" / "children" / "n1" / "mounted" + ) + assert (child_mounted / "verifier" / "test-stdout.txt").read_text() == ( + "boom: exit 137\n" + ) + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" + + +async def test_a_child_that_raises_still_leaves_the_partial_lineage(tmp_path: Path): + """A fork that died mid-way must still be auditable. + + Guards the fix from "fix(branch): persist partial lineage when a branch + child raises". A child failure other than UnscoredChildError propagates + out of _run_children, and control never reached the lineage write below — + so the run had no tree.json and no per-child provenance at all, while + ``run_ablation`` (which catches exactly this exception) went on to report + the arm that scored, the arm that raised and the arms that never ran. The + partial tree is what those reported arms are evidence *of*. + """ + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + parent = rollout._cursor + ran: list[str] = [] + + async def run_child(child): + ran.append(child.id) + if len(ran) == 1: + return 1.0 + raise RuntimeError("agent connection lost") + + with pytest.raises(RuntimeError, match="agent connection lost"): + await rollout.branch(3, run_child=run_child) + + run_dir = rollout._rollout_dir + nodes = { + node["id"]: node + for node in json.loads((run_dir / "tree.json").read_text())["nodes"] + } + # Two children were attached before the fork died; the third never was. + first, second = (child.id for child in parent.children) + assert len(parent.children) == 2 + assert nodes[first]["reward"] == 1.0 + # The failing child is in the tree with neither a reward nor an unscored + # reason, and the parent carries no V — that is what marks it partial. + assert "reward" not in nodes[second] + assert "value" not in nodes[parent.id] + + children_dir = run_dir / "branches" / parent.id / "children" + provenance = json.loads((children_dir / first / "provenance.json").read_text()) + assert provenance["kind"] == "benchflow-branch" + assert json.loads((children_dir / first / "reward.json").read_text()) == { + "reward": 1.0 + } + # The child that raised has its provenance too — it is the arm the caller + # reports as errored — and no reward.json, because it produced no reward. + assert (children_dir / second / "provenance.json").is_file() + assert not (children_dir / second / "reward.json").exists() + + +async def test_nothing_on_the_failure_path_masks_the_child_failure( + tmp_path: Path, monkeypatch +): + """The isolation half of "fix(branch): persist partial lineage when a + branch child raises": both steps taken on the way out — restoring the + parent and persisting the evidence — are best-effort, and neither may + replace the caller's real diagnosis. A full disk must not turn "agent + connection lost" into "OSError". + """ + from benchflow.rollout_branch import _LinearState + + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + + def boom(**_kwargs): + raise OSError("No space left on device") + + monkeypatch.setattr("benchflow.rollout_branch.write_branch_artifacts", boom) + + # Restore #1 is the per-child one before the child runs (it must work, or + # the child never gets to fail); #2 is the one on the failure path. + restores = {"n": 0} + real_restore = _LinearState.restore_onto + + def flaky_restore(self, target): + restores["n"] += 1 + if restores["n"] >= 2: + raise RuntimeError("restore exploded") + real_restore(self, target) + + monkeypatch.setattr(_LinearState, "restore_onto", flaky_restore) + + async def run_child(child): + raise RuntimeError("agent connection lost") + + with pytest.raises(RuntimeError, match="agent connection lost"): + await rollout.branch(2, run_child=run_child) + assert restores["n"] == 2 + + +async def test_a_rollout_without_a_run_directory_still_branches(tmp_path: Path): + """The custody is a no-op when there is nowhere to keep anything — a + branch-first rollout that never ran setup() must not start failing.""" + rollout = _rollout(tmp_path) + rollout._environment = FakeEnvironment() + + async def run_child(child): + return 0.5 + + assert await rollout.branch(2, run_child=run_child) == 0.5 + + +# Custody fails closed +# +# Guards "fix(branch): artifact custody fails closed and preserves the hold +# directory". MountedArtifacts used to catch every failure of its own moves, +# log a warning, and in release() unconditionally rmtree the hold directory — +# deleting the parent's evidence exactly when the move back had failed. + + +def _held_run_dir(tmp_path: Path): + """A run dir with the parent's evidence at its canonical mount paths.""" + from benchflow.task.paths import RolloutPaths + + run_dir = tmp_path / "run" + paths = RolloutPaths(rollout_dir=run_dir) + paths.mkdir() + (paths.verifier_dir / "reward.txt").write_text("parent-1.0\n") + (paths.agent_dir / "acp_trajectory.jsonl").write_text('{"parent": true}\n') + (paths.artifacts_dir / "answer.md").write_text("the parent's answer\n") + return run_dir, paths + + +def test_release_failure_preserves_the_hold_directory_and_raises_typed( + tmp_path: Path, monkeypatch, caplog +): + """Guards "fix(branch): artifact custody fails closed and preserves the + hold directory": a release() whose move-back fails must raise + ArtifactCustodyError, keep the hold directory (never rmtree contents that + were not confirmed moved), and log the preserved path.""" + from benchflow.branch_artifacts import ( + ArtifactCustodyError, + MountedArtifacts, + parent_hold_dir, + ) + + run_dir, _paths = _held_run_dir(tmp_path) + holder = MountedArtifacts.hold(run_dir=run_dir, parent_id="root") + hold_dir = parent_hold_dir(run_dir, "root") + assert (hold_dir / "verifier" / "reward.txt").read_text() == "parent-1.0\n" + + def boom(*_args, **_kwargs): + raise OSError("No space left on device") + + monkeypatch.setattr("benchflow.branch_artifacts.shutil.move", boom) + + with ( + caplog.at_level("ERROR", logger="benchflow.branch_artifacts"), + pytest.raises(ArtifactCustodyError) as excinfo, + ): + holder.release() + + # The parent's evidence is still on disk, in the preserved hold directory. + assert hold_dir.is_dir() + assert (hold_dir / "verifier" / "reward.txt").read_text() == "parent-1.0\n" + assert ( + hold_dir / "agent" / "acp_trajectory.jsonl" + ).read_text() == '{"parent": true}\n' + # The typed error and the log both name the preserved path. + assert excinfo.value.hold_dir == hold_dir + assert str(hold_dir) in str(excinfo.value) + assert any(str(hold_dir) in record.getMessage() for record in caplog.records) + + +def test_release_success_removes_the_hold_directory(tmp_path: Path): + """The other half of "fix(branch): artifact custody fails closed and + preserves the hold directory": a clean release restores the parent's + entries and removes the (now confirmed-empty) hold directory.""" + from benchflow.branch_artifacts import MountedArtifacts, parent_hold_dir + + run_dir, paths = _held_run_dir(tmp_path) + holder = MountedArtifacts.hold(run_dir=run_dir, parent_id="root") + + holder.release() + + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" + assert (paths.artifacts_dir / "answer.md").read_text() == "the parent's answer\n" + assert not parent_hold_dir(run_dir, "root").exists() + + +def test_hold_failure_fails_closed_and_preserves_what_was_held( + tmp_path: Path, monkeypatch +): + """Guards "fix(branch): artifact custody fails closed and preserves the + hold directory": a hold() that cannot move the parent's entries aside must + refuse to run children over them — typed error, partial hold preserved.""" + import benchflow.branch_artifacts as branch_artifacts + from benchflow.branch_artifacts import ( + ArtifactCustodyError, + MountedArtifacts, + parent_hold_dir, + ) + + run_dir, paths = _held_run_dir(tmp_path) + real_move = branch_artifacts._move_entries + + def flaky_move(source, target): + if source.name == "verifier": + raise OSError("Permission denied") + return real_move(source, target) + + monkeypatch.setattr(branch_artifacts, "_move_entries", flaky_move) + + with pytest.raises(ArtifactCustodyError) as excinfo: + MountedArtifacts.hold(run_dir=run_dir, parent_id="root") + + hold_dir = parent_hold_dir(run_dir, "root") + assert str(hold_dir) in str(excinfo.value) + # What was already held (agent, artifacts precede verifier) is preserved + # in the hold directory, not deleted and not rolled back half-way. + assert ( + hold_dir / "agent" / "acp_trajectory.jsonl" + ).read_text() == '{"parent": true}\n' + # The verifier entries never moved — still at their canonical path. + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" + + +def _fail_hand_off_moves(monkeypatch): + """Make _move_entries fail exactly on hand_off targets (children/…).""" + import benchflow.branch_artifacts as branch_artifacts + + real_move = branch_artifacts._move_entries + + def flaky_move(source, target): + if "children" in Path(target).parts: + raise OSError("No space left on device") + return real_move(source, target) + + monkeypatch.setattr(branch_artifacts, "_move_entries", flaky_move) + + +async def test_a_custody_failure_does_not_mask_the_childs_own_exception( + tmp_path: Path, monkeypatch, caplog +): + """Guards "fix(branch): artifact custody fails closed and preserves the + hold directory" (the invariant half): a hand_off failure while the child's + own exception is propagating is recorded and logged — with the preserved + path — never raised over the child's failure, which stays the caller's + diagnosis. The parent's evidence still comes back.""" + from benchflow.task.paths import RolloutPaths + + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + paths = RolloutPaths(rollout_dir=rollout._rollout_dir) + _fail_hand_off_moves(monkeypatch) + + async def run_child(child): + (paths.verifier_dir / "child-only.txt").write_text("child evidence\n") + raise RuntimeError("verifier crashed") + + with ( + caplog.at_level("ERROR", logger="benchflow.branch_artifacts"), + pytest.raises(RuntimeError, match="verifier crashed"), + ): + await rollout.branch(2, run_child=run_child) + + # The custody failure was logged, naming the preserved hold path. + hold_dir = rollout._rollout_dir / "branches" / "root" / "parent" + assert any(str(hold_dir) in record.getMessage() for record in caplog.records) + # And the parent's own evidence is back at its canonical path. + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" + + +async def test_a_hand_off_failure_without_a_child_exception_fails_the_fork_closed( + tmp_path: Path, monkeypatch +): + """Guards "fix(branch): artifact custody fails closed and preserves the + hold directory": when no child failure is in flight, a hand_off failure is + raised as ArtifactCustodyError before the next child can inherit (and + destroy) the files left in the mounts. The completed child's reward is + recorded first, so the observation is not lost.""" + from benchflow.branch_artifacts import ArtifactCustodyError + from benchflow.task.paths import RolloutPaths + + rollout = _mounted_rollout(tmp_path) + rollout._environment = FakeEnvironment() + parent = rollout._cursor + paths = RolloutPaths(rollout_dir=rollout._rollout_dir) + _fail_hand_off_moves(monkeypatch) + ran: list[str] = [] + + async def run_child(child): + ran.append(child.id) + (paths.verifier_dir / "reward.txt").write_text("child-score\n") + return 1.0 + + with pytest.raises(ArtifactCustodyError): + await rollout.branch(2, run_child=run_child) + + # The fork stopped after the first child — the second never ran over the + # first one's leaked files. + assert len(ran) == 1 + assert parent.children[0].state["reward"] == 1.0 + # The parent's evidence was still restored on the way out. + assert (paths.verifier_dir / "reward.txt").read_text() == "parent-1.0\n" diff --git a/tests/test_runtime_capabilities.py b/tests/test_runtime_capabilities.py index 9763edda1..63c149a49 100644 --- a/tests/test_runtime_capabilities.py +++ b/tests/test_runtime_capabilities.py @@ -14,6 +14,7 @@ TaskConfig, TaskRuntimeView, UnsupportedTaskFeatureError, + compile_document_user_runtime, validate_task_runtime_support, ) @@ -1000,34 +1001,96 @@ def test_validator_rejects_invalid_document_user_confirmation_policy( assert "confirmation_policy" in user_issue.reason -def test_validator_rejects_forked_document_user_branch_execution( +_FORKED_SNAPSHOT_FRONTMATTER = dedent( + """\ + agents: + roles: + solver: + agent: codex + scenes: + - name: solve + roles: [solver] + user: + model: claude-haiku + stop_rule: satisfied-or-3-rounds + private_facts: + hidden_need: Use the quarterly file. + benchflow: + nudges: + mode: simulated-user + nudge_budget: 2 + branchable: true + branch_execution: forked-snapshot + """ +) + + +def test_validator_accepts_forked_snapshot_branch_execution( tmp_path: Path, ) -> None: - """Forked branch execution is not implied by preserving option kinds.""" + """A forked-snapshot task.md validates on a snapshot-capable sandbox. + + Guards "feat(task): accept branch_execution forked-snapshot now the + engine supports it": the value used to fail closed as "not implemented" + even though the Environment/sandbox snapshot branch engine ships — the + declaration now compiles to a stage-capture request instead. + """ task_dir = tmp_path / "forked-branch-user" + _write_task_md(task_dir, frontmatter=_FORKED_SNAPSHOT_FRONTMATTER) + task = Task(task_dir) + assert task.document is not None + + issues = validate_task_runtime_support( + task.document, + sandbox="docker", + task_dir=task_dir, + ) + + assert issues == [] + runtime = compile_document_user_runtime(task.document) + assert runtime.contract.status == "supported" + assert runtime.contract.branch_execution == "forked-snapshot" + # The surfaced half: the task requests the auto-capturable boundaries so + # a branch/ablation (or continue --cut-stage) can fork them later. + assert runtime.snapshot_stages == frozenset( + {"env-ready", "pre-verify", "post-verify"} + ) + + +def test_validator_rejects_forked_snapshot_on_a_snapshotless_sandbox( + tmp_path: Path, +) -> None: + """The fail-closed half the engine genuinely cannot honor: a backend whose + sandboxes cannot take container snapshots would fail the requested stage + captures at the first boundary, so validation says so before launch.""" + task_dir = tmp_path / "forked-branch-user-modal" + _write_task_md(task_dir, frontmatter=_FORKED_SNAPSHOT_FRONTMATTER) + task = Task(task_dir) + assert task.document is not None + + issues = validate_task_runtime_support( + task.document, + sandbox="modal", + task_dir=task_dir, + ) + + issue = next( + issue for issue in issues if issue.path == "benchflow.nudges.branch_execution" + ) + assert "container-snapshot-capable sandbox" in issue.reason + assert "modal" in issue.reason + + +def test_validator_rejects_forked_snapshot_without_branchable( + tmp_path: Path, +) -> None: + """forked-snapshot still requires branchable: true, like every other + branch_execution value.""" + task_dir = tmp_path / "forked-not-branchable" _write_task_md( task_dir, - frontmatter=dedent( - """\ - agents: - roles: - solver: - agent: codex - scenes: - - name: solve - roles: [solver] - user: - model: claude-haiku - stop_rule: satisfied-or-3-rounds - private_facts: - hidden_need: Use the quarterly file. - benchflow: - nudges: - mode: simulated-user - nudge_budget: 2 - branchable: true - branch_execution: forked-snapshot - """ + frontmatter=_FORKED_SNAPSHOT_FRONTMATTER.replace( + "branchable: true", "branchable: false" ), ) task = Task(task_dir) @@ -1040,8 +1103,29 @@ def test_validator_rejects_forked_document_user_branch_execution( ) user_issue = next(issue for issue in issues if issue.path == "user") - assert "forked branch execution is not implemented" in user_issue.reason - assert any(issue.path == "benchflow.nudges" for issue in issues) + assert "requires branchable: true" in user_issue.reason + + +def test_validator_rejects_unknown_branch_stages(tmp_path: Path) -> None: + """branch_stages is validated against the branch-stage taxonomy.""" + task_dir = tmp_path / "forked-bad-stages" + _write_task_md( + task_dir, + frontmatter=_FORKED_SNAPSHOT_FRONTMATTER + + " branch_stages: [env-ready, mid-flight]\n", + ) + task = Task(task_dir) + assert task.document is not None + + issues = validate_task_runtime_support( + task.document, + sandbox="docker", + task_dir=task_dir, + ) + + user_issue = next(issue for issue in issues if issue.path == "user") + assert "branch_stages" in user_issue.reason + assert "mid-flight" in user_issue.reason def test_validator_rejects_malformed_document_user_runtime_types( diff --git a/tests/test_sandbox_snapshot_contract.py b/tests/test_sandbox_snapshot_contract.py index aa014b70e..c9d9d22a1 100644 --- a/tests/test_sandbox_snapshot_contract.py +++ b/tests/test_sandbox_snapshot_contract.py @@ -225,6 +225,371 @@ async def _runner(child): assert value == pytest.approx(0.5) +# 4b. Docker restore re-creates an equivalent container + + +_LIVE_CONTAINER = { + "Id": "abc123", + "Config": {"WorkingDir": "/app", "User": "agent", "Env": ["FOO=bar"]}, + "HostConfig": { + "NetworkMode": "bf-proj_default", + "NanoCpus": 2_000_000_000, + "Memory": 4294967296, + }, + "Mounts": [ + { + "Type": "bind", + "Source": "/host/jobs/run/verifier", + "Destination": "/logs/verifier", + "RW": True, + }, + { + "Type": "bind", + "Source": "/host/jobs/run/agent", + "Destination": "/logs/agent", + "RW": True, + }, + { + "Type": "bind", + "Source": "/host/readonly", + "Destination": "/opt/fixtures", + "RW": False, + }, + { + "Type": "volume", + "Name": "pgdata", + "Source": "/var/lib/docker/volumes/pgdata/_data", + "Destination": "/var/lib/postgresql/data", + "RW": True, + }, + ], +} + + +def _docker_sandbox(tmp_path): + """A DockerSandbox instance — construction only, no daemon contact.""" + from benchflow.sandbox.docker import DockerSandbox + from benchflow.task.config import SandboxConfig + from benchflow.task.paths import RolloutPaths + + env_dir = tmp_path / "environment" + env_dir.mkdir(parents=True, exist_ok=True) + (env_dir / "Dockerfile").write_text("FROM alpine:3.20\n") + rollout_paths = RolloutPaths(rollout_dir=tmp_path / "run") + rollout_paths.mkdir() + return DockerSandbox( + environment_dir=env_dir, + environment_name="snapshot-contract", + session_id="bf-snapshot-contract", + rollout_paths=rollout_paths, + task_env_config=SandboxConfig(), + ) + + +class TestReplayedRunArgs: + """The pure reconstruction: inspect output -> ``docker run`` flags.""" + + def test_bind_mounts_are_replayed_with_their_host_paths(self): + from benchflow.sandbox.docker import _replayed_run_args + + args = _replayed_run_args(_LIVE_CONTAINER, default_network="bf-proj_default") + + assert "--mount" in args + specs = [args[i + 1] for i, a in enumerate(args) if a == "--mount"] + assert "type=bind,src=/host/jobs/run/verifier,dst=/logs/verifier" in specs + assert "type=bind,src=/host/jobs/run/agent,dst=/logs/agent" in specs + + def test_a_read_only_bind_stays_read_only(self): + from benchflow.sandbox.docker import _replayed_run_args + + args = _replayed_run_args(_LIVE_CONTAINER, default_network="bf-proj_default") + specs = [args[i + 1] for i, a in enumerate(args) if a == "--mount"] + + assert "type=bind,src=/host/readonly,dst=/opt/fixtures,readonly" in specs + + def test_a_named_volume_is_replayed_by_name_not_by_its_data_path(self): + from benchflow.sandbox.docker import _replayed_run_args + + args = _replayed_run_args(_LIVE_CONTAINER, default_network="bf-proj_default") + specs = [args[i + 1] for i, a in enumerate(args) if a == "--mount"] + + assert "type=volume,src=pgdata,dst=/var/lib/postgresql/data" in specs + + def test_resource_limits_and_the_project_network_are_replayed(self): + from benchflow.sandbox.docker import _replayed_run_args + + args = _replayed_run_args(_LIVE_CONTAINER, default_network="bf-proj_default") + + assert args[:2] == ["--network", "bf-proj_default"] + assert args[args.index("--cpus") + 1] == "2" + assert args[args.index("--memory") + 1] == "4294967296" + + def test_a_container_with_no_network_does_not_get_one_back(self): + """A task that opted out of networking must stay opted out.""" + from benchflow.sandbox.docker import _replayed_run_args + + container = {**_LIVE_CONTAINER, "HostConfig": {"NetworkMode": "none"}} + + args = _replayed_run_args(container, default_network="bf-proj_default") + + assert args[:2] == ["--network", "none"] + assert "bf-proj_default" not in args + + def test_an_empty_host_config_replays_only_the_default_network(self): + from benchflow.sandbox.docker import _replayed_run_args + + args = _replayed_run_args({}, default_network="bf-proj_default") + + assert args == ["--network", "bf-proj_default"] + + def test_a_tmpfs_mount_is_replayed_as_tmpfs(self): + from benchflow.sandbox.docker import _replayed_run_args + + container = {"Mounts": [{"Type": "tmpfs", "Destination": "/scratch"}]} + + args = _replayed_run_args(container, default_network="net") + + assert args[-2:] == ["--tmpfs", "/scratch"] + + +class TestDockerRestoreRebuildsTheContainer: + """``restore()`` must produce a container equivalent to the one it replaces. + + The live regression: the replacement was created with only ``--network`` + and two labels, so the rollout's ``/logs`` bind mounts were dropped. The + verifier then wrote ``reward.txt`` into a container-local directory, the + host saw nothing, and the branch child was reported as ``0.00``. + """ + + async def _restore_with( + self, tmp_path, monkeypatch, inspect_result, container_id="abc123" + ): + import json as _json + + from benchflow.sandbox._base import ExecResult + + sandbox = _docker_sandbox(tmp_path) + calls: list[list[str]] = [] + + async def fake_main_container_id(): + return container_id + + async def fake_docker_cli(args, check=True): + calls.append(list(args)) + if args[0] == "inspect": + if isinstance(inspect_result, int): + return ExecResult(stdout="", stderr="boom", return_code=1) + return ExecResult( + stdout=_json.dumps(inspect_result), stderr="", return_code=0 + ) + return ExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(sandbox, "_main_container_id", fake_main_container_id) + monkeypatch.setattr(sandbox, "_docker_cli", fake_docker_cli) + await sandbox.restore(SandboxImage(provider="docker", ref="bf-snap-x")) + return calls + + async def test_restore_replays_the_original_bind_mounts( + self, tmp_path, monkeypatch + ): + calls = await self._restore_with(tmp_path, monkeypatch, [_LIVE_CONTAINER]) + + run_cmd = next(call for call in calls if call[0] == "run") + specs = [run_cmd[i + 1] for i, a in enumerate(run_cmd) if a == "--mount"] + assert "type=bind,src=/host/jobs/run/verifier,dst=/logs/verifier" in specs + assert "type=bind,src=/host/jobs/run/agent,dst=/logs/agent" in specs + # the compose identity the engine relies on is still there + assert "com.docker.compose.service=main" in run_cmd + assert run_cmd[-2:] == ["sleep", "infinity"] + + async def test_the_container_is_inspected_before_it_is_removed( + self, tmp_path, monkeypatch + ): + """Ordering is the whole trick: a removed container has no host config.""" + calls = await self._restore_with(tmp_path, monkeypatch, [_LIVE_CONTAINER]) + + verbs = [call[0] for call in calls] + assert verbs.index("inspect") < verbs.index("stop") < verbs.index("rm") + + async def test_restore_fails_closed_when_the_container_cannot_be_inspected( + self, tmp_path, monkeypatch + ): + """Unknown host config is not a licence to create a container without one.""" + with pytest.raises(RuntimeError, match="docker inspect"): + await self._restore_with(tmp_path, monkeypatch, 1) + + async def test_restore_fails_closed_when_there_is_no_container_to_inspect( + self, tmp_path, monkeypatch + ): + """The other half of "fails closed" — no ``main`` container at all. + + Guards the fix from "fix(sandbox): restore fails closed when the + container cannot be inspected". The inspect-failure path above raised, + but the *unresolvable* path logged a warning and fell through to + ``replayed = ["--network", default_network]`` — a container with no + bind mounts, which is precisely the regression the mount replay was + added to fix: the verifier writes ``reward.txt`` into a container-local + ``/logs``, the host sees nothing, and the branch child is reported as + ``0.00``. A restore that cannot read the host config cannot reproduce + it, so it must not create the replacement. + """ + from benchflow.sandbox.protocol import SandboxRestoreHostConfigUnavailable + + with pytest.raises(SandboxRestoreHostConfigUnavailable) as excinfo: + await self._restore_with( + tmp_path, monkeypatch, [_LIVE_CONTAINER], container_id=None + ) + + # names what could not be resolved, not just that something failed + assert "'main'" in str(excinfo.value) + assert "bf-snap-x" in str(excinfo.value) + assert isinstance(excinfo.value, RuntimeError) + + async def test_a_restore_that_fails_closed_leaves_no_container_behind( + self, tmp_path, monkeypatch + ): + """Failing closed means nothing was created *or* destroyed. + + Guards the fix from "fix(sandbox): restore fails closed when the + container cannot be inspected". A raise that had already run ``rm -f`` + would trade a mountless container for no container at all. + """ + from benchflow.sandbox._base import ExecResult + from benchflow.sandbox.protocol import SandboxRestoreHostConfigUnavailable + + sandbox = _docker_sandbox(tmp_path) + calls: list[list[str]] = [] + + async def fake_main_container_id(): + return None + + async def fake_docker_cli(args, check=True): + calls.append(list(args)) + return ExecResult(stdout="", stderr="", return_code=0) + + monkeypatch.setattr(sandbox, "_main_container_id", fake_main_container_id) + monkeypatch.setattr(sandbox, "_docker_cli", fake_docker_cli) + + with pytest.raises(SandboxRestoreHostConfigUnavailable): + await sandbox.restore(SandboxImage(provider="docker", ref="bf-snap-x")) + + assert calls == [] + + +class TestVerifierMountDecisionIsLive: + """The "are the outputs already on the host" question is asked, not assumed.""" + + async def test_a_stale_is_mounted_declaration_falls_back_to_downloading( + self, tmp_path + ): + """``is_mounted`` is static; ``restore()`` can invalidate it mid-run.""" + from benchflow.task.paths import RolloutPaths + from benchflow.task.verifier import Verifier + + class _Sandbox: + is_mounted = True + + async def has_host_mount(self, *, host_dir, container_dir) -> bool: + return False + + rollout_paths = RolloutPaths(rollout_dir=tmp_path / "run") + rollout_paths.mkdir() + verifier = Verifier(object(), rollout_paths, _Sandbox()) + + assert await verifier._verifier_outputs_are_mounted("main") is False + + async def test_a_live_mount_still_skips_the_download(self, tmp_path): + from benchflow.task.paths import RolloutPaths + from benchflow.task.verifier import Verifier + + class _Sandbox: + is_mounted = True + + async def has_host_mount(self, *, host_dir, container_dir) -> bool: + return True + + rollout_paths = RolloutPaths(rollout_dir=tmp_path / "run") + rollout_paths.mkdir() + verifier = Verifier(object(), rollout_paths, _Sandbox()) + + assert await verifier._verifier_outputs_are_mounted("main") is True + + async def test_a_live_check_that_raises_falls_back_to_downloading(self, tmp_path): + """Fail toward the download: redundant is cheap, skipped loses the score.""" + from benchflow.task.paths import RolloutPaths + from benchflow.task.verifier import Verifier + + class _Sandbox: + is_mounted = True + + async def has_host_mount(self, *, host_dir, container_dir) -> bool: + raise RuntimeError("daemon unreachable") + + rollout_paths = RolloutPaths(rollout_dir=tmp_path / "run") + rollout_paths.mkdir() + verifier = Verifier(object(), rollout_paths, _Sandbox()) + + assert await verifier._verifier_outputs_are_mounted("main") is False + + async def test_a_backend_without_a_live_check_keeps_its_declaration(self, tmp_path): + from benchflow.task.paths import RolloutPaths + from benchflow.task.verifier import Verifier + + class _Sandbox: + is_mounted = True + + rollout_paths = RolloutPaths(rollout_dir=tmp_path / "run") + rollout_paths.mkdir() + verifier = Verifier(object(), rollout_paths, _Sandbox()) + + assert await verifier._verifier_outputs_are_mounted("main") is True + # a target service is never host-mounted, live check or not + assert await verifier._verifier_outputs_are_mounted("target") is False + + async def test_docker_has_host_mount_reads_the_live_container( + self, tmp_path, monkeypatch + ): + import json as _json + + from benchflow.sandbox._base import ExecResult + + sandbox = _docker_sandbox(tmp_path) + mounted = tmp_path / "run" / "verifier" + container = { + "Mounts": [ + { + "Type": "bind", + "Source": str(mounted), + "Destination": "/logs/verifier", + "RW": True, + } + ] + } + + async def fake_main_container_id(): + return "abc123" + + async def fake_docker_cli(args, check=True): + return ExecResult(stdout=_json.dumps([container]), stderr="", return_code=0) + + monkeypatch.setattr(sandbox, "_main_container_id", fake_main_container_id) + monkeypatch.setattr(sandbox, "_docker_cli", fake_docker_cli) + + assert ( + await sandbox.has_host_mount( + host_dir=mounted, container_dir="/logs/verifier" + ) + is True + ) + # same destination, a different host directory: not our mount + assert ( + await sandbox.has_host_mount( + host_dir=tmp_path / "elsewhere", container_dir="/logs/verifier" + ) + is False + ) + + # 5. Workspace helper is scope-renamed, alias preserved diff --git a/tests/test_snapshot_import.py b/tests/test_snapshot_import.py new file mode 100644 index 000000000..773d00bc6 --- /dev/null +++ b/tests/test_snapshot_import.py @@ -0,0 +1,315 @@ +"""The import half of ``--keep-snapshots``, and the eval-run retention flag. + +Guards "feat(rollout): stage snapshots record their lifetime; +--keep-snapshots on bench eval run with a tested import path" (PR #1046 +second review, P1-B): a completed run's exported stage snapshot must be +loadable later — sha256-verified tar, ``docker load``, and the loaded image +id checked against the recorded one — and the ``bench eval run`` flag must +reach ``RolloutConfig.keep_snapshots`` through the whole planning stack. + +Unit tests against a fake docker runner — no Docker daemon, no API keys. The +live end-to-end proof (real ``docker save`` → ``rmi`` → import) lives in +``tests/test_branch_composed_docker.py``. +""" + +from __future__ import annotations + +import hashlib +import json +import subprocess +from pathlib import Path +from typing import Any + +import pytest +from typer.testing import CliRunner + +from benchflow.cli.main import app +from benchflow.snapshot_import import ( + ImportedSnapshot, + SnapshotImportError, + import_stage_snapshots, +) + +runner = CliRunner() + +_IMAGE_ID = "sha256:" + "ab" * 32 +_REF = "bf-snap-demo-1234" + + +def _sha256(data: bytes) -> str: + return "sha256:" + hashlib.sha256(data).hexdigest() + + +def _run_dir( + tmp_path: Path, + *, + tar_bytes: bytes = b"exported-image-bytes", + image_id: str | None = _IMAGE_ID, + recorded_sha: str | None = None, + exported: bool = True, + stage: str = "pre-verify", +) -> Path: + """A completed run directory: stage_snapshots.json (+ exported tar).""" + run_dir = tmp_path / "run" + entry: dict[str, Any] = { + "environment_ref": None, + "sandbox_ref": _REF, + "layers": ["sandbox"], + "exchanges_completed": None, + "ephemeral": not exported, + "exported": None, + } + if exported: + tar_path = run_dir / "snapshots" / f"{_REF}.tar" + tar_path.parent.mkdir(parents=True, exist_ok=True) + tar_path.write_bytes(tar_bytes) + entry["exported"] = { + "path": str(tar_path), + "sha256": recorded_sha or _sha256(tar_bytes), + "image_id": image_id, + } + run_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "stage_snapshots.json").write_text( + json.dumps({"schema_version": 1, "stages": {stage: entry}}, indent=2) + "\n" + ) + return run_dir + + +class _FakeDocker: + """Records docker CLI calls; serves configurable load/inspect results.""" + + def __init__( + self, *, loaded_id: str = _IMAGE_ID, load_rc: int = 0, inspect_rc: int = 0 + ) -> None: + self.calls: list[list[str]] = [] + self.loaded_id = loaded_id + self.load_rc = load_rc + self.inspect_rc = inspect_rc + + def __call__(self, args: list[str]) -> subprocess.CompletedProcess[str]: + self.calls.append(args) + if args[0] == "load": + return subprocess.CompletedProcess( + args, self.load_rc, stdout="", stderr="boom" if self.load_rc else "" + ) + assert args[:3] == ["image", "inspect", "--format"] + return subprocess.CompletedProcess( + args, + self.inspect_rc, + stdout=f"{self.loaded_id}\n" if self.inspect_rc == 0 else "", + stderr="No such image" if self.inspect_rc else "", + ) + + +# 1. The load path: verify → load → identity-check + + +def test_import_loads_verifies_and_returns_the_recorded_image(tmp_path: Path): + run_dir = _run_dir(tmp_path) + docker = _FakeDocker() + + imported = import_stage_snapshots(run_dir, run_docker=docker) + + tar_path = run_dir / "snapshots" / f"{_REF}.tar" + assert imported == [ + ImportedSnapshot( + stage="pre-verify", + sandbox_ref=_REF, + image_id=_IMAGE_ID, + tar_path=tar_path, + ) + ] + assert docker.calls == [ + ["load", "-i", str(tar_path)], + ["image", "inspect", "--format", "{{.Id}}", _REF], + ] + + +def test_a_loaded_id_that_differs_from_the_recorded_one_fails_closed( + tmp_path: Path, +): + """The identity check is the point: a ref resolving to a *different* + image than the run snapshotted must not be reported restored.""" + other_id = "sha256:" + "ff" * 32 + run_dir = _run_dir(tmp_path) + + with pytest.raises(SnapshotImportError, match="different world"): + import_stage_snapshots(run_dir, run_docker=_FakeDocker(loaded_id=other_id)) + + +def test_a_run_recorded_without_an_image_id_still_verifies_resolution( + tmp_path: Path, +): + """image_id is best-effort at export time (an unreadable tar records + null); the import then verifies the ref resolves and reports the id + docker actually loaded — never a guess, never a skipped load check.""" + run_dir = _run_dir(tmp_path, image_id=None) + + [imported] = import_stage_snapshots(run_dir, run_docker=_FakeDocker()) + + assert imported.image_id == _IMAGE_ID + + +def test_a_tampered_tar_is_refused_before_docker_is_ever_invoked(tmp_path: Path): + run_dir = _run_dir(tmp_path, recorded_sha=_sha256(b"what-the-run-exported")) + docker = _FakeDocker() + + with pytest.raises(SnapshotImportError, match="does not match its record"): + import_stage_snapshots(run_dir, run_docker=docker) + + assert docker.calls == [] + + +def test_a_failed_docker_load_surfaces_its_stderr(tmp_path: Path): + run_dir = _run_dir(tmp_path) + + with pytest.raises(SnapshotImportError, match=r"docker load .* failed: boom"): + import_stage_snapshots(run_dir, run_docker=_FakeDocker(load_rc=1)) + + +def test_a_ref_that_does_not_resolve_after_load_fails_closed(tmp_path: Path): + run_dir = _run_dir(tmp_path) + + with pytest.raises(SnapshotImportError, match="does not resolve"): + import_stage_snapshots(run_dir, run_docker=_FakeDocker(inspect_rc=1)) + + +def test_a_relocated_run_dir_finds_the_tar_beside_the_file(tmp_path: Path): + """The export records an absolute path on the eval machine; a copied run + folder keeps the tar at /snapshots/, and the import + must find it there instead of dying on the stale absolute path.""" + run_dir = _run_dir(tmp_path) + moved = tmp_path / "copied-elsewhere" / "run" + moved.parent.mkdir() + run_dir.rename(moved) + + [imported] = import_stage_snapshots(moved, run_docker=_FakeDocker()) + + assert imported.tar_path == moved / "snapshots" / f"{_REF}.tar" + + +# 2. Honest refusals: ephemeral refs and unknown stages + + +def test_an_ephemeral_entry_fails_closed_naming_the_flag(tmp_path: Path): + """A plain run's refs are marked ephemeral at cleanup; asking to import + one must say why there is nothing to import and how to get it.""" + run_dir = _run_dir(tmp_path, exported=False) + docker = _FakeDocker() + + with pytest.raises(SnapshotImportError, match="--keep-snapshots"): + import_stage_snapshots(run_dir, stages=["pre-verify"], run_docker=docker) + with pytest.raises(SnapshotImportError, match="--keep-snapshots"): + import_stage_snapshots(run_dir, run_docker=docker) + + assert docker.calls == [] + + +def test_an_unrecorded_stage_lists_what_the_run_captured(tmp_path: Path): + run_dir = _run_dir(tmp_path) + + with pytest.raises(SnapshotImportError, match=r"\['pre-verify'\]"): + import_stage_snapshots(run_dir, stages=["env-ready"], run_docker=_FakeDocker()) + + +def test_a_directory_without_the_artifact_is_not_a_snapshot_run(tmp_path: Path): + with pytest.raises(SnapshotImportError, match=r"no stage_snapshots\.json"): + import_stage_snapshots(tmp_path, run_docker=_FakeDocker()) + + +# 3. CLI surface: bench eval import-snapshots + + +def test_cli_import_snapshots_prints_each_restored_ref(tmp_path: Path, monkeypatch): + run_dir = _run_dir(tmp_path) + monkeypatch.setattr("benchflow.snapshot_import._run_docker", _FakeDocker()) + + result = runner.invoke(app, ["eval", "import-snapshots", str(run_dir)]) + + assert result.exit_code == 0, result.output + assert _REF in result.output + assert _IMAGE_ID in result.output + + +def test_cli_import_snapshots_fails_cleanly_on_ephemeral_runs( + tmp_path: Path, monkeypatch +): + run_dir = _run_dir(tmp_path, exported=False) + monkeypatch.setattr("benchflow.snapshot_import._run_docker", _FakeDocker()) + + result = runner.invoke(app, ["eval", "import-snapshots", str(run_dir)]) + + assert result.exit_code == 1 + assert "--keep-snapshots" in result.output + + +# 4. The eval-run flag reaches RolloutConfig through the planning stack + + +def _task_dir(tmp_path: Path) -> Path: + task = tmp_path / "tasks" / "demo-task" + task.mkdir(parents=True) + (task / "task.toml").write_text('version = "1.0"\n', encoding="utf-8") + (task / "instruction.md").write_text("solve it\n", encoding="utf-8") + return task + + +def test_keep_snapshots_threads_from_eval_request_to_rollout_config( + tmp_path: Path, +): + """--keep-snapshots → EvalCreateRequest → EvaluationConfig → + task_rollout_config → RolloutConfig.keep_snapshots, and its absence means + False — the same semantics as bench eval ablate's flag.""" + from benchflow.eval_plan import EvalCreateRequest, build_eval_plan + from benchflow.evaluation import task_rollout_config + + task = _task_dir(tmp_path) + for flag in (True, False): + plan = build_eval_plan( + EvalCreateRequest( + tasks_dir=task, jobs_dir=str(tmp_path / "jobs"), keep_snapshots=flag + ) + ) + eval_config = plan.make_eval_config() + assert eval_config.keep_snapshots is flag + rollout_config = task_rollout_config( + eval_config, task, job_name="j", jobs_dir=tmp_path / "jobs" + ) + assert rollout_config.keep_snapshots is flag + + +def test_cli_keep_snapshots_reaches_the_evaluation_config(tmp_path: Path, monkeypatch): + """The flag half at the CLI surface: `bench eval run --keep-snapshots` + lands on the EvaluationConfig the batch runner receives.""" + seen: list[Any] = [] + monkeypatch.setattr( + "benchflow.cli.main.run_batch_eval", + lambda plan, tasks_dir, config: seen.append(config), + ) + task = _task_dir(tmp_path) + + result = runner.invoke( + app, + ["eval", "run", "--tasks-dir", str(task), "--keep-snapshots"], + ) + assert result.exit_code == 0, result.output + result = runner.invoke(app, ["eval", "run", "--tasks-dir", str(task)]) + assert result.exit_code == 0, result.output + + assert [config.keep_snapshots for config in seen] == [True, False] + + +def test_worker_payload_round_trips_keep_snapshots(): + """Sharded worker runs must not silently drop retention: the payload the + parent writes and the config the worker rebuilds agree on the flag.""" + from benchflow.eval_sharding import EvalShard, _config_payload + from benchflow.eval_worker import _evaluation_config + from benchflow.evaluation import EvaluationConfig + + config = EvaluationConfig(keep_snapshots=True) + shard = EvalShard(index=0, task_names=["t"], concurrency=1) + payload = _config_payload(config, shard=shard) + assert payload["keep_snapshots"] is True + assert _evaluation_config(payload).keep_snapshots is True + payload["keep_snapshots"] = False + assert _evaluation_config(payload).keep_snapshots is False diff --git a/tests/test_task_document.py b/tests/test_task_document.py index 0a396db57..4d01a03f3 100644 --- a/tests/test_task_document.py +++ b/tests/test_task_document.py @@ -1003,6 +1003,72 @@ def test_rollout_config_compiles_model_document_user_runtime( assert config.user.branch_execution == "option-kinds-preserved" +_FORKED_SNAPSHOT_TASK_MD = """--- +agents: + roles: + solver: + agent: codex +scenes: + - name: solve + roles: [solver] +user: + model: claude-haiku + stop_rule: satisfied-or-3-rounds + private_facts: + hidden_need: Use the quarterly file. +benchflow: + nudges: + mode: simulated-user + nudge_budget: 2 + branchable: true + branch_execution: forked-snapshot +--- +## prompt + +Base instruction. +""" + + +def test_rollout_config_adopts_forked_snapshot_stage_request( + tmp_path: Path, +) -> None: + """A forked-snapshot task's launch policy requests stage capture. + + Guards "feat(task): accept branch_execution forked-snapshot now the + engine supports it": the declaration must *do* something — a plain + evaluation of the task captures the auto stages (container layer, since + no Environment plane is bound), so its run folder carries the + stage_snapshots.json a branch/ablation or continue --cut-stage forks. + """ + (tmp_path / "task.md").write_text(_FORKED_SNAPSHOT_TASK_MD) + + config = RolloutConfig.from_legacy(task_path=tmp_path) + + assert isinstance(config.user, ModelDocumentNudgeUser) + assert config.user.branch_execution == "forked-snapshot" + assert config.snapshot_stages == frozenset( + {"env-ready", "pre-verify", "post-verify"} + ) + assert config.snapshot_layers == frozenset({"sandbox"}) + + +def test_rollout_config_explicit_snapshot_request_beats_the_tasks( + tmp_path: Path, +) -> None: + """The run-level request wins: a caller that set snapshot_stages keeps + exactly what it asked for, layers included.""" + (tmp_path / "task.md").write_text(_FORKED_SNAPSHOT_TASK_MD) + + config = RolloutConfig.from_legacy( + task_path=tmp_path, + snapshot_stages={"env-ready"}, + snapshot_layers={"environment", "sandbox"}, + ) + + assert config.snapshot_stages == frozenset({"env-ready"}) + assert config.snapshot_layers == frozenset({"environment", "sandbox"}) + + def test_rollout_config_compiles_multi_scene_document_user_runtime( tmp_path: Path, ) -> None: