From 84a550e6351724ce8315c3d958dfe32edf2d5094 Mon Sep 17 00:00:00 2001 From: Dongkeun Lee Date: Fri, 17 Jul 2026 23:02:35 -0400 Subject: [PATCH 01/54] feat: import and harden Agent Fleet --- .github/workflows/ci.yml | 20 + .no-mistakes.yaml | 19 +- AGENTS.md | 5 +- CONTRIBUTING.md | 15 +- README.md | 3 +- docs/configuration.md | 2 +- tools/agent-fleet/.gitignore | 6 + tools/agent-fleet/AGENTS.md | 17 + tools/agent-fleet/LICENSE | 21 + tools/agent-fleet/PROVENANCE.md | 26 + tools/agent-fleet/README.md | 226 ++++ tools/agent-fleet/RELEASING.md | 37 + tools/agent-fleet/pyproject.toml | 37 + tools/agent-fleet/src/agent_fleet/__init__.py | 3 + tools/agent-fleet/src/agent_fleet/__main__.py | 4 + tools/agent-fleet/src/agent_fleet/audit.py | 47 + tools/agent-fleet/src/agent_fleet/cli.py | 1079 +++++++++++++++++ tools/agent-fleet/src/agent_fleet/config.py | 429 +++++++ .../agent-fleet/src/agent_fleet/cooldowns.py | 57 + tools/agent-fleet/src/agent_fleet/doctor.py | 116 ++ .../agent-fleet/src/agent_fleet/enrollment.py | 488 ++++++++ tools/agent-fleet/src/agent_fleet/identity.py | 248 ++++ tools/agent-fleet/src/agent_fleet/leases.py | 134 ++ tools/agent-fleet/src/agent_fleet/locks.py | 195 +++ tools/agent-fleet/src/agent_fleet/models.py | 76 ++ tools/agent-fleet/src/agent_fleet/output.py | 50 + tools/agent-fleet/src/agent_fleet/paths.py | 34 + .../agent-fleet/src/agent_fleet/providers.py | 110 ++ .../agent-fleet/src/agent_fleet/provision.py | 271 +++++ tools/agent-fleet/src/agent_fleet/quota.py | 524 ++++++++ .../agent-fleet/src/agent_fleet/scheduler.py | 381 ++++++ tools/agent-fleet/src/agent_fleet/sessions.py | 129 ++ tools/agent-fleet/src/agent_fleet/status.py | 151 +++ tools/agent-fleet/src/agent_fleet/util.py | 93 ++ tools/agent-fleet/tests/conftest.py | 109 ++ .../tests/test_config_and_provision.py | 120 ++ .../agent-fleet/tests/test_contract_status.py | 232 ++++ tools/agent-fleet/tests/test_exec_sessions.py | 291 +++++ .../agent-fleet/tests/test_quota_scheduler.py | 498 ++++++++ .../tests/test_safety_transactions.py | 352 ++++++ tools/agent-fleet/uv.lock | 108 ++ 41 files changed, 6749 insertions(+), 14 deletions(-) create mode 100644 tools/agent-fleet/.gitignore create mode 100644 tools/agent-fleet/AGENTS.md create mode 100644 tools/agent-fleet/LICENSE create mode 100644 tools/agent-fleet/PROVENANCE.md create mode 100644 tools/agent-fleet/README.md create mode 100644 tools/agent-fleet/RELEASING.md create mode 100644 tools/agent-fleet/pyproject.toml create mode 100644 tools/agent-fleet/src/agent_fleet/__init__.py create mode 100644 tools/agent-fleet/src/agent_fleet/__main__.py create mode 100644 tools/agent-fleet/src/agent_fleet/audit.py create mode 100644 tools/agent-fleet/src/agent_fleet/cli.py create mode 100644 tools/agent-fleet/src/agent_fleet/config.py create mode 100644 tools/agent-fleet/src/agent_fleet/cooldowns.py create mode 100644 tools/agent-fleet/src/agent_fleet/doctor.py create mode 100644 tools/agent-fleet/src/agent_fleet/enrollment.py create mode 100644 tools/agent-fleet/src/agent_fleet/identity.py create mode 100644 tools/agent-fleet/src/agent_fleet/leases.py create mode 100644 tools/agent-fleet/src/agent_fleet/locks.py create mode 100644 tools/agent-fleet/src/agent_fleet/models.py create mode 100644 tools/agent-fleet/src/agent_fleet/output.py create mode 100644 tools/agent-fleet/src/agent_fleet/paths.py create mode 100644 tools/agent-fleet/src/agent_fleet/providers.py create mode 100644 tools/agent-fleet/src/agent_fleet/provision.py create mode 100644 tools/agent-fleet/src/agent_fleet/quota.py create mode 100644 tools/agent-fleet/src/agent_fleet/scheduler.py create mode 100644 tools/agent-fleet/src/agent_fleet/sessions.py create mode 100644 tools/agent-fleet/src/agent_fleet/status.py create mode 100644 tools/agent-fleet/src/agent_fleet/util.py create mode 100644 tools/agent-fleet/tests/conftest.py create mode 100644 tools/agent-fleet/tests/test_config_and_provision.py create mode 100644 tools/agent-fleet/tests/test_contract_status.py create mode 100644 tools/agent-fleet/tests/test_exec_sessions.py create mode 100644 tools/agent-fleet/tests/test_quota_scheduler.py create mode 100644 tools/agent-fleet/tests/test_safety_transactions.py create mode 100644 tools/agent-fleet/uv.lock diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 15511f7f5d7..1fe998c7e29 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -59,6 +59,26 @@ jobs: "$test_script" done + agent-fleet: + name: Agent Fleet package + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: "3.11" + - name: Install pinned uv + run: python -m pip install uv==0.9.10 + - name: Verify locked package + working-directory: tools/agent-fleet + run: | + set -eu + uv sync --locked + uv run --locked ruff check . + uv run --locked pytest + uv run --locked python -m compileall -q src + uv build --out-dir "$RUNNER_TEMP/agent-fleet-dist" + invariants: name: Repo invariants runs-on: ubuntu-latest diff --git a/.no-mistakes.yaml b/.no-mistakes.yaml index 8fc97893ab2..ac5e3c33f34 100644 --- a/.no-mistakes.yaml +++ b/.no-mistakes.yaml @@ -14,16 +14,17 @@ disable_project_settings: true # commands.lint, the gate's lint step never ran the deterministic # `shellcheck bin/*.sh bin/backends/*.sh tests/*.sh` that CI runs, so info-level # ShellCheck findings (e.g. SC2015) were not surfaced locally before CI rejected -# them. commands.lint delegates to bin/fm-lint.sh, the single owner of the lint -# definition that .github/workflows/ci.yml also invokes, so local can never -# diverge from CI again (parity asserted by tests/fm-lint.test.sh). -# The test command mirrors .github/workflows/ci.yml: iterate every -# tests/*.test.sh, run each, and fail the step if any one exits non-zero (an -# agent-driven test step has crashed the daemon). The e2e tests need tmux on -# PATH, which the firstmate environment provides. +# them. commands.lint delegates the shell definition to bin/fm-lint.sh and runs +# the locked Agent Fleet ruff check that .github/workflows/ci.yml also invokes. +# The shell definition stays single-owned, so local and CI cannot diverge +# (parity asserted by tests/fm-lint.test.sh). +# The test command mirrors the source checks in .github/workflows/ci.yml: +# iterate every tests/*.test.sh, then run Agent Fleet's locked pytest and +# compileall checks, and fail if any check exits non-zero. The e2e tests need +# tmux on PATH, which the firstmate environment provides. commands: - lint: 'bin/fm-lint.sh' - test: 'command -v tmux >/dev/null || { echo "tmux is required for e2e tests" >&2; exit 1; }; tmux -V; rc=0; for t in tests/*.test.sh; do echo "== $t =="; bash "$t" || rc=1; done; exit "$rc"' + lint: 'bin/fm-lint.sh && uv run --directory tools/agent-fleet --locked ruff check .' + test: 'command -v tmux >/dev/null || { echo "tmux is required for e2e tests" >&2; exit 1; }; tmux -V; rc=0; for t in tests/*.test.sh; do echo "== $t =="; bash "$t" || rc=1; done; uv run --directory tools/agent-fleet --locked pytest || rc=1; uv run --directory tools/agent-fleet --locked python -m compileall -q src || rc=1; exit "$rc"' # Keep test evidence out of this repo; it stays in a temp dir instead. test: diff --git a/AGENTS.md b/AGENTS.md index 7b822ad814d..2273fc9e193 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -44,7 +44,7 @@ Hard rules, in priority order: You may freely write to this repo itself (backlog, briefs, state, even this file when the captain approves a change). Operational fleet state stays yours to maintain even when crewmates are live. -Shared, tracked material means `AGENTS.md`, `README.md`, `CONTRIBUTING.md`, `.tasks.toml`, `.github/workflows/`, `bin/`, `.agents/skills/`, and public `skills/`. +Shared, tracked material means `AGENTS.md`, `README.md`, `CONTRIBUTING.md`, `.tasks.toml`, `.github/workflows/`, `bin/`, `.agents/skills/`, public `skills/`, and provider-neutral components under `tools/`. When one or more crewmates are in flight, delegate changes to shared, tracked material to a crewmate through the normal scout or ship machinery instead of hand-editing them yourself. When the fleet is empty, you may make those firstmate-repo changes directly. Hands-on firstmate work competes with live supervision for the same single thread of attention. @@ -73,6 +73,7 @@ README.md public overview and development notes .claude/skills symlink to .agents/skills for claude compatibility skills/ standalone public installer-facing skills, committed; not loaded by firstmate bin/ helper scripts, committed; read each script's header before first use +tools/ independently versioned provider-neutral components, committed; Agent Fleet lives under tools/agent-fleet .env optional X-mode pairing token; LOCAL, gitignored; presence-gates section 14 config/crew-harness crewmate harness override; LOCAL, gitignored; absent or "default" = same as firstmate. Inherited as the literal file: a concrete primary adapter value also controls a secondmate home's own crewmates (section 4) config/crew-dispatch.json optional crewmate dispatch profiles; LOCAL, gitignored; firstmate-maintained but human-editable natural-language rules that choose a per-task harness/model/effort profile (section 4). Inherited by secondmate homes @@ -763,7 +764,7 @@ Adjust the other sections only when the task genuinely deviates from the standar ## 12. Self-update -firstmate is its own repo behind the no-mistakes gate, so improvements to `AGENTS.md`, `bin/`, `.agents/skills/`, and public `skills/` reach `main` and then wait for each running firstmate to pull them. +firstmate is its own repo behind the no-mistakes gate, so improvements to `AGENTS.md`, `bin/`, `.agents/skills/`, public `skills/`, and `tools/` reach `main` and then wait for each running firstmate to pull them. Only `AGENTS.md`, `bin/`, and `.agents/skills/` are a running firstmate instruction surface; public `skills/` is tracked for installers and is not loaded by firstmate. When the captain invokes `/updatefirstmate` or asks to update firstmate, load the `/updatefirstmate` skill. It performs only fast-forward self-updates of firstmate and registered secondmate homes, re-reads `AGENTS.md` when needed, nudges updated live secondmates, and never touches anything under `projects/`. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 70b4b18a1ef..65aa9682450 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -34,7 +34,7 @@ See the [no-mistakes quick start](https://kunchenguid.github.io/no-mistakes/star - This repo is a template for running a firstmate orchestrator agent. `AGENTS.md` is the agent's main job description and names when to load bundled firstmate skills; `CLAUDE.md` is a symlink to it, and `.claude/skills` is a symlink to `.agents/skills`. -- Only shared material is tracked: `AGENTS.md`, `README.md`, `CONTRIBUTING.md`, `.tasks.toml`, `.github/workflows/`, `bin/`, `.agents/skills/`, and `skills/`. +- Only shared material is tracked: `AGENTS.md`, `README.md`, `CONTRIBUTING.md`, `.tasks.toml`, `.github/workflows/`, `bin/`, `.agents/skills/`, `skills/`, and provider-neutral components under `tools/`. `.agents/skills/` holds agent-loaded skills that assume a live firstmate home and carry `metadata.internal: true` so installers such as [skills.sh](https://skills.sh) hide them from discovery; `skills/` holds standalone, installer-facing public skills with no firstmate dependency (see the README's "Two-tier skill layout"). Everything personal to one captain's fleet (`.env`, `data/`, `state/`, `config/`, `projects/`, `.no-mistakes/`) is gitignored; never commit it. The root `.tasks.toml` is tracked `tasks-axi` config for `data/backlog.md`; compatible `tasks-axi` is the default backend for routine backlog mutations, with the compatibility definition owned by [`docs/configuration.md`](docs/configuration.md) ("Backlog backend"). @@ -76,6 +76,19 @@ for test_script in tests/*.test.sh; do bash "$test_script"; done # behavior te tmp=$(mktemp -d) && printf 'done: smoke\n' > "$tmp/smoke.status" && FM_STATE_OVERRIDE="$tmp" FM_SIGNAL_GRACE=1 FM_POLL=1 FM_HEARTBEAT=999999 bin/fm-watch-arm.sh # watcher re-arm smoke test (prints arm status, then an actionable signal) ``` +Agent Fleet is independently packaged under `tools/agent-fleet` and requires Python 3.11 or newer plus `uv`. +Run its locked checks from that directory before pushing: + +```sh +uv sync --locked +uv run --locked ruff check . +uv run --locked pytest +uv run --locked python -m compileall -q src +uv build --out-dir dist +``` + +The complete versioning, tagging, and clean-install verification procedure lives in [`tools/agent-fleet/RELEASING.md`](tools/agent-fleet/RELEASING.md). + Discover tests by listing `tests/*.test.sh`: each is a self-contained bash script named `.test.sh`, and its header comment describes what it covers, so run one directly to focus on a subject. Tests that need a real optional backend, an explicit opt-in, or an ambient toolchain capability (real herdr/zellij/cmux smoke tests, the live Pi regression, the Pi TypeScript-extension checks when node cannot import `.ts` modules directly) skip themselves and print the tool or environment gate needed to enable them, so the run-all loop above is always safe. diff --git a/README.md b/README.md index a6295b0417d..65ca3a883e0 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,7 @@ Launching a supported harness inside it instantiates your first mate - and makes - **Two task shapes** - ship tasks deliver a change; scout tasks investigate, plan, reproduce, or audit and leave a report. - **Explicit project modes** - each project ships via `no-mistakes`, `direct-PR`, or `local-only`, with an optional `+yolo` autonomy flag. - **Optional secondmates** - opt in to persistent domain supervisors that run from isolated firstmate homes with their own `FM_HOME`, state, projects, and session lock, supervising project clones or a project-less firstmate-repo domain, kept on the primary firstmate version by guarded local fast-forwards and checked for live agent processes at session start. -- **Optional multi-account routing** - route Claude and Codex crews through explicit or pooled Agent Fleet profiles while preserving sticky recovery and provider-neutral continuation. +- **Optional multi-account routing** - route Claude and Codex crews through explicit or pooled [Agent Fleet](tools/agent-fleet/) profiles while preserving sticky recovery and provider-neutral continuation. - **Durable completion reports** - every new ship and scout task publishes a detailed, visual, searchable report to one machine-global stack through fail-closed teardown. - **Event-driven, zero-token supervision** - a bash watcher sleeps on the fleet and wakes the first mate only when something needs you; verified primary harnesses also get a turn-end backstop that blocks or follows up on a blind stop when work is in flight and supervision is not live. - **Optional X mode** - opt in with one local `.env` token so firstmate can answer your public `@myfirstmate` mentions, act on normal reversible mention requests through the same lifecycle as chat requests, acknowledge spawned work, and post up to three public-safe completion follow-ups within seven days for genuine milestones and the final outcome without changing non-X behavior; dry-run preview records would-be replies and dismissals locally before go-live. @@ -200,6 +200,7 @@ Firstmate's skills live in two separate places with different audiences: - [docs/codex-app-backend.md](docs/codex-app-backend.md) - Codex App backend boundary, evidence, and rollout contract. - [docs/turnend-guard.md](docs/turnend-guard.md) - the primary session's structural "no turn ends blind" backstop: verified per-harness hook mechanisms, scoping, loop safety, and fail-open tradeoffs. - [docs/report-stack.md](docs/report-stack.md) - completion-report requirements, publication safety, stored artifacts, and browsing commands. +- [tools/agent-fleet/](tools/agent-fleet/) - public source, installation, tests, and release procedure for the provider-neutral account router used by optional multi-account routing. - [docs/supervision-protocols/](docs/supervision-protocols/) - rendered primary-harness watcher protocols for Claude, Codex, OpenCode, Pi, Grok, and unknown harness fallback. - [docs/scripts.md](docs/scripts.md) - the `bin/` toolbelt reference. - [`AGENTS.md`](AGENTS.md) - the distro's core instruction file and the first mate's full operating manual. diff --git a/docs/configuration.md b/docs/configuration.md index bf0d5789011..77ad32f3f83 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -166,7 +166,7 @@ For Pi secondmate launches, `fm-spawn.sh` starts Pi with `-e` pointed at the sec ## Agent Fleet account routing Firstmate can route Claude and Codex launches through the machine-global `agent-fleet` CLI without reading profile homes, credentials, quota caches, or Agent Fleet state directly. -Agent Fleet is a private CLI with no Firstmate-managed installer; obtain an approved release bundle from the maintainer, follow that bundle's installation instructions, and ensure `agent-fleet` is on `PATH` before enabling routing. +Agent Fleet's public, provider-neutral source and installation instructions live under [`tools/agent-fleet`](../tools/agent-fleet/README.md); install an immutable release tag and ensure `agent-fleet` is on `PATH` before enabling routing. Account routing is default-off, so an unchanged installation makes no Agent Fleet calls, does not wrap the provider launch, and adds no managed account-routing fields to task metadata. Routing never retrofits live task metadata or migrates existing sessions; the selected runtime backend, including Herdr, remains the observation and attachment layer rather than an account authority. The effective mode resolves in this order: explicit `--account-pool` or `--account-profile` enforces routing for that spawn, `--no-account-routing` disables it for that spawn, `FM_ACCOUNT_ROUTING`, the single value in local `config/account-routing-mode`, then `off`. diff --git a/tools/agent-fleet/.gitignore b/tools/agent-fleet/.gitignore new file mode 100644 index 00000000000..912469c64f7 --- /dev/null +++ b/tools/agent-fleet/.gitignore @@ -0,0 +1,6 @@ +.venv/ +dist/ +*.egg-info/ +__pycache__/ +.pytest_cache/ +.ruff_cache/ diff --git a/tools/agent-fleet/AGENTS.md b/tools/agent-fleet/AGENTS.md new file mode 100644 index 00000000000..93626165789 --- /dev/null +++ b/tools/agent-fleet/AGENTS.md @@ -0,0 +1,17 @@ +# Agent Fleet contributor rules + +- Agent Fleet is machine-global and provider-neutral. Firstmate is a client, not a dependency. +- Work only from a feature-branch worktree. Never edit, commit, or push `main` directly. +- Never read, print, copy, back up, or commit provider credential values. +- Profile ids are stable opaque local labels; never use account emails as ids. +- Missing profiles, session mappings, leases, or auth must fail closed in enforce/resume paths. +- Selection and lease acquisition are one atomic operation using a macOS-portable lock. +- JSON is the shell integration contract. TOON is the default agent-facing structured output. +- Tests must use temporary homes and fake provider/quota binaries. They must not touch real provider homes. + +Run before commit: + +```sh +uv run pytest +uv run python -m compileall -q src +``` diff --git a/tools/agent-fleet/LICENSE b/tools/agent-fleet/LICENSE new file mode 100644 index 00000000000..e43f5a97743 --- /dev/null +++ b/tools/agent-fleet/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Dongkeun Lee + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/tools/agent-fleet/PROVENANCE.md b/tools/agent-fleet/PROVENANCE.md new file mode 100644 index 00000000000..63a5aa1c9eb --- /dev/null +++ b/tools/agent-fleet/PROVENANCE.md @@ -0,0 +1,26 @@ +# Provenance + +Agent Fleet was developed as a provider-neutral local component before its source was imported into the public Firstmate repository. +The initial imported tree was the exact tracked tree from standalone commit `863188085af088371cf1e2b29e800f1ff1533e27`. +No provider credentials, Fleet state, build outputs, caches, or virtual environments were imported. +Version 0.2.0 adds the post-import safety work in this repository: remote +identity-set verification, Desktop/base identity anchors, provider maintenance +serialization, staged and journaled Codex enrollment, fail-closed process +ownership, and login/logout isolation. After this import, +`tools/agent-fleet` in `ruby-dlee/firstmate` is the canonical source and +standalone release boundary. + +The original standalone history was authored by Dongkeun Lee: + +| Commit | Date | Subject | +| --- | --- | --- | +| `333805f651c695cee4c7eab2a04004dcdee34eae` | 2026-07-13 | Initialize Agent Fleet repository | +| `de02ab1a5195a9c39dc4a1b942a5c1816cbf778d` | 2026-07-13 | Add dynamic local agent account routing | +| `7f29a4e8aa10c851dcb1ea76a37615775c3af425` | 2026-07-13 | Verify Fleet integration health | +| `33c8552c32ede7eaad65b50622fa3768365b6000` | 2026-07-13 | Preserve unavailable quota status | +| `2392f8cbbaba064c2d439a16efaba0bdef79d68b` | 2026-07-13 | Make Fleet lifecycle recovery safe | +| `5eecf0d6ed490dcb12127fd91a789db2d5ae1a67` | 2026-07-13 | Reserve sticky account recovery | +| `add094815d4fa159ff6a86821484ba81ee4d4294` | 2026-07-17 | Fail closed on Fleet authentication readiness | +| `863188085af088371cf1e2b29e800f1ff1533e27` | 2026-07-17 | Keep Fleet version metadata aligned | + +The component is licensed under the MIT license in [LICENSE](LICENSE). diff --git a/tools/agent-fleet/README.md b/tools/agent-fleet/README.md new file mode 100644 index 00000000000..c53bbbcd01b --- /dev/null +++ b/tools/agent-fleet/README.md @@ -0,0 +1,226 @@ +# Agent Fleet + +Agent Fleet is a machine-global account routing layer for local Claude Code and +Codex CLI agents. It gives an orchestrator a dynamic pool of named account +profiles without making the orchestrator, tmux, or Herdr responsible for +credentials. + +The profile registry contains policy and opaque labels only. Provider logins +remain in private per-profile homes owned by the provider CLI. + +## Source and installation + +Agent Fleet ships as a provider-neutral component of the public Firstmate +repository under `tools/agent-fleet`. +Its original standalone history and import boundary are recorded in +[PROVENANCE.md](PROVENANCE.md). + +Install a tagged release directly from the public repository: + +```sh +uv tool install --force \ + "agent-fleet @ git+https://github.com/ruby-dlee/firstmate.git@agent-fleet-v0.2.0#subdirectory=tools/agent-fleet" +agent-fleet version +``` + +For development from a Firstmate checkout, install the local package instead: + +```sh +uv tool install --force ./tools/agent-fleet +``` + +Reinstalling is idempotent and updates the existing `agent-fleet` tool +environment in place. +It does not modify provider profile homes or Fleet state. + +## Boundaries + +The launch stack is deliberately split: + +1. Firstmate (or another client) chooses the task, provider, and account pool. +2. Agent Fleet atomically selects and leases one concrete profile. +3. tmux or Herdr hosts the process. +4. Claude Code or Codex CLI reads credentials from that profile's isolated home. + +Herdr native restore must remain disabled for multi-account panes. A bare +`codex resume` or `claude --resume` does not contain enough information to +recover the account home. Agent Fleet records task-to-profile and +task-to-provider-session mappings and resumes through the original profile. + +## Initial setup + +```sh +agent-fleet init --claude 3 --codex 2 +agent-fleet profile enroll claude-1 +agent-fleet profile enroll codex-1 +``` + +`profile enroll` is the login transaction: the profile must already be disabled +and every same-provider Fleet lease must be drained. It provisions the isolated +home and hooks, runs provider login, verifies the credential against the live +quota endpoint, and checks for duplicate provider identities. The profile stays +disabled after both success and failure. Enable it only in a separate command +after reviewing verification: + +```sh +agent-fleet profile verify claude-1 --allow-keychain-prompt +agent-fleet profile enable claude-1 +``` + +Codex enrollment uses device authorization by default and a fresh staging home, +so a cancelled or failed attempt cannot revoke the target profile's installed +credential. Complete the device page in a fresh private/Guest browser context, +close that entire context after success, and never select **Log out**. Login is +not idempotent at the provider: raw Codex login/logout can revoke a refresh +session. Agent Fleet therefore refuses enrollment while any same-provider Fleet +lease is active, verifies the staged identity, atomically promotes it, and keeps +a durable recovery journal until commit. Do not reauthenticate a Codex Desktop +account while it has live tasks. + +For ChatGPT Business or Enterprise accounts that support Codex access tokens, +the browser-independent form is: + +```sh +printenv CODEX_ACCESS_TOKEN | agent-fleet profile enroll codex-1 --access-token +``` + +The token is consumed directly by Codex from stdin; Agent Fleet never reads or +logs it. Profiles are disabled by default, and login remains the only interactive +step for browser/device authorization. + +Add or remove profiles at any time; the registry has no fixed account count: + +```sh +agent-fleet profile add claude-4 --provider claude +agent-fleet profile add codex-3 --provider codex --max-concurrent 3 +agent-fleet profile disable codex-3 +agent-fleet profile remove codex-3 +``` + +Profile ids are operational labels such as `claude-1`; emails are rejected. +Registry mutations share the provider maintenance lock with enrollment and +selection, so a concurrent `profile add`, verify, login, quota refresh, or worker +start either completes under the lock or fails closed. + +Worker profiles must use an identity distinct from the provider's base CLI and +Desktop identities. Use `manual_only` or `desktop_shared` for a human reserve; +those policies cannot join a crew pool or be enabled for worker routing. Claude +worker launches also set `DISABLE_LOGIN_COMMAND=1` and +`DISABLE_LOGOUT_COMMAND=1`, preventing an in-flight session from replacing its +managed login. Agent Fleet re-reads Claude Desktop's identity immediately before +each selection rather than trusting a TTL cache. + +## Routing contract + +Shell integrations should request compact JSON explicitly: + +```sh +agent-fleet --format json lease choose \ + --task fm:crew:example --pool codex-crew --provider codex + +agent-fleet exec \ + --task fm:crew:example --pool codex-crew --provider codex -- \ + --full-auto + +agent-fleet resume --task fm:crew:example -- --full-auto +agent-fleet --format json session status --task fm:crew:example +agent-fleet --format json lease recover --task fm:crew:example +agent-fleet --format json lease release --task fm:crew:example +agent-fleet --format json session remove --task fm:crew:example +``` + +`exec` binds the lease to its own PID and then replaces itself with the provider +CLI, preserving a verifiable owner identity. A standalone `lease choose` +creates a short reservation that must be followed by `exec`; expired or dead +owners are reclaimed under the same portable directory lock used for +selection. A reservation may be released normally when pane creation fails, +while a live worker lease requires an explicit forced release after the worker +has stopped. Worker `exec` and `resume` require a managed task id, and every live +lease and lock requires a verified process-start token; missing tokens fail +closed. A live task can never be rebound to a different process. + +Task resume is fail-closed: it requires the recorded provider session and +reuses that session's exact profile. The new-task quota reserve does not block +recovery of an existing conversation, but disabled, unprovisioned, cooled-down, +or capacity-exhausted profiles still do. + +Orchestrators that create a terminal endpoint before launching its command use +`lease recover` for an atomic, below-reserve-safe recovery reservation, then +launch `resume --task` inside the endpoint. Recovery refuses when the task still +has a live worker lease. + +TOON is the default structured output for agent-facing inspection. Use +`--format json` for shell parsing and `--format human` for a simple terminal +view. + +## Selection policy + +`quota refresh` invokes `quota-axi` inside each profile environment and stores +only normalized status, an opaque provider-identity fingerprint when available, +and percentage windows. Before every selection, old quota evidence is refreshed +automatically. `auth_required`, `rate_limited`, provider errors, duplicate +identities, missing remote verification, and expired verification are hard +blocks. A cached response whose live probe says sign-in is required is also a +hard block; cached percentages can never hide revoked credentials. + +When fresh quota exists, Agent Fleet selects the highest safe headroom after +reserve and active-lease penalties. A transient stale/unavailable response may +use weighted least-active/LRU fallback only when that exact profile has a recent +successful remote verification (24 hours by default). A fresh profile at or +below its reserve is never selected for new work. + +```sh +agent-fleet quota refresh --all +agent-fleet quota show --all +agent-fleet profile verify --all --allow-keychain-prompt +agent-fleet doctor +agent-fleet doctor --workspace ~/firstmate +``` + +`profile verify --all` remotely rechecks every profile while all profiles remain +disabled. Enable each verified worker explicitly afterward; verification and +routing activation are deliberately separate phases. On macOS, +`--allow-keychain-prompt` is the explicit one-time Claude Keychain grant; choose +**Always Allow** so later automatic checks stay non-interactive. Bare `enable` +never requests Keychain access and fails closed when verification is unavailable. + +`doctor` verifies private profile homes, the Agent Fleet SessionStart hook, +current Herdr session-identity hooks, inherited workflow hooks/assets, provider +auth state, pinned binaries, and (when `--workspace` is supplied) that both +Claude and Codex supervision hooks expose `PreToolUse` and `Stop` there. + +## Isolation and shared workflow assets + +Claude profiles set `CLAUDE_CONFIG_DIR`; Codex profiles set `CODEX_HOME` and +force file-backed credentials inside that home. Homes and state are mode 0700; +registry, lease, quota, and session files are mode 0600. Ambient provider API +key and access-token variables are removed before every provider launch. + +The provider registry can declare a base home, hook source, and simple shared +entries. Initial profiles share only account-neutral workflow assets: + +- Claude: `CLAUDE.md`, skills, plugins, and hook definitions. +- Codex: `AGENTS.md`, skills, plugins, rules, and hook definitions. + +Auth files, sessions, histories, logs, caches, databases, and provider state are +never shared. + +## State + +Defaults: + +- Registry: `~/.config/agent-fleet/accounts.toml` +- Profile homes: `~/.local/share/agent-fleet/accounts//` +- Leases, normalized quota, and session mappings: + `~/.local/state/agent-fleet` +- Neutral provider runtimes: `~/.local/libexec/agent-fleet/runtime` +- Pinned profile-aware quota reader: + `~/.local/libexec/agent-fleet/quota-axi/current/bin/quota-axi` + +All paths can be redirected in tests through `AGENT_FLEET_CONFIG`, +`AGENT_FLEET_STATE_DIR`, and `AGENT_FLEET_SHARE_DIR`. + +## Development and releases + +The locked test, lint, build, versioning, tagging, and installation verification +procedure lives in [RELEASING.md](RELEASING.md). diff --git a/tools/agent-fleet/RELEASING.md b/tools/agent-fleet/RELEASING.md new file mode 100644 index 00000000000..57a36f14790 --- /dev/null +++ b/tools/agent-fleet/RELEASING.md @@ -0,0 +1,37 @@ +# Releasing Agent Fleet + +Agent Fleet remains independently versioned and installable even though its canonical source lives inside Firstmate. +Release tags use `agent-fleet-v` and point at the Firstmate merge commit containing that exact component version. + +## Prepare the release + +1. Update `version` in `pyproject.toml` and `__version__` in `src/agent_fleet/__init__.py` to the same semantic version. +2. Run `uv lock` from this directory and commit the updated lockfile with the source change. +3. Run the complete local verification from this directory: + + ```sh + uv sync --locked + uv run --locked ruff check . + uv run --locked pytest + uv run --locked python -m compileall -q src + uv build --out-dir dist + ``` + +4. Ship the Firstmate branch through no-mistakes and merge only after the repository checks are green and the owner authorizes the merge. + +## Tag and verify the release + +1. Create the annotated `agent-fleet-v` tag at the verified Firstmate merge commit and push that tag to `ruby-dlee/firstmate`. +2. Create the GitHub release from the tag and attach both files from `tools/agent-fleet/dist/`. +3. Install from the immutable tag in a clean tool environment: + + ```sh + uv tool install --force \ + "agent-fleet @ git+https://github.com/ruby-dlee/firstmate.git@agent-fleet-v#subdirectory=tools/agent-fleet" + agent-fleet version + ``` + +4. Confirm that the reported CLI version matches the tag before announcing the release. + +The Git tag and GitHub release are distribution records only. +Provider homes, credentials, registry data, leases, quota evidence, and session mappings are never release inputs or artifacts. diff --git a/tools/agent-fleet/pyproject.toml b/tools/agent-fleet/pyproject.toml new file mode 100644 index 00000000000..33a77536357 --- /dev/null +++ b/tools/agent-fleet/pyproject.toml @@ -0,0 +1,37 @@ +[build-system] +requires = ["hatchling>=1.27"] +build-backend = "hatchling.build" + +[project] +name = "agent-fleet" +version = "0.2.0" +description = "Machine-global account profile routing for local agent CLIs" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "MIT" } +authors = [{ name = "Dongkeun Lee" }] +dependencies = [] + +[project.urls] +Repository = "https://github.com/ruby-dlee/firstmate/tree/main/tools/agent-fleet" + +[project.scripts] +agent-fleet = "agent_fleet.cli:main" + +[tool.hatch.build.targets.wheel] +packages = ["src/agent_fleet"] + +[tool.pytest.ini_options] +addopts = "-q" +testpaths = ["tests"] +pythonpath = ["src"] + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "I", "UP", "B", "SIM"] + +[dependency-groups] +dev = ["pytest>=8.3", "ruff>=0.12"] diff --git a/tools/agent-fleet/src/agent_fleet/__init__.py b/tools/agent-fleet/src/agent_fleet/__init__.py new file mode 100644 index 00000000000..7ea4575e02b --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/__init__.py @@ -0,0 +1,3 @@ +"""Agent Fleet.""" + +__version__ = "0.2.0" diff --git a/tools/agent-fleet/src/agent_fleet/__main__.py b/tools/agent-fleet/src/agent_fleet/__main__.py new file mode 100644 index 00000000000..bfdcd0c1158 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/__main__.py @@ -0,0 +1,4 @@ +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/agent-fleet/src/agent_fleet/audit.py b/tools/agent-fleet/src/agent_fleet/audit.py new file mode 100644 index 00000000000..a8fe15cbb9e --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/audit.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +import json +import os +from collections import deque +from pathlib import Path +from typing import Any + +from .models import Registry +from .paths import ensure_private_dir +from .util import utc_now + + +def audit_path(registry: Registry) -> Path: + return registry.settings.state_dir / "audit.jsonl" + + +def append_audit(registry: Registry, event: str, fields: dict[str, Any]) -> None: + path = audit_path(registry) + ensure_private_dir(path.parent) + payload = {"schema": 1, "at": utc_now(), "event": event, **fields} + encoded = (json.dumps(payload, sort_keys=True, separators=(",", ":")) + "\n").encode() + descriptor = os.open(path, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o600) + try: + os.write(descriptor, encoded) + os.fsync(descriptor) + os.fchmod(descriptor, 0o600) + finally: + os.close(descriptor) + + +def read_audit(registry: Registry, *, limit: int = 100) -> list[dict[str, Any]]: + if limit < 1 or limit > 10_000: + raise ValueError("audit limit must be between 1 and 10000") + path = audit_path(registry) + if not path.exists(): + return [] + rows: deque[dict[str, Any]] = deque(maxlen=limit) + with path.open("r", encoding="utf-8") as handle: + for line in handle: + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(value, dict): + rows.append(value) + return list(rows) diff --git a/tools/agent-fleet/src/agent_fleet/cli.py b/tools/agent-fleet/src/agent_fleet/cli.py new file mode 100644 index 00000000000..f263e958a5e --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/cli.py @@ -0,0 +1,1079 @@ +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from collections.abc import Callable, Iterator +from contextlib import ExitStack, contextmanager +from pathlib import Path +from typing import Any + +from . import __version__ +from .audit import read_audit +from .config import ( + initial_registry, + load_registry, + save_registry, + set_profile_enabled, + set_profile_safety_policy, + with_profile, + without_profile, +) +from .cooldowns import clear_cooldown, set_cooldown +from .doctor import run_doctor +from .enrollment import ( + CodexAuthTransaction, + activate_codex_promotion, + create_codex_login_stage, + discard_codex_promotion, + discard_codex_stage, + finalize_codex_promotion, + prepare_codex_promotion, + recover_pending_codex_transaction, + recover_pending_codex_transactions, + rollback_codex_promotion, +) +from .identity import ( + identity_conflict, + refresh_provider_identity_anchors, + refresh_provider_identity_anchors_if_due, +) +from .leases import active_leases, release_lease +from .locks import provider_enrollment_lock, state_lock +from .models import PROFILE_SAFETY_POLICIES, SUPPORTED_PROVIDERS, Profile, Registry +from .output import emit +from .paths import default_config_path, expand_path +from .providers import ( + auth_probe, + auth_status, + login_argv, + provider_argv, + provider_environment, + resume_argv, +) +from .provision import profile_is_provisioned, provision_profile +from .quota import ( + probe_quota, + quota_routeability, + read_quota, + refresh_due_quotas, + refresh_quota, + snapshot_quota_cache, + store_quota, +) +from .scheduler import select_and_acquire +from .sessions import get_session, read_hook_payload, record_session_from_hook, remove_session +from .status import pool_status, profile_status +from .util import validate_id + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="agent-fleet") + parser.add_argument("--config", type=Path, default=default_config_path()) + parser.add_argument( + "--format", + choices=("toon", "json", "human"), + default=os.environ.get("AGENT_FLEET_FORMAT", "toon"), + ) + commands = parser.add_subparsers(dest="command", required=True) + + init = commands.add_parser("init", help="create a disabled profile registry") + init.add_argument("--claude", type=int, default=0) + init.add_argument("--codex", type=int, default=0) + init.add_argument("--force", action="store_true") + + profile = commands.add_parser("profile") + profile_commands = profile.add_subparsers(dest="profile_command", required=True) + profile_commands.add_parser("list") + add = profile_commands.add_parser("add") + add.add_argument("profile_id") + add.add_argument("--provider", choices=SUPPORTED_PROVIDERS, required=True) + add.add_argument("--home", type=Path) + add.add_argument("--pool", action="append", dest="pools") + add.add_argument("--weight", type=int, default=1) + add.add_argument("--max-concurrent", type=int, default=2) + add.add_argument("--reserve-percent", type=int, default=15) + add.add_argument( + "--safety-policy", + choices=PROFILE_SAFETY_POLICIES, + default="worker", + ) + remove = profile_commands.add_parser("remove") + remove.add_argument("profile_id") + for name in ("enable", "disable"): + item = profile_commands.add_parser(name) + item.add_argument("profile_id") + for name in ("login", "enroll"): + item = profile_commands.add_parser(name) + item.add_argument("profile_id") + item.add_argument( + "--browser-login", + action="store_true", + help="use Codex browser callback instead of isolated device authorization", + ) + item.add_argument( + "--access-token", + action="store_true", + help="read a supported Codex access token from non-interactive stdin", + ) + enable_group = item.add_mutually_exclusive_group() + enable_group.add_argument("--enable", action="store_true") + enable_group.add_argument("--no-enable", action="store_true", help=argparse.SUPPRESS) + for name in ("provision", "status", "auth-status", "verify"): + item = profile_commands.add_parser(name) + item.add_argument("profile_id", nargs="?") + item.add_argument("--all", action="store_true") + if name == "verify": + item.add_argument("--allow-keychain-prompt", action="store_true") + cooldown = profile_commands.add_parser("cooldown") + cooldown.add_argument("profile_id") + cooldown.add_argument("--seconds", type=int, required=True) + cooldown.add_argument("--reason", required=True) + cooldown_clear = profile_commands.add_parser("cooldown-clear") + cooldown_clear.add_argument("profile_id") + policy = profile_commands.add_parser("policy") + policy.add_argument("profile_id") + policy.add_argument("--safety-policy", choices=PROFILE_SAFETY_POLICIES, required=True) + + pool = commands.add_parser("pool") + pool_commands = pool.add_subparsers(dest="pool_command", required=True) + pool_status_parser = pool_commands.add_parser("status") + pool_status_parser.add_argument("--pool", required=True) + pool_status_parser.add_argument("--provider", choices=SUPPORTED_PROVIDERS) + + choose = commands.add_parser("choose") + _add_route_arguments(choose) + choose.add_argument("--dry-run", action="store_true") + + quota = commands.add_parser("quota") + quota_commands = quota.add_subparsers(dest="quota_command", required=True) + for name in ("refresh", "show"): + item = quota_commands.add_parser(name) + item.add_argument("profile_id", nargs="?") + item.add_argument("--all", action="store_true") + if name == "refresh": + item.add_argument("--allow-keychain-prompt", action="store_true") + + lease = commands.add_parser("lease") + lease_commands = lease.add_subparsers(dest="lease_command", required=True) + lease_choose = lease_commands.add_parser("choose") + _add_route_arguments(lease_choose) + recover = lease_commands.add_parser("recover") + recover.add_argument("--task", required=True) + acquire = lease_commands.add_parser("acquire") + acquire.add_argument("--profile", required=True) + acquire.add_argument("--task", required=True) + acquire.add_argument("--pool") + lease_commands.add_parser("list") + release = lease_commands.add_parser("release") + release.add_argument("--task", required=True) + release.add_argument("--force", action="store_true") + + execute = commands.add_parser("exec") + _add_route_arguments(execute, required=False) + execute.add_argument("provider_args", nargs=argparse.REMAINDER) + + resume = commands.add_parser("resume") + resume.add_argument("--task") + resume.add_argument("--profile") + resume.add_argument("--session") + resume.add_argument("--pool") + resume.add_argument("provider_args", nargs=argparse.REMAINDER) + + session = commands.add_parser("session") + session_commands = session.add_subparsers(dest="session_command", required=True) + for name in ("status", "remove"): + item = session_commands.add_parser(name) + item.add_argument("--task", required=True) + + hook = commands.add_parser("hook") + hook_commands = hook.add_subparsers(dest="hook_command", required=True) + hook_commands.add_parser("session-start") + + doctor = commands.add_parser("doctor") + doctor.add_argument("--workspace", type=Path) + commands.add_parser("status") + commands.add_parser("contract") + commands.add_parser("version") + audit = commands.add_parser("audit") + audit.add_argument("--limit", type=int, default=100) + return parser + + +def _add_route_arguments( + parser: argparse.ArgumentParser, + *, + required: bool = True, +) -> None: + parser.add_argument("--task", required=required) + parser.add_argument("--pool", required=required) + parser.add_argument("--provider", choices=SUPPORTED_PROVIDERS) + parser.add_argument("--profile") + + +def _task(value: str) -> str: + if not value.strip() or len(value) > 512 or any(ord(char) < 32 for char in value): + raise ValueError("task id must be 1-512 printable characters") + return value + + +def _strip_separator(values: list[str]) -> list[str]: + return values[1:] if values[:1] == ["--"] else values + + +def _mutate(registry: Registry, operation: Callable[[Registry], Registry], path: Path) -> Registry: + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + current = load_registry(path) + updated = operation(current) + save_registry(updated, path) + return updated + + +@contextmanager +def _provider_maintenance( + registry: Registry, + config_path: Path, + providers: set[str], + *, + require_disabled: tuple[str, ...] = (), +) -> Iterator[Registry]: + """Serialize provider maintenance and re-read mutable registry state.""" + + try: + with ExitStack() as locks: + for provider in sorted(providers): + locks.enter_context( + provider_enrollment_lock( + registry.settings.state_dir, + provider, + registry.settings.lock_stale_seconds, + ) + ) + for provider in sorted(providers): + recover_pending_codex_transactions(registry, provider) + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + current = load_registry(config_path) + for profile_id in require_disabled: + if current.require_profile(profile_id).enabled: + raise ValueError(f"disable {profile_id} before provider maintenance") + yield current + except TimeoutError as exc: + names = ", ".join(sorted(providers)) + raise ValueError(f"provider maintenance is already in progress for {names}") from exc + + +def _profiles_for(registry: Registry, profile_id: str | None, all_profiles: bool) -> list[Profile]: + if all_profiles: + if profile_id is not None: + raise ValueError("use a profile id or --all, not both") + return [registry.profiles[key] for key in sorted(registry.profiles)] + if profile_id is None: + raise ValueError("profile id is required unless --all is used") + return [registry.require_profile(profile_id)] + + +def _contract() -> dict[str, Any]: + return { + "contract_version": 1, + "cli_version": __version__, + "formats": ["json", "toon", "human"], + "selection_fields": [ + "profile", + "provider", + "pool", + "decision_reason", + "quota_fresh", + "headroom_percent", + "active_lease_count", + "degraded", + ], + "commands": { + "pool_summary": "pool status --pool [--provider ]", + "dry_run": "choose --pool --task --dry-run", + "atomic_choose": "lease choose --pool --task ", + "explicit_acquire": "lease acquire --profile --task ", + "recover": "lease recover --task ", + "release": "lease release --task ", + "exec": "exec --profile [--task --pool ] -- ", + "resume_task": "resume --task -- ", + "resume_explicit": "managed task/session mapping required", + "session_status": "session status --task ", + "session_remove": "session remove --task ", + "profile_enroll": "profile enroll ", + "profile_verify": "profile verify |--all", + }, + } + + +def _credential_is_remotely_verified(quota: dict[str, Any]) -> bool: + fingerprint = quota.get("identity_fingerprint") + return ( + quota.get("status") == "fresh" + and quota.get("verified_at") is not None + and quota.get("headroom_percent") is not None + and isinstance(quota.get("windows"), list) + and bool(quota["windows"]) + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + ) + + +def _cached_credential_proof_is_usable(quota: dict[str, Any]) -> bool: + fingerprint = quota.get("identity_fingerprint") + return ( + quota.get("verified_recent") is True + and quota.get("status") not in {"auth_required", "rate_limited", "error"} + and quota.get("headroom_percent") is not None + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + ) + + +def _verification_reason(result: dict[str, Any]) -> str: + return str( + result.get("identity_conflict") + or result.get("remote_reason") + or result.get("remote_status") + or result.get("local_auth") + or "remote_unverified" + ) + + +def _evaluate_remote_profile( + registry: Registry, + profile: Profile, + authentication: str, + quota: dict[str, Any], + *, + require_complete_worker_set: bool = True, +) -> dict[str, Any]: + conflict = identity_conflict( + registry, + profile, + quota, + require_complete_worker_set=require_complete_worker_set, + ) + route_conflict = identity_conflict(registry, profile, quota) + remotely_verified = _credential_is_remotely_verified(quota) or ( + quota.get("fresh") is not True and _cached_credential_proof_is_usable(quota) + ) + credential_verified = ( + authentication == "authenticated" and remotely_verified and conflict is None + ) + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication, + ignore_reserve=True, + ) + if route_conflict is not None: + routeability = { + "eligible": False, + "mode": "blocked", + "reason": route_conflict, + } + return { + "profile": profile.id, + "provider": profile.provider, + "local_auth": authentication, + "remote_status": quota.get("status"), + "remote_reason": quota.get("reason"), + "credential_verified": credential_verified, + "identity_conflict": conflict, + "identity_set_block": route_conflict, + "routeability": routeability, + } + + +def _verify_remote_profile( + registry: Registry, + profile: Profile, + *, + allow_keychain_prompt: bool, + require_complete_worker_set: bool = False, +) -> dict[str, Any]: + provision_profile(registry, profile) + authentication = str(auth_probe(registry, profile)["status"]) + refresh_provider_identity_anchors( + registry, + profile.provider, + allow_keychain_prompt=allow_keychain_prompt, + ) + cached = read_quota(registry, profile.id) + candidate = probe_quota( + registry, + profile, + allow_keychain_prompt=allow_keychain_prompt, + ) + if ( + not allow_keychain_prompt + and candidate.get("reason") == "keychain_access_required" + and _cached_credential_proof_is_usable(cached) + ): + quota = cached + else: + store_quota(registry, profile, candidate) + quota = read_quota(registry, profile.id) + return _evaluate_remote_profile( + registry, + profile, + authentication, + quota, + require_complete_worker_set=require_complete_worker_set, + ) + + +def _probe_enrollment_candidate( + registry: Registry, + profile: Profile, +) -> tuple[dict[str, Any], dict[str, Any]]: + authentication = str(auth_probe(registry, profile)["status"]) + quota = probe_quota(registry, profile) + quota = {**quota, "fresh": _credential_is_remotely_verified(quota)} + result = _evaluate_remote_profile( + registry, + profile, + authentication, + quota, + require_complete_worker_set=False, + ) + return result, quota + + +def _provider_has_active_lease(registry: Registry, provider: str) -> bool: + for lease in active_leases(registry): + profile_id = lease.get("profile") + if ( + isinstance(profile_id, str) + and registry.require_profile(profile_id).provider == provider + ): + return True + return False + + +def _enroll_codex_profile( + registry: Registry, + target: Profile, + *, + browser_login: bool, + access_token: bool, +) -> dict[str, Any]: + stage = create_codex_login_stage(target) + promotion: Profile | None = None + transaction: CodexAuthTransaction | None = None + quota_snapshot = snapshot_quota_cache(registry, target.id) + try: + argv = login_argv( + registry, + stage, + browser_login=browser_login, + access_token=access_token, + ) + completed = subprocess.run( + argv, + env=provider_environment(stage), + check=False, + ) + if completed.returncode != 0: + raise ValueError( + f"provider login failed for {target.id}; target credentials were unchanged" + ) + refresh_provider_identity_anchors(registry, target.provider) + staged_result, _ = _probe_enrollment_candidate(registry, stage) + if not staged_result["credential_verified"]: + raise ValueError( + f"staged remote verification failed for {target.id}: " + f"{_verification_reason(staged_result)}; target credentials were unchanged" + ) + promotion = prepare_codex_promotion(registry, target, stage) + promotion_result, _ = _probe_enrollment_candidate(registry, promotion) + if not promotion_result["credential_verified"]: + raise ValueError( + f"promotion verification failed for {target.id}: " + f"{_verification_reason(promotion_result)}; target credentials were unchanged" + ) + transaction = activate_codex_promotion( + registry, + target, + promotion, + quota_snapshot, + ) + try: + target_result, target_quota = _probe_enrollment_candidate(registry, target) + if not target_result["credential_verified"]: + raise ValueError( + f"post-promotion verification failed for {target.id}: " + f"{_verification_reason(target_result)}" + ) + store_quota(registry, target, target_quota) + finalized = _evaluate_remote_profile( + registry, + target, + str(auth_probe(registry, target)["status"]), + read_quota(registry, target.id), + require_complete_worker_set=False, + ) + if not finalized["credential_verified"]: + raise ValueError( + f"stored verification failed for {target.id}: {_verification_reason(finalized)}" + ) + except BaseException: + rollback_codex_promotion(registry, target, transaction) + transaction = None + raise + finalize_codex_promotion(registry, target, transaction) + transaction = None + return finalized + finally: + if transaction is not None: + rollback_codex_promotion(registry, target, transaction) + else: + # Activation can crash/fail after the journal is durable but before + # returning its in-memory handle. Recover that case before deleting + # any staging or promotion artifact. + recover_pending_codex_transaction(registry, target) + if promotion is not None and promotion.home.exists(): + discard_codex_promotion(promotion, target) + if stage.home.exists(): + discard_codex_stage(stage, target) + + +def _run_profile_enrollment( + registry: Registry, + profile: Profile, + args: argparse.Namespace, + config_path: Path, +) -> dict[str, Any]: + if args.browser_login and profile.provider != "codex": + raise ValueError("--browser-login applies only to Codex profiles") + if args.access_token and profile.provider != "codex": + raise ValueError("--access-token applies only to Codex profiles") + if args.browser_login and args.access_token: + raise ValueError("choose --browser-login or --access-token, not both") + if args.enable: + raise ValueError( + "enrollment and routing enable are separate phases; verify, then run " + f"`agent-fleet profile enable {profile.id}`" + ) + try: + enrollment_lock = provider_enrollment_lock( + registry.settings.state_dir, + profile.provider, + registry.settings.lock_stale_seconds, + timeout=0.1, + ) + with enrollment_lock: + recover_pending_codex_transactions(registry, profile.provider) + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + current = load_registry(config_path) + profile = current.require_profile(profile.id) + if profile.enabled: + raise ValueError( + f"disable {profile.id} and drain all {profile.provider} leases " + "before enrollment" + ) + if _provider_has_active_lease(current, profile.provider): + raise ValueError( + f"refusing {profile.provider} login while any same-provider " + "Fleet lease is active" + ) + registry = current + if args.access_token: + print( + "Agent Fleet login safety: Codex will read the access token directly " + "from non-interactive stdin; Agent Fleet does not read or log it.", + file=sys.stderr, + ) + else: + if profile.provider == "claude": + browser_behavior = "Claude normally opens browser login automatically." + elif args.browser_login: + browser_behavior = "Codex will use its browser callback login." + else: + browser_behavior = ( + "Codex device login prints a URL and code; it does not normally " + "open the browser automatically." + ) + print( + "Agent Fleet enrollment safety: login is not generally idempotent; " + "a raw Codex login revokes the selected home's existing refresh token " + "before OAuth, even when OAuth is cancelled. Agent Fleet uses an " + "isolated Codex staging home so retries cannot revoke the target. " + f"{browser_behavior} Use a fresh Guest/private window, close that " + "entire window after success, and never select provider Log out.", + file=sys.stderr, + ) + if profile.provider == "codex": + verified = _enroll_codex_profile( + registry, + profile, + browser_login=args.browser_login, + access_token=args.access_token, + ) + else: + provision_profile(registry, profile) + completed = subprocess.run( + login_argv(registry, profile), + env=provider_environment(profile), + check=False, + ) + if completed.returncode != 0: + raise ValueError( + f"provider login failed for {profile.id}; profile remains disabled" + ) + verified = _verify_remote_profile( + registry, + profile, + allow_keychain_prompt=False, + ) + if ( + not verified["credential_verified"] + and verified.get("remote_reason") == "keychain_access_required" + ): + verified["verification_pending"] = True + verified["next_step"] = ( + f"agent-fleet profile verify {profile.id} --allow-keychain-prompt" + ) + elif not verified["credential_verified"]: + raise ValueError( + f"remote verification failed for {profile.id}: " + f"{_verification_reason(verified)}; profile remains disabled" + ) + verified["enabled"] = False + return verified + except TimeoutError as exc: + raise ValueError( + f"another {profile.provider} enrollment or Fleet selection is in progress" + ) from exc + + +def _require_routeable_profile( + registry: Registry, + profile: Profile, + *, + ignore_reserve: bool = False, +) -> None: + refresh_provider_identity_anchors_if_due(registry, profile.provider) + refresh_due_quotas(registry, [profile]) + authentication = auth_status(registry, profile) + quota = read_quota(registry, profile.id) + conflict = identity_conflict(registry, profile, quota) + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication, + ignore_reserve=ignore_reserve, + ) + if conflict is not None: + raise ValueError(f"profile duplicates enabled provider identity {conflict}") + if not routeability["eligible"]: + raise ValueError(f"profile is not routeable: {routeability['reason']}") + + +def _run(args: argparse.Namespace) -> Any | None: + config_path = expand_path(args.config) + if args.command == "version": + return {"cli_version": __version__, "contract_version": 1} + if args.command == "contract": + return _contract() + if args.command == "init": + if config_path.exists() and not args.force: + raise ValueError(f"registry already exists: {config_path}") + registry = initial_registry(args.claude, args.codex) + save_registry(registry, config_path) + return { + "registry": str(config_path), + "profiles": [registry.profiles[key].public_dict() for key in sorted(registry.profiles)], + "enabled": 0, + } + + registry = load_registry(config_path) + if args.command == "profile": + if args.profile_command == "list": + return { + "profiles": [ + registry.profiles[key].public_dict() for key in sorted(registry.profiles) + ] + } + if args.profile_command == "add": + profile_id = validate_id(args.profile_id, "profile id") + with _provider_maintenance( + registry, + config_path, + {args.provider}, + ) as current: + if profile_id in current.profiles: + raise ValueError(f"profile already exists: {profile_id}") + default_pools = [f"{args.provider}-manual"] + if args.safety_policy == "worker": + default_pools.insert(0, f"{args.provider}-crew") + pools = args.pools or default_pools + pools = [validate_id(pool, "pool id") for pool in pools] + if args.safety_policy != "worker" and f"{args.provider}-crew" in pools: + raise ValueError( + f"{args.safety_policy} profiles cannot join the worker crew pool" + ) + if args.weight < 1 or args.max_concurrent < 1: + raise ValueError("weight and max-concurrent must be positive integers") + if not 0 <= args.reserve_percent <= 100: + raise ValueError("reserve-percent must be between 0 and 100") + home = expand_path( + args.home + or current.settings.share_dir / "accounts" / args.provider / profile_id + ) + profile = Profile( + id=profile_id, + provider=args.provider, + home=home, + pools=tuple(pools), + enabled=False, + weight=args.weight, + max_concurrent=args.max_concurrent, + reserve_percent=args.reserve_percent, + safety_policy=args.safety_policy, + ) + updated = _mutate( + current, + lambda item: with_profile(item, profile), + config_path, + ) + return updated.require_profile(profile_id).public_dict() + if args.profile_command == "remove": + profile = registry.require_profile(args.profile_id) + with _provider_maintenance( + registry, + config_path, + {profile.provider}, + require_disabled=(profile.id,), + ) as current: + if any(lease.get("profile") == profile.id for lease in active_leases(current)): + raise ValueError("cannot remove a profile with an active lease") + _mutate( + current, + lambda item: without_profile(item, profile.id), + config_path, + ) + return {"profile": profile.id, "removed": True, "home_deleted": False} + if args.profile_command in {"provision", "status", "auth-status", "verify"}: + profiles = _profiles_for(registry, args.profile_id, args.all) + if args.profile_command == "provision": + providers = {profile.provider for profile in profiles} + ids = tuple(profile.id for profile in profiles) + with _provider_maintenance( + registry, + config_path, + providers, + require_disabled=ids, + ) as current: + return { + "profiles": [ + provision_profile(current, current.require_profile(profile_id)) + for profile_id in ids + ] + } + if args.profile_command == "status": + providers = {profile.provider for profile in profiles} + with _provider_maintenance(registry, config_path, providers) as current: + return { + "profiles": [profile_status(current, profile.id) for profile in profiles] + } + if args.profile_command == "verify": + providers = {profile.provider for profile in profiles} + ids = tuple(profile.id for profile in profiles) + with _provider_maintenance( + registry, + config_path, + providers, + require_disabled=ids, + ) as current: + results = [ + _verify_remote_profile( + current, + current.require_profile(profile_id), + allow_keychain_prompt=args.allow_keychain_prompt, + ) + for profile_id in ids + ] + ready = all(result["credential_verified"] for result in results) + for result in results: + result["enabled"] = False + return { + "profiles": results, + "ready": ready, + "enabled_as_batch": None, + "next_step": ( + "enable each verified worker profile explicitly" + if ready + else "resolve verification failures while profiles remain disabled" + ), + } + return { + "profiles": [ + {"profile": profile.id, "status": auth_status(registry, profile)} + for profile in profiles + ] + } + if args.profile_command in {"cooldown", "cooldown-clear"}: + profile = registry.require_profile(args.profile_id) + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + if args.profile_command == "cooldown": + return set_cooldown( + registry, + profile.id, + seconds=args.seconds, + reason=args.reason, + ) + return clear_cooldown(registry, profile.id) + profile = registry.require_profile(args.profile_id) + if args.profile_command == "policy": + with _provider_maintenance( + registry, + config_path, + {profile.provider}, + require_disabled=(profile.id,), + ) as current: + updated = _mutate( + current, + lambda item: set_profile_safety_policy( + item, + profile.id, + args.safety_policy, + ), + config_path, + ) + return updated.require_profile(profile.id).public_dict() + if args.profile_command in {"enable", "disable"}: + enabled = args.profile_command == "enable" + with _provider_maintenance( + registry, + config_path, + {profile.provider}, + require_disabled=(profile.id,) if enabled else (), + ) as current: + profile = current.require_profile(profile.id) + if enabled: + if not profile_is_provisioned(profile): + raise ValueError("provision the profile before enabling it") + verified = _verify_remote_profile( + current, + profile, + allow_keychain_prompt=False, + require_complete_worker_set=True, + ) + if not verified["credential_verified"]: + raise ValueError( + f"profile is not remotely verified: {_verification_reason(verified)}" + ) + updated = _mutate( + current, + lambda item: set_profile_enabled(item, profile.id, enabled), + config_path, + ) + return updated.require_profile(profile.id).public_dict() + if args.profile_command in {"login", "enroll"}: + return _run_profile_enrollment(registry, profile, args, config_path) + + if args.command == "pool" and args.pool_command == "status": + validate_id(args.pool, "pool id") + providers = {args.provider} if args.provider else set(SUPPORTED_PROVIDERS) + with _provider_maintenance(registry, config_path, providers) as current: + return pool_status(current, pool=args.pool, provider=args.provider) + + if args.command == "choose": + if not args.dry_run: + raise ValueError("diagnostic choose requires --dry-run") + return select_and_acquire( + registry, + task=_task(args.task), + pool=validate_id(args.pool, "pool id"), + provider=args.provider, + profile_id=args.profile, + dry_run=True, + ) + + if args.command == "quota": + profiles = _profiles_for(registry, args.profile_id, args.all) + if args.quota_command == "refresh": + providers = {profile.provider for profile in profiles} + ids = tuple(profile.id for profile in profiles) + with _provider_maintenance(registry, config_path, providers) as current: + return { + "quota": [ + refresh_quota( + current, + current.require_profile(profile_id), + allow_keychain_prompt=args.allow_keychain_prompt, + ) + for profile_id in ids + ] + } + return {"quota": [read_quota(registry, profile.id) for profile in profiles]} + + if args.command == "lease": + if args.lease_command == "list": + return {"leases": active_leases(registry)} + if args.lease_command == "release": + task = _task(args.task) + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + return release_lease(registry, task, force=args.force) + if args.lease_command == "acquire": + task = _task(args.task) + pool = validate_id(args.pool or "explicit", "pool id") + return select_and_acquire( + registry, + task=task, + pool=pool, + profile_id=args.profile, + explicit_profile=True, + ) + if args.lease_command == "recover": + task = _task(args.task) + mapping = get_session(registry, task) + profile = registry.require_profile(str(mapping.get("profile"))) + if mapping.get("provider") != profile.provider: + raise ValueError("session mapping provider does not match its profile") + pool = mapping.get("pool") + if not isinstance(pool, str) or not pool: + raise ValueError("session mapping has no pool") + return select_and_acquire( + registry, + task=task, + pool=validate_id(pool, "pool id"), + provider=profile.provider, + profile_id=profile.id, + explicit_profile=True, + ignore_reserve=True, + recovery_reservation=True, + ) + task = _task(args.task) + return select_and_acquire( + registry, + task=task, + pool=validate_id(args.pool, "pool id"), + provider=args.provider, + profile_id=args.profile, + ) + + if args.command == "exec": + task = _task(args.task) if args.task else None + if task is None: + raise ValueError("worker exec requires --task so its provider lease is tracked") + pool = args.pool or ("explicit" if args.profile else None) + if pool is None: + raise ValueError("exec with --task requires --pool or --profile") + selected = select_and_acquire( + registry, + task=task, + pool=validate_id(pool, "pool id"), + provider=args.provider, + profile_id=args.profile, + bind_pid=os.getpid(), + explicit_profile=pool == "explicit", + ) + profile = registry.require_profile(str(selected["profile"])) + argv = provider_argv(registry, profile, _strip_separator(args.provider_args)) + os.execvpe(argv[0], argv, provider_environment(profile, task)) + + if args.command == "resume": + task = _task(args.task) if args.task else None + if task is None: + raise ValueError("worker resume requires --task so its provider lease is tracked") + mapping = get_session(registry, task) + profile = registry.require_profile(str(mapping.get("profile"))) + if args.profile and args.profile != profile.id: + raise ValueError("explicit profile does not match task session mapping") + session_id = mapping.get("session_id") + if args.session and args.session != session_id: + raise ValueError("explicit session does not match task session mapping") + pool = args.pool or str(mapping.get("pool")) + if not pool or pool == "None": + raise ValueError("session mapping has no pool; pass --pool") + if not isinstance(session_id, str): + raise ValueError("session mapping has no provider session id") + validate_id(session_id, "session id") + select_and_acquire( + registry, + task=task, + pool=validate_id(str(pool), "pool id"), + profile_id=profile.id, + bind_pid=os.getpid(), + explicit_profile=True, + ignore_reserve=True, + ) + argv = resume_argv( + registry, + profile, + session_id, + _strip_separator(args.provider_args), + ) + os.execvpe(argv[0], argv, provider_environment(profile, task)) + + if args.command == "session": + task = _task(args.task) + if args.session_command == "status": + return get_session(registry, task) + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + return remove_session(registry, task) + + if args.command == "hook" and args.hook_command == "session-start": + record_session_from_hook(registry, read_hook_payload()) + return None + if args.command == "doctor": + return run_doctor( + registry, + config_path, + workspace=expand_path(args.workspace) if args.workspace else None, + ) + if args.command == "status": + with _provider_maintenance( + registry, + config_path, + set(SUPPORTED_PROVIDERS), + ) as current: + return { + "schema": 1, + "profiles": [ + profile_status(current, profile_id) for profile_id in sorted(current.profiles) + ], + "leases": active_leases(current), + } + if args.command == "audit": + return {"schema": 1, "events": read_audit(registry, limit=args.limit)} + raise ValueError("unsupported command") + + +def main(argv: list[str] | None = None) -> int: + parser = _parser() + args = parser.parse_args(argv) + try: + payload = _run(args) + if payload is not None: + emit(payload, args.format) + return 0 + except (ValueError, TimeoutError, OSError) as exc: + if args.format == "json": + emit({"error": str(exc), "ok": False}, "json") + else: + print(f"agent-fleet: {exc}", file=sys.stderr) + return 2 diff --git a/tools/agent-fleet/src/agent_fleet/config.py b/tools/agent-fleet/src/agent_fleet/config.py new file mode 100644 index 00000000000..1b7e2d262d3 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/config.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +import json +import os +import tomllib +from dataclasses import replace +from pathlib import Path +from typing import Any + +from .models import ( + PROFILE_SAFETY_POLICIES, + SUPPORTED_PROVIDERS, + Profile, + ProviderConfig, + Registry, + Settings, +) +from .paths import default_config_path, default_share_dir, default_state_dir, expand_path +from .util import validate_id + + +def _integer(value: Any, name: str, *, minimum: int, maximum: int | None = None) -> int: + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError(f"{name} must be an integer") + if value < minimum or (maximum is not None and value > maximum): + suffix = f"..{maximum}" if maximum is not None else " or greater" + raise ValueError(f"{name} must be {minimum}{suffix}") + return value + + +def _profile_from_toml(profile_id: str, raw: dict[str, Any], share_dir: Path) -> Profile: + validate_id(profile_id, "profile id") + provider = str(raw.get("provider", "")) + if provider not in SUPPORTED_PROVIDERS: + raise ValueError(f"profile {profile_id}: unsupported provider: {provider}") + home_raw = raw.get("home") + if not isinstance(home_raw, str) or not home_raw: + home = share_dir / "accounts" / provider / profile_id + else: + home = expand_path(home_raw) + pools_raw = raw.get("pools", []) + if not isinstance(pools_raw, list) or not pools_raw: + raise ValueError(f"profile {profile_id}: pools must be a non-empty array") + pools = tuple(validate_id(str(pool), "pool id") for pool in pools_raw) + enabled = raw.get("enabled", False) + if not isinstance(enabled, bool): + raise ValueError(f"profile {profile_id}: enabled must be boolean") + safety_policy = raw.get("safety_policy", "worker") + if safety_policy not in PROFILE_SAFETY_POLICIES: + choices = ", ".join(PROFILE_SAFETY_POLICIES) + raise ValueError(f"profile {profile_id}: safety_policy must be one of {choices}") + if safety_policy != "worker": + pools = tuple(pool for pool in pools if pool != f"{provider}-crew") + if not pools: + pools = (f"{provider}-manual",) + if enabled: + raise ValueError(f"profile {profile_id}: {safety_policy} profiles cannot be enabled") + return Profile( + id=profile_id, + provider=provider, + home=home, + pools=pools, + enabled=enabled, + weight=_integer(raw.get("weight", 1), f"profile {profile_id}: weight", minimum=1), + max_concurrent=_integer( + raw.get("max_concurrent", 2), + f"profile {profile_id}: max_concurrent", + minimum=1, + ), + reserve_percent=_integer( + raw.get("reserve_percent", 15), + f"profile {profile_id}: reserve_percent", + minimum=0, + maximum=100, + ), + safety_policy=str(safety_policy), + ) + + +def load_registry(path: Path | None = None) -> Registry: + config_path = path or default_config_path() + try: + raw = tomllib.loads(config_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise ValueError(f"registry not found: {config_path}; run `agent-fleet init`") from exc + except tomllib.TOMLDecodeError as exc: + raise ValueError(f"invalid registry TOML at {config_path}: {exc}") from exc + + version = _integer(raw.get("version", 0), "version", minimum=1) + if version != 1: + raise ValueError(f"unsupported registry version: {version}") + settings_raw = raw.get("settings", {}) + if not isinstance(settings_raw, dict): + raise ValueError("settings must be a TOML table") + state_dir = expand_path(settings_raw.get("state_dir", default_state_dir())) + share_dir = expand_path(settings_raw.get("share_dir", default_share_dir())) + quota_binary = expand_path( + settings_raw.get( + "quota_binary", + "~/.local/libexec/agent-fleet/quota-axi/current/bin/quota-axi", + ) + ) + settings = Settings( + state_dir=state_dir, + share_dir=share_dir, + quota_binary=quota_binary, + quota_stale_seconds=_integer( + settings_raw.get("quota_stale_seconds", 300), + "quota_stale_seconds", + minimum=0, + ), + quota_verification_grace_seconds=_integer( + settings_raw.get("quota_verification_grace_seconds", 86400), + "quota_verification_grace_seconds", + minimum=0, + ), + lease_grace_seconds=_integer( + settings_raw.get("lease_grace_seconds", 30), + "lease_grace_seconds", + minimum=0, + ), + active_lease_penalty=_integer( + settings_raw.get("active_lease_penalty", 8), + "active_lease_penalty", + minimum=0, + ), + lock_stale_seconds=_integer( + settings_raw.get("lock_stale_seconds", 30), + "lock_stale_seconds", + minimum=1, + ), + ) + + providers_raw = raw.get("providers", {}) + if not isinstance(providers_raw, dict): + raise ValueError("providers must be a TOML table") + providers: dict[str, ProviderConfig] = {} + for provider in SUPPORTED_PROVIDERS: + item = providers_raw.get(provider, {}) + if not isinstance(item, dict): + raise ValueError(f"providers.{provider} must be a TOML table") + binary_raw = item.get("binary") + if not isinstance(binary_raw, str) or not binary_raw: + raise ValueError(f"providers.{provider}.binary is required") + base_home_raw = item.get("base_home") + hooks_source_raw = item.get("hooks_source") + desktop_identity_file_raw = item.get( + "desktop_identity_file", + "~/Library/Application Support/Claude/config.json" if provider == "claude" else None, + ) + shared_raw = item.get("shared_entries", []) + if base_home_raw is not None and not isinstance(base_home_raw, str): + raise ValueError(f"providers.{provider}.base_home must be a path string") + if hooks_source_raw is not None and not isinstance(hooks_source_raw, str): + raise ValueError(f"providers.{provider}.hooks_source must be a path string") + if desktop_identity_file_raw is not None and not isinstance( + desktop_identity_file_raw, (str, bool) + ): + raise ValueError( + f"providers.{provider}.desktop_identity_file must be a path string or false" + ) + if desktop_identity_file_raw is True: + raise ValueError( + f"providers.{provider}.desktop_identity_file=true is ambiguous; " + "provide a path or false" + ) + if desktop_identity_file_raw == "": + raise ValueError( + f"providers.{provider}.desktop_identity_file cannot be empty; use false to opt out" + ) + if not isinstance(shared_raw, list) or not all( + isinstance(entry, str) and "/" not in entry and entry not in {"", ".", ".."} + for entry in shared_raw + ): + raise ValueError(f"providers.{provider}.shared_entries must contain simple file names") + providers[provider] = ProviderConfig( + provider, + expand_path(binary_raw), + expand_path(base_home_raw) if base_home_raw else None, + expand_path(hooks_source_raw) if hooks_source_raw else None, + tuple(shared_raw), + ( + expand_path(desktop_identity_file_raw) + if isinstance(desktop_identity_file_raw, str) and desktop_identity_file_raw + else None + ), + ) + + profiles_raw = raw.get("profiles", {}) + if not isinstance(profiles_raw, dict): + raise ValueError("profiles must be a TOML table") + profiles = { + profile_id: _profile_from_toml(profile_id, item, share_dir) + for profile_id, item in profiles_raw.items() + if isinstance(item, dict) + } + if len(profiles) != len(profiles_raw): + raise ValueError("each profiles entry must be a TOML table") + registry = Registry(version, settings, providers, profiles) + _validate_profile_invariants(registry) + return registry + + +def initial_registry(claude_count: int, codex_count: int) -> Registry: + if claude_count < 0 or codex_count < 0 or claude_count + codex_count == 0: + raise ValueError("at least one profile is required") + state_dir = default_state_dir() + share_dir = default_share_dir() + providers = { + "claude": ProviderConfig( + "claude", + expand_path(os.environ.get("AGENT_FLEET_CLAUDE_BIN", "~/.local/bin/claude")), + expand_path("~/.claude"), + expand_path("~/.claude/settings.json"), + ("CLAUDE.md", "skills", "plugins"), + expand_path("~/Library/Application Support/Claude/config.json"), + ), + "codex": ProviderConfig( + "codex", + expand_path( + os.environ.get( + "AGENT_FLEET_CODEX_BIN", + "~/.local/libexec/agent-fleet/runtime/codex", + ) + ), + expand_path("~/.codex"), + expand_path("~/.codex/hooks.json"), + ("AGENTS.md", "skills", "plugins", "rules"), + ), + } + profiles: dict[str, Profile] = {} + for provider, count in (("claude", claude_count), ("codex", codex_count)): + for index in range(1, count + 1): + profile_id = f"{provider}-{index}" + safety_policy = "worker" + pools = [f"{provider}-crew", f"{provider}-manual"] + if provider == "claude": + pools.append("claude-captain") + profiles[profile_id] = Profile( + id=profile_id, + provider=provider, + home=share_dir / "accounts" / provider / str(index), + pools=tuple(pools), + enabled=False, + max_concurrent=2 if provider == "claude" else 3, + safety_policy=safety_policy, + ) + quota_binary = expand_path( + os.environ.get( + "AGENT_FLEET_QUOTA_BIN", + "~/.local/libexec/agent-fleet/quota-axi/current/bin/quota-axi", + ) + ) + registry = Registry(1, Settings(state_dir, share_dir, quota_binary), providers, profiles) + _validate_profile_invariants(registry) + return registry + + +def with_profile(registry: Registry, profile: Profile) -> Registry: + profiles = dict(registry.profiles) + profiles[profile.id] = profile + updated = replace(registry, profiles=profiles) + _validate_profile_invariants(updated) + return updated + + +def without_profile(registry: Registry, profile_id: str) -> Registry: + registry.require_profile(profile_id) + profiles = dict(registry.profiles) + del profiles[profile_id] + return replace(registry, profiles=profiles) + + +def _paths_overlap(first: Path, second: Path) -> bool: + return first == second or first in second.parents or second in first.parents + + +def _validate_profile_invariants(registry: Registry) -> None: + profiles = list(registry.profiles.values()) + for profile in profiles: + if profile.safety_policy not in PROFILE_SAFETY_POLICIES: + raise ValueError( + f"profile {profile.id}: unsupported safety policy {profile.safety_policy}" + ) + if profile.safety_policy != "worker": + if profile.enabled: + raise ValueError( + f"profile {profile.id}: {profile.safety_policy} profiles cannot be enabled" + ) + if f"{profile.provider}-crew" in profile.pools: + raise ValueError( + f"profile {profile.id}: {profile.safety_policy} profiles " + "cannot join a worker crew pool" + ) + home = profile.home.resolve() + state_dir = registry.settings.state_dir.resolve() + share_dir = registry.settings.share_dir.resolve() + if _paths_overlap(home, state_dir): + raise ValueError(f"profile {profile.id}: home overlaps Agent Fleet state directory") + if home == share_dir or home in share_dir.parents: + raise ValueError(f"profile {profile.id}: home is too broad for Agent Fleet share data") + provider = registry.require_provider(profile.provider) + if provider.base_home is not None and _paths_overlap( + home, + provider.base_home.resolve(), + ): + raise ValueError(f"profile {profile.id}: home overlaps the provider base/Desktop home") + if provider.desktop_identity_file is not None and home in ( + provider.desktop_identity_file.resolve().parents + ): + raise ValueError( + f"profile {profile.id}: home contains the provider Desktop identity file" + ) + for index, profile in enumerate(profiles): + for other in profiles[index + 1 :]: + if _paths_overlap(profile.home.resolve(), other.home.resolve()): + raise ValueError(f"profile homes must not overlap: {profile.id}, {other.id}") + + +def set_profile_enabled(registry: Registry, profile_id: str, enabled: bool) -> Registry: + profile = registry.require_profile(profile_id) + if enabled and profile.safety_policy != "worker": + raise ValueError( + f"{profile.safety_policy} profile {profile.id} cannot be enabled for routing" + ) + return with_profile(registry, replace(profile, enabled=enabled)) + + +def set_profile_safety_policy( + registry: Registry, + profile_id: str, + safety_policy: str, +) -> Registry: + if safety_policy not in PROFILE_SAFETY_POLICIES: + choices = ", ".join(PROFILE_SAFETY_POLICIES) + raise ValueError(f"safety policy must be one of {choices}") + profile = registry.require_profile(profile_id) + pools = profile.pools + if safety_policy != "worker": + pools = tuple(pool for pool in pools if pool != f"{profile.provider}-crew") + if not pools: + pools = (f"{profile.provider}-manual",) + return with_profile( + registry, + replace( + profile, + enabled=False if safety_policy != "worker" else profile.enabled, + pools=pools, + safety_policy=safety_policy, + ), + ) + + +def _toml_string(value: str | Path) -> str: + return json.dumps(str(value), ensure_ascii=False) + + +def save_registry(registry: Registry, path: Path | None = None) -> Path: + _validate_profile_invariants(registry) + config_path = path or default_config_path() + config_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + config_path.parent.chmod(0o700) + lines = [ + "# Agent Fleet registry. Contains labels and policy only; never credentials.", + f"version = {registry.version}", + "", + "[settings]", + f"state_dir = {_toml_string(registry.settings.state_dir)}", + f"share_dir = {_toml_string(registry.settings.share_dir)}", + f"quota_binary = {_toml_string(registry.settings.quota_binary)}", + f"quota_stale_seconds = {registry.settings.quota_stale_seconds}", + f"quota_verification_grace_seconds = {registry.settings.quota_verification_grace_seconds}", + f"lease_grace_seconds = {registry.settings.lease_grace_seconds}", + f"active_lease_penalty = {registry.settings.active_lease_penalty}", + f"lock_stale_seconds = {registry.settings.lock_stale_seconds}", + ] + for provider in SUPPORTED_PROVIDERS: + provider_config = registry.require_provider(provider) + lines.extend( + [ + "", + f"[providers.{provider}]", + f"binary = {_toml_string(provider_config.binary)}", + *( + [f"base_home = {_toml_string(provider_config.base_home)}"] + if provider_config.base_home + else [] + ), + *( + [f"hooks_source = {_toml_string(provider_config.hooks_source)}"] + if provider_config.hooks_source + else [] + ), + *( + [ + "desktop_identity_file = " + f"{_toml_string(provider_config.desktop_identity_file)}" + ] + if provider_config.desktop_identity_file + else (["desktop_identity_file = false"] if provider == "claude" else []) + ), + "shared_entries = [" + + ", ".join(_toml_string(entry) for entry in provider_config.shared_entries) + + "]", + ] + ) + for profile_id in sorted(registry.profiles): + profile = registry.profiles[profile_id] + pools = ", ".join(_toml_string(pool) for pool in profile.pools) + lines.extend( + [ + "", + f"[profiles.{_toml_string(profile_id)}]", + f"provider = {_toml_string(profile.provider)}", + f"home = {_toml_string(profile.home)}", + f"pools = [{pools}]", + f"enabled = {'true' if profile.enabled else 'false'}", + f"weight = {profile.weight}", + f"max_concurrent = {profile.max_concurrent}", + f"reserve_percent = {profile.reserve_percent}", + f"safety_policy = {_toml_string(profile.safety_policy)}", + ] + ) + temp = config_path.with_name(f".{config_path.name}.{os.getpid()}.tmp") + temp.write_text("\n".join(lines) + "\n", encoding="utf-8") + temp.chmod(0o600) + os.replace(temp, config_path) + config_path.chmod(0o600) + return config_path diff --git a/tools/agent-fleet/src/agent_fleet/cooldowns.py b/tools/agent-fleet/src/agent_fleet/cooldowns.py new file mode 100644 index 00000000000..f6fedada9b6 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/cooldowns.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import json +import time +from pathlib import Path +from typing import Any + +from .models import Registry +from .util import atomic_write_json, utc_now, validate_id + + +def cooldown_path(registry: Registry, profile_id: str) -> Path: + return registry.settings.state_dir / "cooldowns" / f"{profile_id}.json" + + +def read_cooldown(registry: Registry, profile_id: str) -> dict[str, Any] | None: + path = cooldown_path(registry, profile_id) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + if not isinstance(value, dict) or value.get("profile") != profile_id: + return None + expires = value.get("expires_unix") + if not isinstance(expires, (int, float)) or float(expires) <= time.time(): + return None + return value + + +def set_cooldown( + registry: Registry, + profile_id: str, + *, + seconds: int, + reason: str, +) -> dict[str, Any]: + validate_id(profile_id, "profile id") + if seconds < 1 or seconds > 86_400: + raise ValueError("cooldown seconds must be between 1 and 86400") + if not reason or len(reason) > 128 or any(ord(char) < 32 for char in reason): + raise ValueError("cooldown reason must be 1-128 printable characters") + payload = { + "schema": 1, + "profile": profile_id, + "reason": reason, + "created_at": utc_now(), + "expires_unix": time.time() + seconds, + } + atomic_write_json(cooldown_path(registry, profile_id), payload) + return payload + + +def clear_cooldown(registry: Registry, profile_id: str) -> dict[str, Any]: + path = cooldown_path(registry, profile_id) + existed = path.exists() + path.unlink(missing_ok=True) + return {"profile": profile_id, "cooldown_cleared": existed} diff --git a/tools/agent-fleet/src/agent_fleet/doctor.py b/tools/agent-fleet/src/agent_fleet/doctor.py new file mode 100644 index 00000000000..1e79cdf3f51 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/doctor.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import json +import os +import shutil +import stat +from pathlib import Path +from typing import Any + +from .models import Registry +from .providers import auth_status +from .provision import ( + profile_hook_health, + profile_is_provisioned, + profile_shared_assets_healthy, +) + + +def _mode(path: Path) -> str | None: + try: + return oct(stat.S_IMODE(path.stat().st_mode)) + except FileNotFoundError: + return None + + +def _workspace_hook_events(path: Path) -> set[str]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return set() + hooks = payload.get("hooks", {}) if isinstance(payload, dict) else {} + return set(hooks) if isinstance(hooks, dict) else set() + + +def run_doctor( + registry: Registry, + config_path: Path, + *, + workspace: Path | None = None, +) -> dict[str, Any]: + checks: list[dict[str, Any]] = [] + + def add(name: str, ok: bool, detail: str, *, required: bool = True) -> None: + checks.append({"name": name, "ok": ok, "required": required, "detail": detail}) + + add( + "registry-permissions", + _mode(config_path) == "0o600", + f"{config_path} mode={_mode(config_path)} expected=0o600", + ) + add("toon", shutil.which("toon") is not None, "TOON encoder on PATH") + add( + "quota-axi", + registry.settings.quota_binary.is_file() + and os.access(registry.settings.quota_binary, os.X_OK), + f"pinned quota reader: {registry.settings.quota_binary}", + ) + for provider, config in registry.providers.items(): + binary_ok = config.binary.is_file() and os.access(config.binary, os.X_OK) + add( + f"binary:{provider}", + binary_ok, + f"{config.binary}", + ) + for profile in sorted(registry.profiles.values(), key=lambda item: item.id): + provisioned = profile_is_provisioned(profile) + add( + f"profile:{profile.id}:provisioned", + provisioned or not profile.enabled, + "provisioned" if provisioned else "not provisioned (allowed while disabled)", + ) + status = auth_status(registry, profile) if provisioned else "not-provisioned" + add( + f"profile:{profile.id}:auth", + status == "authenticated" or not profile.enabled, + status, + ) + if provisioned: + home_mode = _mode(profile.home) + add( + f"profile:{profile.id}:home-permissions", + home_mode == "0o700", + f"mode={home_mode} expected=0o700", + ) + hook_health = profile_hook_health(registry, profile) + for name, healthy in hook_health.items(): + add( + f"profile:{profile.id}:{name.replace('_', '-')}", + healthy, + "present" if healthy else "missing", + ) + shared_healthy = profile_shared_assets_healthy(registry, profile) + add( + f"profile:{profile.id}:shared-workflow-assets", + shared_healthy, + "healthy" if shared_healthy else "missing or redirected link", + ) + if workspace is not None: + workspace = workspace.resolve() + claude_events = _workspace_hook_events(workspace / ".claude" / "settings.json") + codex_events = _workspace_hook_events(workspace / ".codex" / "hooks.json") + for provider, events in (("claude", claude_events), ("codex", codex_events)): + required = {"PreToolUse", "Stop"} + add( + f"workspace:{provider}:supervision-hooks", + required.issubset(events), + f"{workspace}: events={','.join(sorted(events)) or 'none'}", + ) + required_failures = [check for check in checks if check["required"] and not check["ok"]] + return { + "healthy": not required_failures, + "profiles": len(registry.profiles), + "enabled_profiles": sum(profile.enabled for profile in registry.profiles.values()), + "workspace": str(workspace) if workspace else None, + "checks": checks, + } diff --git a/tools/agent-fleet/src/agent_fleet/enrollment.py b/tools/agent-fleet/src/agent_fleet/enrollment.py new file mode 100644 index 00000000000..49262ef6231 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/enrollment.py @@ -0,0 +1,488 @@ +from __future__ import annotations + +import base64 +import json +import os +import shutil +import stat +import tempfile +import uuid +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +from .models import Profile, Registry +from .paths import ensure_private_dir +from .provision import provision_profile +from .quota import QuotaCacheSnapshot, restore_quota_cache +from .util import atomic_write_json + +CODEX_AUTH_FILE = "auth.json" + + +@dataclass(frozen=True) +class CodexAuthTransaction: + target_auth: Path + backup_auth: Path | None + installed_stat: tuple[int, int, int, int] + journal_path: Path + + +def _regular_file_stat(path: Path, label: str) -> os.stat_result: + try: + current = path.lstat() + except FileNotFoundError as exc: + raise ValueError(f"{label} is missing: {path}") from exc + if not stat.S_ISREG(current.st_mode): + raise ValueError(f"{label} must be a regular file: {path}") + if current.st_uid != os.getuid(): + raise ValueError(f"{label} must be owned by the current user: {path}") + return current + + +def _private_directory(path: Path, label: str) -> None: + try: + current = path.lstat() + except FileNotFoundError as exc: + raise ValueError(f"{label} is missing: {path}") from exc + if not stat.S_ISDIR(current.st_mode): + raise ValueError(f"{label} must be a regular directory: {path}") + if current.st_uid != os.getuid(): + raise ValueError(f"{label} must be owned by the current user: {path}") + if stat.S_IMODE(current.st_mode) & 0o077: + raise ValueError(f"{label} must not grant group/world access: {path}") + + +def _stat_identity(current: os.stat_result) -> tuple[int, int, int, int]: + return (current.st_dev, current.st_ino, current.st_size, current.st_mtime_ns) + + +def _encoded_stat(current: tuple[int, int, int, int] | None) -> list[int] | None: + return list(current) if current is not None else None + + +def _decoded_stat(value: Any, label: str) -> tuple[int, int, int, int] | None: + if value is None: + return None + if ( + not isinstance(value, list) + or len(value) != 4 + or not all(isinstance(item, int) and not isinstance(item, bool) for item in value) + ): + raise ValueError(f"invalid {label} in Codex auth transaction journal") + return tuple(value) # type: ignore[return-value] + + +def _journal_path(registry: Registry, target: Profile) -> Path: + return registry.settings.state_dir / "transactions" / f"codex-auth-{target.id}.json" + + +def _snapshot_payload(snapshot: QuotaCacheSnapshot) -> dict[str, Any]: + return { + "existed": snapshot.existed, + "payload": base64.b64encode(snapshot.payload).decode("ascii"), + "mode": snapshot.mode, + } + + +def _snapshot_from_payload(value: Any) -> QuotaCacheSnapshot: + if not isinstance(value, dict) or not isinstance(value.get("existed"), bool): + raise ValueError("invalid quota snapshot in Codex auth transaction journal") + encoded = value.get("payload", "") + mode = value.get("mode", 0o600) + if not isinstance(encoded, str) or not isinstance(mode, int) or isinstance(mode, bool): + raise ValueError("invalid quota snapshot in Codex auth transaction journal") + if mode < 0 or mode > 0o777 or mode & 0o077: + raise ValueError("unsafe quota snapshot mode in Codex auth transaction journal") + try: + payload = base64.b64decode(encoded, validate=True) + except ValueError as exc: + raise ValueError("invalid quota snapshot in Codex auth transaction journal") from exc + return QuotaCacheSnapshot(bool(value["existed"]), payload, mode) + + +def _read_journal(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise ValueError(f"Codex auth transaction journal is missing: {path}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"Codex auth transaction journal is invalid: {path}") from exc + if not isinstance(value, dict) or value.get("schema") != 1: + raise ValueError(f"Codex auth transaction journal is invalid: {path}") + return value + + +def _validate_journal_paths( + registry: Registry, + target: Profile, + journal: dict[str, Any], +) -> tuple[Path, Path | None, Path]: + if journal.get("profile") != target.id or journal.get("provider") != "codex": + raise ValueError("Codex auth transaction journal targets another profile") + target_auth = Path(str(journal.get("target_auth", ""))) + if target_auth != target.home / CODEX_AUTH_FILE: + raise ValueError("Codex auth transaction journal target is outside the profile") + backup_raw = journal.get("backup_auth") + if backup_raw is not None and not isinstance(backup_raw, str): + raise ValueError("Codex auth transaction backup path is invalid") + backup = Path(backup_raw) if isinstance(backup_raw, str) else None + if backup is not None and ( + backup.parent != target.home or not backup.name.startswith(f".{CODEX_AUTH_FILE}.backup-") + ): + raise ValueError("Codex auth transaction backup is outside the profile") + temporary = Path(str(journal.get("temporary_auth", ""))) + if temporary.parent != target.home or not temporary.name.startswith(f".{CODEX_AUTH_FILE}.new-"): + raise ValueError("Codex auth transaction temporary file is outside the profile") + expected_journal = _journal_path(registry, target) + if Path(str(journal.get("journal_path", ""))) != expected_journal: + raise ValueError("Codex auth transaction journal path is inconsistent") + return target_auth, backup, temporary + + +def _open_no_follow(path: Path) -> int: + flags = os.O_RDONLY + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + return os.open(path, flags) + + +def _copy_regular_file(source: Path, destination: Path) -> None: + source_stat = _regular_file_stat(source, "Codex staged auth") + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + source_fd = _open_no_follow(source) + try: + if _stat_identity(os.fstat(source_fd)) != _stat_identity(source_stat): + raise ValueError("Codex staged auth changed while opening it") + destination_fd = os.open(destination, flags, 0o600) + try: + while True: + chunk = os.read(source_fd, 1024 * 1024) + if not chunk: + break + view = memoryview(chunk) + while view: + written = os.write(destination_fd, view) + view = view[written:] + os.fchmod(destination_fd, 0o600) + os.fsync(destination_fd) + finally: + os.close(destination_fd) + finally: + os.close(source_fd) + + +def create_codex_login_stage(profile: Profile) -> Profile: + ensure_private_dir(profile.home.parent) + _private_directory(profile.home.parent, "managed Codex accounts directory") + stage = Path( + tempfile.mkdtemp( + prefix=f".{profile.home.name}.login-", + dir=profile.home.parent, + ) + ) + stage.chmod(0o700) + config = stage / "config.toml" + config.write_text('cli_auth_credentials_store = "file"\n', encoding="utf-8") + config.chmod(0o600) + return replace(profile, home=stage, enabled=False) + + +def _remove_private_tree(path: Path) -> None: + if path.is_symlink(): + path.unlink() + elif path.exists(): + shutil.rmtree(path) + + +def discard_codex_stage(stage: Profile, target: Profile) -> None: + expected_prefix = f".{target.home.name}.login-" + if stage.home.parent != target.home.parent or not stage.home.name.startswith(expected_prefix): + raise ValueError(f"refusing to discard unrecognized Codex stage: {stage.home}") + _remove_private_tree(stage.home) + + +def discard_codex_promotion(promotion: Profile, target: Profile) -> None: + expected_prefix = f".{target.home.name}.promote-" + if promotion.home.parent != target.home.parent or not promotion.home.name.startswith( + expected_prefix + ): + raise ValueError(f"refusing to discard unrecognized Codex promotion: {promotion.home}") + _remove_private_tree(promotion.home) + + +def prepare_codex_promotion( + registry: Registry, + target: Profile, + stage: Profile, +) -> Profile: + staged_auth = stage.home / CODEX_AUTH_FILE + _private_directory(stage.home, "Codex staging home") + _regular_file_stat(staged_auth, "Codex staged auth") + if target.home.is_symlink(): + raise ValueError(f"managed Codex profile home cannot be a symlink: {target.home}") + provision_profile(registry, target) + _private_directory(target.home, "managed Codex profile home") + promotion = Path( + tempfile.mkdtemp( + prefix=f".{target.home.name}.promote-", + dir=target.home.parent, + ) + ) + promotion.chmod(0o700) + try: + shutil.copytree( + target.home, + promotion, + symlinks=True, + copy_function=shutil.copy2, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns(".agent-fleet-quota-cache"), + ) + destination = promotion / CODEX_AUTH_FILE + if destination.exists() or destination.is_symlink(): + if destination.is_symlink() or not destination.is_file(): + raise ValueError("managed Codex auth.json must be a regular file") + destination.unlink() + temporary = promotion / f".{CODEX_AUTH_FILE}.new-{uuid.uuid4().hex}" + _copy_regular_file(staged_auth, temporary) + os.replace(temporary, destination) + destination.chmod(0o600) + _fsync_directory(promotion) + promotion_profile = replace(target, home=promotion, enabled=False) + provision_profile(registry, promotion_profile) + return promotion_profile + except BaseException: + _remove_private_tree(promotion) + raise + + +def _fsync_directory(path: Path) -> None: + descriptor = os.open(path, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _validate_promotion(target: Profile, promotion: Profile) -> Path: + expected_prefix = f".{target.home.name}.promote-" + if promotion.home.parent != target.home.parent or not promotion.home.name.startswith( + expected_prefix + ): + raise ValueError(f"refusing to activate unrecognized Codex promotion: {promotion.home}") + _private_directory(target.home, "managed Codex profile home") + _private_directory(promotion.home, "Codex promotion home") + source = promotion.home / CODEX_AUTH_FILE + _regular_file_stat(source, "Codex promotion auth") + return source + + +def activate_codex_promotion( + registry: Registry, + target: Profile, + promotion: Profile, + quota_snapshot: QuotaCacheSnapshot, +) -> CodexAuthTransaction: + source = _validate_promotion(target, promotion) + target_auth = target.home / CODEX_AUTH_FILE + original_stat: os.stat_result | None = None + backup: Path | None = None + if target_auth.exists() or target_auth.is_symlink(): + original_stat = _regular_file_stat(target_auth, "managed Codex auth") + if stat.S_IMODE(original_stat.st_mode) & 0o077: + raise ValueError("managed Codex auth must not grant group/world access") + backup = target.home / f".{CODEX_AUTH_FILE}.backup-{uuid.uuid4().hex}" + _copy_regular_file(target_auth, backup) + temporary = target.home / f".{CODEX_AUTH_FILE}.new-{uuid.uuid4().hex}" + journal_path = _journal_path(registry, target) + if journal_path.exists() or journal_path.is_symlink(): + raise ValueError(f"unfinished Codex auth transaction requires recovery: {journal_path}") + try: + _copy_regular_file(source, temporary) + if original_stat is not None: + if _stat_identity( + _regular_file_stat(target_auth, "managed Codex auth") + ) != _stat_identity(original_stat): + raise ValueError("managed Codex auth changed during promotion") + elif target_auth.exists() or target_auth.is_symlink(): + raise ValueError("managed Codex auth appeared during promotion") + installed_stat = _stat_identity(_regular_file_stat(temporary, "prepared Codex auth")) + journal = { + "schema": 1, + "profile": target.id, + "provider": "codex", + "phase": "prepared", + "target_auth": str(target_auth), + "backup_auth": str(backup) if backup is not None else None, + "temporary_auth": str(temporary), + "original_stat": _encoded_stat( + _stat_identity(original_stat) if original_stat is not None else None + ), + "installed_stat": _encoded_stat(installed_stat), + "quota_snapshot": _snapshot_payload(quota_snapshot), + "journal_path": str(journal_path), + } + # This durable intent record must reach disk before auth.json can change. + # A process crash after the replace is therefore always recoverable. + atomic_write_json(journal_path, journal) + _fsync_directory(journal_path.parent) + os.replace(temporary, target_auth) + target_auth.chmod(0o600) + installed = _regular_file_stat(target_auth, "promoted Codex auth") + if _stat_identity(installed) != installed_stat: + raise ValueError("promoted Codex auth identity changed during replace") + _fsync_directory(target.home) + return CodexAuthTransaction( + target_auth=target_auth, + backup_auth=backup, + installed_stat=_stat_identity(installed), + journal_path=journal_path, + ) + except BaseException: + temporary.unlink(missing_ok=True) + if backup is not None and not journal_path.exists(): + backup.unlink(missing_ok=True) + raise + + +def _validate_transaction( + registry: Registry, + target: Profile, + transaction: CodexAuthTransaction, +) -> dict[str, Any]: + expected = target.home / CODEX_AUTH_FILE + if transaction.target_auth != expected: + raise ValueError("Codex auth transaction does not belong to the target profile") + if transaction.journal_path != _journal_path(registry, target): + raise ValueError("Codex auth transaction journal does not belong to the target profile") + journal = _read_journal(transaction.journal_path) + target_auth, backup, _ = _validate_journal_paths(registry, target, journal) + if target_auth != transaction.target_auth or backup != transaction.backup_auth: + raise ValueError("Codex auth transaction disagrees with its durable journal") + if _decoded_stat(journal.get("installed_stat"), "installed stat") != ( + transaction.installed_stat + ): + raise ValueError("Codex auth transaction stat disagrees with its durable journal") + current = _regular_file_stat(expected, "promoted Codex auth") + if _stat_identity(current) != transaction.installed_stat: + raise ValueError("promoted Codex auth changed before transaction completion") + return journal + + +def rollback_codex_promotion( + registry: Registry, + target: Profile, + transaction: CodexAuthTransaction, +) -> None: + _private_directory(target.home, "managed Codex profile home") + _validate_transaction(registry, target, transaction) + backup = transaction.backup_auth + if backup is None: + transaction.target_auth.unlink() + else: + expected_prefix = f".{CODEX_AUTH_FILE}.backup-" + if backup.parent != target.home or not backup.name.startswith(expected_prefix): + raise ValueError("Codex auth backup is outside the target profile") + _regular_file_stat(backup, "Codex auth backup") + os.replace(backup, transaction.target_auth) + transaction.target_auth.chmod(0o600) + _fsync_directory(target.home) + journal = _read_journal(transaction.journal_path) + restore_quota_cache( + registry, + target.id, + _snapshot_from_payload(journal.get("quota_snapshot")), + ) + transaction.journal_path.unlink() + _fsync_directory(transaction.journal_path.parent) + + +def finalize_codex_promotion( + registry: Registry, + target: Profile, + transaction: CodexAuthTransaction, +) -> None: + _private_directory(target.home, "managed Codex profile home") + journal = _validate_transaction(registry, target, transaction) + journal["phase"] = "committed" + atomic_write_json(transaction.journal_path, journal) + _fsync_directory(transaction.journal_path.parent) + backup = transaction.backup_auth + if backup is not None: + expected_prefix = f".{CODEX_AUTH_FILE}.backup-" + if backup.parent != target.home or not backup.name.startswith(expected_prefix): + raise ValueError("Codex auth backup is outside the target profile") + _regular_file_stat(backup, "Codex auth backup") + backup.unlink() + _fsync_directory(target.home) + transaction.journal_path.unlink() + _fsync_directory(transaction.journal_path.parent) + + +def recover_pending_codex_transaction(registry: Registry, target: Profile) -> bool: + """Recover one interrupted promotion while the caller owns the provider lock.""" + + if target.provider != "codex": + return False + journal_path = _journal_path(registry, target) + if not journal_path.exists(): + return False + journal = _read_journal(journal_path) + target_auth, backup, temporary = _validate_journal_paths(registry, target, journal) + installed_stat = _decoded_stat(journal.get("installed_stat"), "installed stat") + original_stat = _decoded_stat(journal.get("original_stat"), "original stat") + if installed_stat is None: + raise ValueError("Codex auth transaction has no installed stat") + phase = journal.get("phase") + if phase not in {"prepared", "committed"}: + raise ValueError("Codex auth transaction has an invalid phase") + _private_directory(target.home, "managed Codex profile home") + + current_stat: tuple[int, int, int, int] | None + if target_auth.exists() or target_auth.is_symlink(): + current_stat = _stat_identity(_regular_file_stat(target_auth, "managed Codex auth")) + else: + current_stat = None + + if phase == "committed": + if current_stat != installed_stat: + raise ValueError("committed Codex auth changed before crash recovery") + else: + if current_stat == installed_stat: + if backup is None: + target_auth.unlink() + else: + _regular_file_stat(backup, "Codex auth backup") + os.replace(backup, target_auth) + target_auth.chmod(0o600) + _fsync_directory(target.home) + elif current_stat != original_stat: + raise ValueError("Codex auth changed outside the interrupted transaction") + # Restore normalized quota evidence before deleting any recovery artifact. + restore_quota_cache( + registry, + target.id, + _snapshot_from_payload(journal.get("quota_snapshot")), + ) + + # Cleanup happens only after rollback+quota restore or a durable commit. + temporary.unlink(missing_ok=True) + if backup is not None: + backup.unlink(missing_ok=True) + journal_path.unlink() + _fsync_directory(target.home) + _fsync_directory(journal_path.parent) + return True + + +def recover_pending_codex_transactions(registry: Registry, provider: str) -> list[str]: + if provider != "codex": + return [] + recovered: list[str] = [] + for profile in registry.profiles.values(): + if profile.provider == provider and recover_pending_codex_transaction(registry, profile): + recovered.append(profile.id) + return recovered diff --git a/tools/agent-fleet/src/agent_fleet/identity.py b/tools/agent-fleet/src/agent_fleet/identity.py new file mode 100644 index 00000000000..3f3fe8891bc --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/identity.py @@ -0,0 +1,248 @@ +from __future__ import annotations + +import json +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from .models import Profile, Registry +from .providers import identity_fingerprint +from .quota import probe_quota, read_quota +from .util import atomic_write_json, utc_now + + +def _anchor_path(registry: Registry, provider: str, kind: str) -> Path: + return registry.settings.state_dir / "identity-anchors" / f"{provider}-{kind}.json" + + +def _age_seconds(value: Any) -> int | None: + if not isinstance(value, str): + return None + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=UTC) + return max(0, int((datetime.now(UTC) - parsed.astimezone(UTC)).total_seconds())) + + +def _read_anchor(registry: Registry, provider: str, kind: str) -> dict[str, Any]: + try: + payload = json.loads(_anchor_path(registry, provider, kind).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {"status": "unavailable", "reason": "identity_anchor_missing"} + return ( + payload + if isinstance(payload, dict) + else { + "status": "unavailable", + "reason": "identity_anchor_invalid", + } + ) + + +def _quota_identity_is_verified(quota: dict[str, Any]) -> bool: + fingerprint = quota.get("identity_fingerprint") + return ( + quota.get("status") == "fresh" + and quota.get("verified_at") is not None + and quota.get("headroom_percent") is not None + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + ) + + +def _managed_identity_has_recent_proof(quota: dict[str, Any]) -> bool: + return _quota_identity_is_verified(quota) or quota.get("verified_recent") is True + + +def _anchor_is_fresh(registry: Registry, anchor: dict[str, Any]) -> bool: + age = _age_seconds(anchor.get("refreshed_at")) + return age is not None and age <= registry.settings.quota_stale_seconds + + +def _refresh_desktop_identity_anchor( + registry: Registry, + provider: str, + desktop_file: Path, +) -> dict[str, Any]: + try: + payload = json.loads(desktop_file.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + payload = None + identifier = payload.get("lastKnownAccountUuid") if isinstance(payload, dict) else None + if not desktop_file.exists(): + desktop_status = "absent" + elif isinstance(identifier, str) and identifier: + desktop_status = "present" + else: + desktop_status = "error" + result = { + "schema": 1, + "provider": provider, + "kind": "desktop", + "status": desktop_status, + "identity_fingerprint": ( + identity_fingerprint(provider, identifier) + if isinstance(identifier, str) and identifier + else None + ), + "refreshed_at": utc_now(), + } + atomic_write_json(_anchor_path(registry, provider, "desktop"), result) + return result + + +def refresh_provider_identity_anchors( + registry: Registry, + provider: str, + *, + allow_keychain_prompt: bool = False, + timeout: int = 30, +) -> dict[str, dict[str, Any]]: + provider_config = registry.require_provider(provider) + results: dict[str, dict[str, Any]] = {} + if provider_config.base_home is not None and provider_config.base_home.exists(): + base = Profile( + id=f"{provider}-base-anchor", + provider=provider, + home=provider_config.base_home, + pools=(f"{provider}-manual",), + enabled=False, + safety_policy="desktop_shared", + ) + try: + quota = probe_quota( + registry, + base, + timeout=timeout, + allow_keychain_prompt=allow_keychain_prompt, + ) + except (OSError, TimeoutError, ValueError) as exc: + base_result = { + "schema": 1, + "provider": provider, + "kind": "base", + "status": "error", + "reason": type(exc).__name__, + "identity_fingerprint": None, + "refreshed_at": utc_now(), + } + else: + status = str(quota.get("status", "unavailable")) + fingerprint = quota.get("identity_fingerprint") + reason = quota.get("reason") + if status == "fresh" and not _quota_identity_is_verified(quota): + status = "error" + reason = "base_identity_unavailable" + fingerprint = None + base_result = { + "schema": 1, + "provider": provider, + "kind": "base", + "status": status, + "reason": reason, + "identity_fingerprint": fingerprint, + "refreshed_at": utc_now(), + } + atomic_write_json(_anchor_path(registry, provider, "base"), base_result) + results["base"] = base_result + if provider == "claude" and provider_config.desktop_identity_file is not None: + desktop_result = _refresh_desktop_identity_anchor( + registry, + provider, + provider_config.desktop_identity_file, + ) + results["desktop"] = desktop_result + return results + + +def refresh_provider_identity_anchors_if_due( + registry: Registry, + provider: str, + *, + timeout: int = 4, +) -> None: + provider_config = registry.require_provider(provider) + due = False + if provider_config.base_home is not None and provider_config.base_home.exists(): + current = _read_anchor(registry, provider, "base") + age = _age_seconds(current.get("refreshed_at")) + due = age is None or age > registry.settings.quota_stale_seconds + if provider == "claude" and provider_config.desktop_identity_file is not None: + # Desktop can switch accounts between two route attempts. This local + # JSON read is cheap and must not inherit the base quota anchor's TTL. + _refresh_desktop_identity_anchor( + registry, + provider, + provider_config.desktop_identity_file, + ) + if due: + refresh_provider_identity_anchors(registry, provider, timeout=timeout) + + +def identity_conflict( + registry: Registry, + profile: Profile, + quota: dict[str, Any], + *, + require_complete_worker_set: bool = True, +) -> str | None: + fingerprint = quota.get("identity_fingerprint") + if not isinstance(fingerprint, str) or len(fingerprint) != 64: + return "identity_unavailable" + for other in registry.profiles.values(): + if other.id == profile.id or other.provider != profile.provider: + continue + other_quota = read_quota(registry, other.id) + other_fingerprint = other_quota.get("identity_fingerprint") + has_recent_proof = _managed_identity_has_recent_proof(other_quota) + if require_complete_worker_set and other.safety_policy == "worker" and not has_recent_proof: + return f"managed_identity_unverified:{other.id}" + if not has_recent_proof: + continue + if not isinstance(other_fingerprint, str) or len(other_fingerprint) != 64: + if require_complete_worker_set and other.safety_policy == "worker": + return f"managed_identity_unverified:{other.id}" + continue + if other_fingerprint == fingerprint: + return f"managed:{other.id}" + provider_config = registry.require_provider(profile.provider) + if provider_config.base_home is not None and profile.home == provider_config.base_home: + return "base_home_overlap" + if provider_config.base_home is not None and provider_config.base_home.exists(): + base = _read_anchor(registry, profile.provider, "base") + base_status = str(base.get("status", "unavailable")) + base_fingerprint = base.get("identity_fingerprint") + if ( + base_status != "fresh" + or not _anchor_is_fresh(registry, base) + or not isinstance(base_fingerprint, str) + or len(base_fingerprint) != 64 + ): + return f"base_identity_unverified:{base.get('reason') or base_status}" + if base_fingerprint == fingerprint: + return "base_identity" + if profile.provider == "claude" and provider_config.desktop_identity_file is not None: + desktop = _read_anchor(registry, profile.provider, "desktop") + desktop_status = desktop.get("status") + desktop_fingerprint = desktop.get("identity_fingerprint") + if not _anchor_is_fresh(registry, desktop): + return "desktop_identity_unverified" + # A configured Desktop anchor is a required safety boundary. Missing + # state is not evidence that Desktop is signed out: it can also mean + # the configured path moved or became unreadable. Operators who do not + # have/use Desktop must opt out explicitly with + # desktop_identity_file = false. + if desktop_status == "absent": + return "desktop_identity_unverified" + if ( + desktop_status != "present" + or not isinstance(desktop_fingerprint, str) + or len(desktop_fingerprint) != 64 + ): + return "desktop_identity_unverified" + if desktop_fingerprint == fingerprint: + return "desktop_identity" + return None diff --git a/tools/agent-fleet/src/agent_fleet/leases.py b/tools/agent-fleet/src/agent_fleet/leases.py new file mode 100644 index 00000000000..8d0b0148e06 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/leases.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import json +import socket +import time +from pathlib import Path +from typing import Any + +from .audit import append_audit +from .models import Registry +from .util import atomic_write_json, process_matches, process_start_token, task_key, utc_now + + +def lease_path(registry: Registry, task: str) -> Path: + return registry.settings.state_dir / "leases" / f"{task_key(task)}.json" + + +def _read(path: Path) -> dict[str, Any] | None: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + return value if isinstance(value, dict) else None + + +def lease_is_active(lease: dict[str, Any], *, grace_seconds: int) -> bool: + state = lease.get("state") + if state == "running": + pid = lease.get("pid") + process_start = lease.get("process_start") + return ( + isinstance(pid, int) + and isinstance(process_start, str) + and bool(process_start) + and process_matches(pid, process_start) + ) + if state == "reserved": + created = lease.get("created_unix") + return isinstance(created, (int, float)) and time.time() - float(created) <= grace_seconds + return False + + +def active_leases(registry: Registry, *, prune: bool = False) -> list[dict[str, Any]]: + directory = registry.settings.state_dir / "leases" + if not directory.exists(): + return [] + active: list[dict[str, Any]] = [] + for path in sorted(directory.glob("*.json")): + lease = _read(path) + if lease is not None and lease_is_active( + lease, grace_seconds=registry.settings.lease_grace_seconds + ): + active.append(lease) + elif prune: + path.unlink(missing_ok=True) + return active + + +def get_active_lease(registry: Registry, task: str) -> dict[str, Any] | None: + path = lease_path(registry, task) + lease = _read(path) + if lease is None: + return None + if lease.get("task") != task: + raise ValueError(f"task hash collision or corrupt lease: {path}") + if lease_is_active(lease, grace_seconds=registry.settings.lease_grace_seconds): + return lease + return None + + +def new_lease(task: str, profile_id: str, pool: str, *, pid: int | None) -> dict[str, Any]: + process_start = process_start_token(pid) if pid is not None else None + if pid is not None and process_start is None: + raise ValueError("cannot bind worker lease without a verified process start token") + payload: dict[str, Any] = { + "schema": 1, + "task": task, + "profile": profile_id, + "pool": pool, + "state": "reserved" if pid is None else "running", + "pid": pid, + "process_start": process_start, + "hostname": socket.gethostname(), + "created_at": utc_now(), + "created_unix": time.time(), + "bound_at": utc_now() if pid is not None else None, + } + return payload + + +def write_lease(registry: Registry, lease: dict[str, Any]) -> None: + atomic_write_json(lease_path(registry, str(lease["task"])), lease) + + +def bind_lease(registry: Registry, lease: dict[str, Any], pid: int) -> dict[str, Any]: + process_start = process_start_token(pid) + if process_start is None: + raise ValueError("cannot bind worker lease without a verified process start token") + bound = dict(lease) + bound.update( + { + "state": "running", + "pid": pid, + "process_start": process_start, + "hostname": socket.gethostname(), + "bound_at": utc_now(), + } + ) + write_lease(registry, bound) + return bound + + +def release_lease(registry: Registry, task: str, *, force: bool = False) -> dict[str, Any]: + path = lease_path(registry, task) + lease = _read(path) + if lease is None or lease.get("task") != task: + raise ValueError(f"no lease for task: {task}") + running_live = lease.get("state") == "running" and lease_is_active( + lease, grace_seconds=registry.settings.lease_grace_seconds + ) + if running_live and not force: + raise ValueError("lease is active; pass --force only after confirming the worker stopped") + path.unlink() + append_audit( + registry, + "lease-released", + { + "task_key": task_key(task), + "profile": lease.get("profile"), + "pool": lease.get("pool"), + "forced": force, + }, + ) + return {"task": task, "profile": lease.get("profile"), "released": True} diff --git a/tools/agent-fleet/src/agent_fleet/locks.py b/tools/agent-fleet/src/agent_fleet/locks.py new file mode 100644 index 00000000000..3da6673af96 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/locks.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import json +import os +import shutil +import socket +import time +import uuid +from pathlib import Path +from typing import Any + +from .paths import ensure_private_dir +from .util import atomic_write_json, process_matches, process_start_token + + +class DirectoryLock: + """Portable inter-process lock based on atomic directory creation.""" + + def __init__( + self, + path: Path, + *, + stale_seconds: int, + timeout: float = 10.0, + purpose: str = "exclusive", + ): + self.path = path + self.stale_seconds = stale_seconds + self.timeout = timeout + self.nonce = uuid.uuid4().hex + process_start = process_start_token(os.getpid()) + if process_start is None: + raise RuntimeError("cannot acquire a Fleet lock without a verified process start token") + self.owner = { + "schema": 1, + "pid": os.getpid(), + "process_start": process_start, + "hostname": socket.gethostname(), + "created_unix": time.time(), + "nonce": self.nonce, + "purpose": purpose, + } + self.acquired = False + + @property + def owner_path(self) -> Path: + return self.path / "owner.json" + + def _owner(self) -> dict[str, Any] | None: + try: + value = json.loads(self.owner_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + return value if isinstance(value, dict) else None + + def _age(self, owner: dict[str, Any] | None) -> float: + if owner is not None and isinstance(owner.get("created_unix"), (int, float)): + return max(0.0, time.time() - float(owner["created_unix"])) + try: + return max(0.0, time.time() - self.path.stat().st_mtime) + except FileNotFoundError: + return 0.0 + + def _reclaim_if_stale(self) -> bool: + owner = self._owner() + if self._age(owner) <= self.stale_seconds: + return False + if owner is not None: + if owner.get("hostname") != socket.gethostname(): + return False + pid = owner.get("pid") + if isinstance(pid, int) and process_matches(pid, owner.get("process_start")): + return False + tombstone = self.path.with_name(f".{self.path.name}.stale.{uuid.uuid4().hex}") + try: + os.replace(self.path, tombstone) + except FileNotFoundError: + return True + except OSError: + return False + shutil.rmtree(tombstone, ignore_errors=True) + return True + + def held(self) -> bool: + """Return whether a live owner holds this lock, reclaiming stale state.""" + + if not self.path.exists(): + return False + self._reclaim_if_stale() + return self.path.exists() + + def acquire(self) -> None: + ensure_private_dir(self.path.parent) + deadline = time.monotonic() + self.timeout + while True: + try: + self.path.mkdir(mode=0o700) + except FileExistsError: + self._reclaim_if_stale() + if time.monotonic() >= deadline: + raise TimeoutError(f"timed out acquiring state lock: {self.path}") from None + time.sleep(0.05) + continue + try: + atomic_write_json(self.owner_path, self.owner) + except BaseException: + shutil.rmtree(self.path, ignore_errors=True) + raise + self.acquired = True + return + + def release(self) -> None: + if not self.acquired: + return + owner = self._owner() + if owner is None or owner.get("nonce") != self.nonce: + raise RuntimeError(f"state lock ownership changed unexpectedly: {self.path}") + tombstone = self.path.with_name(f".{self.path.name}.release.{self.nonce}") + os.replace(self.path, tombstone) + shutil.rmtree(tombstone, ignore_errors=False) + self.acquired = False + + def __enter__(self) -> DirectoryLock: + self.acquire() + return self + + def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: + self.release() + + +def state_lock(state_dir: Path, stale_seconds: int, *, timeout: float = 10.0) -> DirectoryLock: + return DirectoryLock( + state_dir / "locks" / "registry.lock", + stale_seconds=stale_seconds, + timeout=timeout, + ) + + +def provider_enrollment_lock( + state_dir: Path, + provider: str, + stale_seconds: int, + *, + timeout: float = 10.0, +) -> DirectoryLock: + return DirectoryLock( + state_dir / "locks" / f"provider-enrollment-{provider}.lock", + stale_seconds=stale_seconds, + timeout=timeout, + purpose="maintenance", + ) + + +def provider_selection_refresh_lock( + state_dir: Path, + provider: str, + stale_seconds: int, + *, + timeout: float = 10.0, +) -> DirectoryLock: + return DirectoryLock( + state_dir / "locks" / f"provider-enrollment-{provider}.lock", + stale_seconds=stale_seconds, + timeout=timeout, + purpose="selection-refresh", + ) + + +def provider_maintenance_active( + state_dir: Path, + provider: str, + stale_seconds: int, +) -> bool: + lock = provider_enrollment_lock( + state_dir, + provider, + stale_seconds, + timeout=0, + ) + if not lock.held(): + return False + owner = lock._owner() + # mkdir necessarily precedes the atomic owner.json write. Avoid treating + # that tiny initialization window as auth maintenance and spuriously + # rejecting a concurrent selection; an owner that stays absent remains + # fail-closed. + for _ in range(10): + if owner is not None: + break + time.sleep(0.005) + owner = lock._owner() + # A missing/corrupt owner is fail-closed. Selection refreshes use the same + # exclusion primitive to serialize quota writes, but they are not auth + # maintenance and therefore do not block another lease commit. + return owner is None or owner.get("purpose") != "selection-refresh" diff --git a/tools/agent-fleet/src/agent_fleet/models.py b/tools/agent-fleet/src/agent_fleet/models.py new file mode 100644 index 00000000000..8742288cb2c --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/models.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +SUPPORTED_PROVIDERS = ("claude", "codex") +PROFILE_SAFETY_POLICIES = ("worker", "manual_only", "desktop_shared") + + +@dataclass(frozen=True) +class ProviderConfig: + name: str + binary: Path + base_home: Path | None = None + hooks_source: Path | None = None + shared_entries: tuple[str, ...] = () + desktop_identity_file: Path | None = None + + +@dataclass(frozen=True) +class Profile: + id: str + provider: str + home: Path + pools: tuple[str, ...] + enabled: bool = False + weight: int = 1 + max_concurrent: int = 2 + reserve_percent: int = 15 + safety_policy: str = "worker" + + def public_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "provider": self.provider, + "home": str(self.home), + "pools": list(self.pools), + "enabled": self.enabled, + "weight": self.weight, + "max_concurrent": self.max_concurrent, + "reserve_percent": self.reserve_percent, + "safety_policy": self.safety_policy, + } + + +@dataclass(frozen=True) +class Settings: + state_dir: Path + share_dir: Path + quota_binary: Path + quota_stale_seconds: int = 300 + quota_verification_grace_seconds: int = 86400 + lease_grace_seconds: int = 30 + active_lease_penalty: int = 8 + lock_stale_seconds: int = 30 + + +@dataclass(frozen=True) +class Registry: + version: int + settings: Settings + providers: dict[str, ProviderConfig] = field(default_factory=dict) + profiles: dict[str, Profile] = field(default_factory=dict) + + def require_profile(self, profile_id: str) -> Profile: + try: + return self.profiles[profile_id] + except KeyError as exc: + raise ValueError(f"unknown profile: {profile_id}") from exc + + def require_provider(self, provider: str) -> ProviderConfig: + try: + return self.providers[provider] + except KeyError as exc: + raise ValueError(f"provider binary is not configured: {provider}") from exc diff --git a/tools/agent-fleet/src/agent_fleet/output.py b/tools/agent-fleet/src/agent_fleet/output.py new file mode 100644 index 00000000000..3100a583b1f --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/output.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import json +import shutil +import subprocess +import sys +from typing import Any + + +def emit(payload: Any, output_format: str) -> None: + if output_format == "json": + json.dump(payload, sys.stdout, sort_keys=True, separators=(",", ":")) + sys.stdout.write("\n") + return + if output_format == "toon": + toon = shutil.which("toon") + if toon is None: + raise ValueError("TOON output requested but `toon` is not on PATH; use --format json") + encoded = json.dumps(payload, separators=(",", ":")) + result = subprocess.run( + [toon], + input=encoded, + text=True, + capture_output=True, + check=False, + ) + if result.returncode != 0: + raise ValueError(f"TOON encoder failed: {result.stderr.strip()}") + sys.stdout.write(result.stdout) + return + _emit_human(payload) + + +def _emit_human(payload: Any, prefix: str = "") -> None: + if isinstance(payload, dict): + for key, value in payload.items(): + if isinstance(value, (dict, list)): + print(f"{prefix}{key}:") + _emit_human(value, prefix + " ") + else: + print(f"{prefix}{key}: {value}") + elif isinstance(payload, list): + for item in payload: + if isinstance(item, (dict, list)): + print(f"{prefix}-") + _emit_human(item, prefix + " ") + else: + print(f"{prefix}- {item}") + else: + print(f"{prefix}{payload}") diff --git a/tools/agent-fleet/src/agent_fleet/paths.py b/tools/agent-fleet/src/agent_fleet/paths.py new file mode 100644 index 00000000000..1620c76648c --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/paths.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import os +from pathlib import Path + + +def expand_path(value: str | Path) -> Path: + return Path(os.path.expandvars(os.path.expanduser(str(value)))).resolve() + + +def default_config_path() -> Path: + override = os.environ.get("AGENT_FLEET_CONFIG") + if override: + return expand_path(override) + return Path.home() / ".config" / "agent-fleet" / "accounts.toml" + + +def default_state_dir() -> Path: + override = os.environ.get("AGENT_FLEET_STATE_DIR") + if override: + return expand_path(override) + return Path.home() / ".local" / "state" / "agent-fleet" + + +def default_share_dir() -> Path: + override = os.environ.get("AGENT_FLEET_SHARE_DIR") + if override: + return expand_path(override) + return Path.home() / ".local" / "share" / "agent-fleet" + + +def ensure_private_dir(path: Path) -> None: + path.mkdir(parents=True, exist_ok=True, mode=0o700) + path.chmod(0o700) diff --git a/tools/agent-fleet/src/agent_fleet/providers.py b/tools/agent-fleet/src/agent_fleet/providers.py new file mode 100644 index 00000000000..b1c821d28ba --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/providers.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import os +import shlex +import shutil +import subprocess +from collections.abc import Iterable +from hashlib import sha256 + +from .models import Profile, Registry + + +def identity_fingerprint(provider: str, identifier: str) -> str: + return sha256(f"{provider}:{identifier}".encode()).hexdigest() + + +def provider_environment(profile: Profile, task: str | None = None) -> dict[str, str]: + env = dict(os.environ) + for name in tuple(env): + if name.startswith(("ANTHROPIC_", "CLAUDE_", "OPENAI_", "CODEX_")): + env.pop(name, None) + env["AGENT_FLEET_PROFILE"] = profile.id + env["AGENT_FLEET_PROVIDER"] = profile.provider + if task: + env["AGENT_FLEET_TASK_ID"] = task + else: + env.pop("AGENT_FLEET_TASK_ID", None) + if profile.provider == "claude": + env["CLAUDE_CONFIG_DIR"] = str(profile.home) + env["DISABLE_LOGIN_COMMAND"] = "1" + env["DISABLE_LOGOUT_COMMAND"] = "1" + elif profile.provider == "codex": + env["CODEX_HOME"] = str(profile.home) + env["CODEX_SQLITE_HOME"] = str(profile.home) + return env + + +def provider_argv(registry: Registry, profile: Profile, command: Iterable[str] = ()) -> list[str]: + binary = registry.require_provider(profile.provider).binary + suffix = list(command) + return [str(binary), *suffix] + + +def login_argv( + registry: Registry, + profile: Profile, + *, + browser_login: bool = False, + access_token: bool = False, +) -> list[str]: + if profile.provider == "claude": + suffix = ["auth", "login"] + elif access_token: + suffix = ["login", "--with-access-token"] + else: + suffix = ["login"] if browser_login else ["login", "--device-auth"] + return provider_argv(registry, profile, suffix) + + +def auth_probe(registry: Registry, profile: Profile) -> dict[str, str | None]: + binary = registry.require_provider(profile.provider).binary + if not binary.exists(): + return { + "status": "binary-missing", + "identity_fingerprint": None, + "identity_source": None, + } + suffix = ["auth", "status", "--json"] if profile.provider == "claude" else ["login", "status"] + try: + result = subprocess.run( + [str(binary), *suffix], + env=provider_environment(profile), + capture_output=True, + text=True, + timeout=15, + check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return { + "status": "unknown", + "identity_fingerprint": None, + "identity_source": None, + } + status = "authenticated" if result.returncode == 0 else "unauthenticated" + return { + "status": status, + "identity_fingerprint": None, + "identity_source": None, + } + + +def auth_status(registry: Registry, profile: Profile) -> str: + return str(auth_probe(registry, profile)["status"]) + + +def session_hook_command() -> str: + configured = os.environ.get("AGENT_FLEET_BIN") + executable = configured or shutil.which("agent-fleet") or "agent-fleet" + return f"{shlex.quote(executable)} --format json hook session-start" + + +def resume_argv( + registry: Registry, + profile: Profile, + session_id: str, + extra: list[str], +) -> list[str]: + if profile.provider == "claude": + return provider_argv(registry, profile, ["--resume", session_id, *extra]) + return provider_argv(registry, profile, ["resume", session_id, *extra]) diff --git a/tools/agent-fleet/src/agent_fleet/provision.py b/tools/agent-fleet/src/agent_fleet/provision.py new file mode 100644 index 00000000000..d61943e6fd5 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/provision.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import json +import os +import stat +import tomllib +from pathlib import Path +from typing import Any + +from .models import Profile, ProviderConfig, Registry +from .paths import ensure_private_dir +from .providers import session_hook_command +from .util import atomic_write_json + +HOOK_MARKER = " hook session-start" + + +def _read_json_object(path: Path) -> dict[str, Any]: + if not path.exists(): + return {} + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"refusing to modify invalid JSON: {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"refusing to modify non-object JSON: {path}") + return value + + +def _merge_source_hooks(payload: dict[str, Any], source: Path | None) -> None: + if source is None or not source.exists(): + return + source_payload = _read_json_object(source) + source_hooks = source_payload.get("hooks", {}) + if not isinstance(source_hooks, dict): + raise ValueError(f"hooks must be an object: {source}") + destination = payload.setdefault("hooks", {}) + if not isinstance(destination, dict): + raise ValueError("destination hooks must be an object") + for event, groups in source_hooks.items(): + if not isinstance(groups, list): + raise ValueError(f"hooks.{event} must be an array: {source}") + target_groups = destination.setdefault(event, []) + if not isinstance(target_groups, list): + raise ValueError(f"destination hooks.{event} must be an array") + existing = {json.dumps(group, sort_keys=True) for group in target_groups} + for group in groups: + encoded = json.dumps(group, sort_keys=True) + if encoded not in existing: + target_groups.append(group) + existing.add(encoded) + + +def _install_session_hook(path: Path, source: Path | None) -> None: + payload = _read_json_object(path) + _merge_source_hooks(payload, source) + hooks = payload.setdefault("hooks", {}) + if not isinstance(hooks, dict): + raise ValueError(f"hooks must be an object: {path}") + groups = hooks.setdefault("SessionStart", []) + if not isinstance(groups, list): + raise ValueError(f"hooks.SessionStart must be an array: {path}") + for group in groups: + if not isinstance(group, dict): + continue + for hook in group.get("hooks", []): + if isinstance(hook, dict) and HOOK_MARKER in str(hook.get("command", "")): + atomic_write_json(path, payload) + return + groups.append( + { + "matcher": "startup|resume|clear|compact", + "hooks": [ + { + "type": "command", + "command": session_hook_command(), + "statusMessage": "Recording Agent Fleet session identity", + } + ], + } + ) + atomic_write_json(path, payload) + + +def _ensure_codex_config(home: Path) -> None: + path = home / "config.toml" + if not path.exists(): + path.write_text( + 'cli_auth_credentials_store = "file"\n\n[features]\nhooks = true\n', + encoding="utf-8", + ) + path.chmod(0o600) + return + try: + raw = tomllib.loads(path.read_text(encoding="utf-8")) + except tomllib.TOMLDecodeError as exc: + raise ValueError(f"invalid Codex config: {path}: {exc}") from exc + store = raw.get("cli_auth_credentials_store") + if store != "file": + raise ValueError( + f'managed Codex profile requires cli_auth_credentials_store="file": {path}' + ) + features = raw.get("features", {}) + if not isinstance(features, dict) or features.get("hooks") is not True: + raise ValueError(f"managed Codex profile requires [features] hooks=true: {path}") + path.chmod(0o600) + + +def _share_workflow_entries(profile: Profile, provider: ProviderConfig) -> list[str]: + if provider.base_home is None: + return [] + shared: list[str] = [] + flags = os.O_RDONLY + if hasattr(os, "O_DIRECTORY"): + flags |= os.O_DIRECTORY + if hasattr(os, "O_NOFOLLOW"): + flags |= os.O_NOFOLLOW + try: + home_fd = os.open(profile.home, flags) + except OSError as exc: + raise ValueError(f"managed profile home is not a safe directory: {profile.home}") from exc + try: + opened = os.fstat(home_fd) + current = profile.home.lstat() + if not stat.S_ISDIR(opened.st_mode) or (opened.st_dev, opened.st_ino) != ( + current.st_dev, + current.st_ino, + ): + raise ValueError(f"managed profile home changed during provisioning: {profile.home}") + for entry in provider.shared_entries: + source = provider.base_home / entry + if not source.exists(): + continue + try: + destination_stat = os.stat(entry, dir_fd=home_fd, follow_symlinks=False) + except FileNotFoundError: + os.symlink( + str(source), + entry, + target_is_directory=source.is_dir(), + dir_fd=home_fd, + ) + destination_stat = os.stat( + entry, + dir_fd=home_fd, + follow_symlinks=False, + ) + if not stat.S_ISLNK(destination_stat.st_mode): + raise ValueError( + f"refusing to replace existing managed workflow path: {profile.home / entry}" + ) + if os.readlink(entry, dir_fd=home_fd) != str(source): + raise ValueError(f"managed shared link points elsewhere: {profile.home / entry}") + shared.append(entry) + finally: + os.close(home_fd) + return shared + + +def provision_profile(registry: Registry, profile: Profile) -> dict[str, Any]: + if profile.home.is_symlink(): + raise ValueError(f"managed profile home cannot be a symlink: {profile.home}") + ensure_private_dir(profile.home) + current = profile.home.lstat() + if not stat.S_ISDIR(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"managed profile home must be a current-user directory: {profile.home}") + provider = registry.require_provider(profile.provider) + shared = _share_workflow_entries(profile, provider) + if profile.provider == "claude": + ensure_private_dir(profile.home / "hooks") + _install_session_hook(profile.home / "settings.json", provider.hooks_source) + elif profile.provider == "codex": + ensure_private_dir(profile.home / "hooks") + _ensure_codex_config(profile.home) + _install_session_hook(profile.home / "hooks.json", provider.hooks_source) + marker = profile.home / ".agent-fleet-profile.json" + atomic_write_json( + marker, + {"schema": 1, "profile": profile.id, "provider": profile.provider}, + ) + os.chmod(marker, 0o600) + return { + "profile": profile.id, + "provider": profile.provider, + "home": str(profile.home), + "provisioned": True, + "shared_entries": shared, + } + + +def profile_is_provisioned(profile: Profile) -> bool: + marker = profile.home / ".agent-fleet-profile.json" + if not profile.home.is_dir() or not marker.is_file(): + return False + try: + raw = json.loads(marker.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return False + return raw.get("profile") == profile.id and raw.get("provider") == profile.provider + + +def profile_hook_health(registry: Registry, profile: Profile) -> dict[str, bool]: + path = profile.home / ("settings.json" if profile.provider == "claude" else "hooks.json") + try: + payload = _read_json_object(path) + except ValueError: + return { + "agent_fleet_session_hook": False, + "herdr_session_hook": False, + "inherited_workflow_hooks": False, + } + hooks = payload.get("hooks", {}) + commands: list[str] = [] + if isinstance(hooks, dict): + for groups in hooks.values(): + if not isinstance(groups, list): + continue + for group in groups: + if not isinstance(group, dict): + continue + entries = group.get("hooks", []) + if not isinstance(entries, list): + continue + commands.extend( + str(entry.get("command", "")) for entry in entries if isinstance(entry, dict) + ) + + source_ok = True + source = registry.require_provider(profile.provider).hooks_source + if source is not None and source.exists(): + try: + source_payload = _read_json_object(source) + except ValueError: + source_ok = False + else: + source_hooks = source_payload.get("hooks", {}) + if not isinstance(source_hooks, dict) or not isinstance(hooks, dict): + source_ok = False + else: + for event, groups in source_hooks.items(): + destination_groups = hooks.get(event, []) + if not isinstance(groups, list) or not isinstance(destination_groups, list): + source_ok = False + break + destination = { + json.dumps(group, sort_keys=True) for group in destination_groups + } + if any( + json.dumps(group, sort_keys=True) not in destination for group in groups + ): + source_ok = False + break + return { + "agent_fleet_session_hook": any(HOOK_MARKER in command for command in commands), + "herdr_session_hook": any("herdr-agent-state" in command for command in commands), + "inherited_workflow_hooks": source_ok, + } + + +def profile_shared_assets_healthy(registry: Registry, profile: Profile) -> bool: + provider = registry.require_provider(profile.provider) + if provider.base_home is None: + return True + for entry in provider.shared_entries: + source = provider.base_home / entry + if not source.exists(): + continue + destination = profile.home / entry + if not destination.is_symlink() or destination.resolve() != source.resolve(): + return False + return True diff --git a/tools/agent-fleet/src/agent_fleet/quota.py b/tools/agent-fleet/src/agent_fleet/quota.py new file mode 100644 index 00000000000..a6bb6b98c03 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/quota.py @@ -0,0 +1,524 @@ +from __future__ import annotations + +import json +import os +import re +import stat +import subprocess +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from .locks import DirectoryLock +from .models import Profile, Registry +from .paths import ensure_private_dir +from .providers import auth_probe, identity_fingerprint, provider_environment +from .util import atomic_write_bytes, atomic_write_json, utc_now + +PRIMARY_WINDOWS = {"five_hour", "seven_day", "weekly", "week"} +AUTH_FAILURE_RE = re.compile(r"sign[- ]?in|required|reauth|token.*(?:revoked|invalid)", re.I) +RATE_LIMIT_RE = re.compile(r"rate.?limit", re.I) +HARD_BLOCKED_STATUSES = {"auth_required", "rate_limited", "error"} +FALLBACK_STATUSES = {"fresh", "stale", "unavailable"} +KNOWN_STATUSES = HARD_BLOCKED_STATUSES | FALLBACK_STATUSES +SAFE_TOKEN_RE = re.compile(r"[A-Za-z0-9_.:-]{1,128}") + + +@dataclass(frozen=True) +class QuotaCacheSnapshot: + existed: bool + payload: bytes = b"" + mode: int = 0o600 + + +def quota_path(registry: Registry, profile_id: str) -> Path: + return registry.settings.state_dir / "quota" / f"{profile_id}.json" + + +def snapshot_quota_cache(registry: Registry, profile_id: str) -> QuotaCacheSnapshot: + path = quota_path(registry, profile_id) + try: + current = path.lstat() + except FileNotFoundError: + return QuotaCacheSnapshot(False) + if not stat.S_ISREG(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"quota cache must be a current-user regular file: {path}") + return QuotaCacheSnapshot( + True, + path.read_bytes(), + stat.S_IMODE(current.st_mode), + ) + + +def restore_quota_cache( + registry: Registry, + profile_id: str, + snapshot: QuotaCacheSnapshot, +) -> None: + path = quota_path(registry, profile_id) + if snapshot.existed: + atomic_write_bytes(path, snapshot.payload, mode=snapshot.mode) + return + try: + current = path.lstat() + except FileNotFoundError: + return + if not stat.S_ISREG(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"refusing to remove unexpected quota cache path: {path}") + path.unlink() + + +def _number(value: Any) -> float | None: + if isinstance(value, (int, float)) and not isinstance(value, bool): + return max(0.0, min(100.0, float(value))) + return None + + +def _parse_time(value: Any) -> datetime | None: + if not isinstance(value, str) or not value: + return None + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=UTC) + + +def _age_seconds(value: Any) -> int | None: + parsed = _parse_time(value) + if parsed is None: + return None + return max(0, int((datetime.now(UTC) - parsed.astimezone(UTC)).total_seconds())) + + +def _effective_status(status: str, error: str | None) -> tuple[str, str | None]: + if status != "stale" or not error: + return status, None + if AUTH_FAILURE_RE.search(error): + return "auth_required", "cached_after_auth_failure" + if RATE_LIMIT_RE.search(error): + return "rate_limited", "cached_after_rate_limit" + return status, None + + +def _safe_token(value: Any, fallback: str) -> str: + if isinstance(value, str) and SAFE_TOKEN_RE.fullmatch(value): + return value + return fallback + + +def _safe_reason(value: Any) -> str | None: + if value is None: + return None + return _safe_token(value, "provider_reported_error") + + +def _normalized_windows(value: Any) -> list[dict[str, Any]]: + windows: list[dict[str, Any]] = [] + if not isinstance(value, list): + return windows + for item in value: + if not isinstance(item, dict): + continue + remaining = _number(item.get("remaining_percent")) + if remaining is None: + continue + resets_at = item.get("resets_at") + windows.append( + { + "id": _safe_token(item.get("id"), "unknown"), + "kind": _safe_token(item.get("kind"), "unknown"), + "remaining_percent": remaining, + "resets_at": resets_at if _parse_time(resets_at) is not None else None, + } + ) + return windows + + +def _identity_fingerprint(profile: Profile, provider: dict[str, Any]) -> str | None: + account = provider.get("account") + if not isinstance(account, dict): + return None + identifier = account.get("accountId") or account.get("account_id") + if not isinstance(identifier, str) or not identifier: + return None + return identity_fingerprint(profile.provider, identifier) + + +def _normalize(profile: Profile, raw: dict[str, Any]) -> dict[str, Any]: + if raw.get("profile") == profile.id and raw.get("schema") == 1: + status = _safe_token(raw.get("status"), "unavailable") + if status not in KNOWN_STATUSES: + status = "unavailable" + fingerprint = raw.get("identity_fingerprint") + if not isinstance(fingerprint, str) or len(fingerprint) != 64: + fingerprint = None + windows = _normalized_windows(raw.get("windows")) + verified_at = raw.get("verified_at") + if _parse_time(verified_at) is None: + verified_at = None + now = utc_now() + if status == "fresh" and verified_at is None: + status = "error" + return { + "schema": 1, + "profile": profile.id, + "provider": profile.provider, + "status": status, + "reported_status": status, + "reason": ( + "missing_remote_verification_timestamp" + if status == "error" and raw.get("status") == "fresh" + else _safe_reason(raw.get("reason")) + ), + "headroom_percent": _number(raw.get("headroom_percent")), + "windows": windows, + "verified_at": verified_at, + "identity_fingerprint": fingerprint, + "identity_source": "quota-account" if fingerprint else None, + "refreshed_at": now, + } + + providers = raw.get("providers", []) + if not isinstance(providers, list): + raise ValueError("quota-axi response has no providers array") + provider = next( + ( + item + for item in providers + if isinstance(item, dict) and item.get("provider") == profile.provider + ), + None, + ) + if provider is None: + raise ValueError(f"quota-axi response has no {profile.provider} result") + state = provider.get("state", {}) + reported_status = _safe_token( + state.get("status") if isinstance(state, dict) else None, + "unavailable", + ) + if reported_status not in KNOWN_STATUSES: + reported_status = "unavailable" + error = state.get("error") if isinstance(state, dict) else None + error = error if isinstance(error, str) else None + status, derived_reason = _effective_status(reported_status, error) + windows_raw = provider.get("windows", []) + windows: list[dict[str, Any]] = [] + if isinstance(windows_raw, list): + for item in windows_raw: + if not isinstance(item, dict): + continue + remaining = _number(item.get("percentRemaining")) + if remaining is None: + used = _number(item.get("percentUsed")) + remaining = None if used is None else 100.0 - used + if remaining is None: + continue + windows.append( + { + "id": _safe_token(item.get("id"), "unknown"), + "kind": _safe_token(item.get("kind"), "unknown"), + "remaining_percent": remaining, + "resets_at": ( + item.get("resetsAt") + if _parse_time(item.get("resetsAt")) is not None + else None + ), + } + ) + primary = [item for item in windows if item["id"] in PRIMARY_WINDOWS] + if not primary: + primary = [item for item in windows if item["kind"] in {"session", "weekly"}] + headroom = min((item["remaining_percent"] for item in primary), default=None) + reason = state.get("reason") if isinstance(state, dict) else None + verified_at = state.get("refreshedAt") if isinstance(state, dict) else None + if _parse_time(verified_at) is None: + verified_at = None + if status == "fresh": + status = "error" + derived_reason = "missing_remote_verification_timestamp" + identity_fingerprint = _identity_fingerprint(profile, provider) + return { + "schema": 1, + "profile": profile.id, + "provider": profile.provider, + "status": status, + "reported_status": reported_status, + "reason": derived_reason or _safe_reason(reason), + "headroom_percent": headroom, + "windows": windows, + "verified_at": verified_at, + "identity_fingerprint": identity_fingerprint, + "identity_source": "quota-account" if identity_fingerprint else None, + "refreshed_at": utc_now(), + } + + +def _load_raw_quota( + registry: Registry, + profile: Profile, + *, + timeout: int = 30, + allow_keychain_prompt: bool = False, +) -> dict[str, Any]: + fixture_dir = os.environ.get("AGENT_FLEET_QUOTA_FIXTURE_DIR") + if fixture_dir: + fixture = Path(fixture_dir) / f"{profile.id}.json" + try: + raw = json.loads(fixture.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise ValueError(f"quota fixture not found: {fixture}") from exc + else: + binary = registry.settings.quota_binary + if not binary.is_file() or not os.access(binary, os.X_OK): + raise ValueError(f"configured quota-axi candidate is not executable: {binary}") + cache = profile.home / ".agent-fleet-quota-cache" + ensure_private_dir(cache) + env = provider_environment(profile) + env["XDG_CACHE_HOME"] = str(cache) + if profile.provider == "codex": + env["QUOTA_AXI_CODEX_BINARY"] = str(registry.require_provider("codex").binary) + try: + argv = [str(binary), "--provider", profile.provider, "--json", "--full"] + if allow_keychain_prompt and profile.provider == "claude": + argv.append("--allow-keychain-prompt") + result = subprocess.run( + argv, + env=env, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except subprocess.TimeoutExpired as exc: + raise ValueError(f"quota refresh timed out for {profile.id}") from exc + try: + raw = json.loads(result.stdout) + except json.JSONDecodeError as exc: + if result.returncode != 0: + detail = result.stderr.strip().splitlines()[-1:] or ["unknown error"] + raise ValueError(f"quota refresh failed for {profile.id}: {detail[0]}") from exc + raise ValueError(f"quota-axi emitted invalid JSON for {profile.id}") from exc + if result.returncode != 0: + providers = raw.get("providers") if isinstance(raw, dict) else None + fresh = isinstance(providers, list) and any( + isinstance(item, dict) + and item.get("provider") == profile.provider + and isinstance(item.get("state"), dict) + and item["state"].get("status") == "fresh" + for item in providers + ) + if fresh: + raise ValueError( + f"quota-axi exited nonzero with untrusted fresh data for {profile.id}" + ) + if not isinstance(raw, dict): + raise ValueError(f"quota response for {profile.id} must be an object") + return raw + + +def probe_quota( + registry: Registry, + profile: Profile, + *, + timeout: int = 30, + allow_keychain_prompt: bool = False, +) -> dict[str, Any]: + raw = _load_raw_quota( + registry, + profile, + timeout=timeout, + allow_keychain_prompt=allow_keychain_prompt, + ) + normalized = _normalize(profile, raw) + if profile.provider == "claude": + probe = auth_probe(registry, profile) + fingerprint = probe.get("identity_fingerprint") + existing = normalized.get("identity_fingerprint") + if fingerprint and existing and fingerprint != existing: + normalized["status"] = "error" + normalized["reason"] = "identity_source_mismatch" + normalized["verified_at"] = None + normalized["identity_fingerprint"] = None + normalized["identity_source"] = None + elif fingerprint: + normalized["identity_fingerprint"] = fingerprint + normalized["identity_source"] = probe.get("identity_source") + return normalized + + +def store_quota( + registry: Registry, + profile: Profile, + normalized: dict[str, Any], +) -> dict[str, Any]: + normalized = dict(normalized) + normalized["profile"] = profile.id + normalized["provider"] = profile.provider + path = quota_path(registry, profile.id) + try: + previous = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + previous = {} + if isinstance(previous, dict): + if not normalized.get("verified_at") and _parse_time(previous.get("verified_at")): + normalized["verified_at"] = previous["verified_at"] + fingerprint = normalized.get("identity_fingerprint") + fingerprint_is_verified = ( + normalized.get("status") == "fresh" + and normalized.get("verified_at") is not None + and normalized.get("headroom_percent") is not None + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + ) + previous_fingerprint = previous.get("identity_fingerprint") + if ( + not fingerprint_is_verified + and isinstance(previous_fingerprint, str) + and len(previous_fingerprint) == 64 + and _parse_time(previous.get("verified_at")) is not None + ): + normalized["identity_fingerprint"] = previous_fingerprint + normalized["identity_source"] = "previous-verified-quota-account" + atomic_write_json(path, normalized) + return normalized + + +def refresh_quota( + registry: Registry, + profile: Profile, + *, + timeout: int = 30, + allow_keychain_prompt: bool = False, +) -> dict[str, Any]: + normalized = probe_quota( + registry, + profile, + timeout=timeout, + allow_keychain_prompt=allow_keychain_prompt, + ) + return store_quota(registry, profile, normalized) + + +def read_quota(registry: Registry, profile_id: str) -> dict[str, Any]: + path = quota_path(registry, profile_id) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return { + "schema": 1, + "profile": profile_id, + "status": "unavailable", + "reason": "no_cache", + "headroom_percent": None, + "windows": [], + "refreshed_at": None, + "verified_at": None, + "verified_recent": False, + "fresh": False, + } + if not isinstance(value, dict): + return {"profile": profile_id, "status": "unavailable", "fresh": False} + age = _age_seconds(value.get("refreshed_at")) + verification_age = _age_seconds(value.get("verified_at")) + value["age_seconds"] = age + value["verification_age_seconds"] = verification_age + value["verified_recent"] = ( + verification_age is not None + and verification_age <= registry.settings.quota_verification_grace_seconds + ) + value["fresh"] = ( + value.get("status") == "fresh" + and age is not None + and age <= registry.settings.quota_stale_seconds + and _number(value.get("headroom_percent")) is not None + ) + return value + + +def quota_routeability( + registry: Registry, + profile: Profile, + *, + quota: dict[str, Any] | None = None, + authentication: str | None = None, + ignore_reserve: bool = False, +) -> dict[str, Any]: + current = quota or read_quota(registry, profile.id) + if profile.safety_policy != "worker": + return { + "eligible": False, + "mode": "blocked", + "reason": f"safety_policy_{profile.safety_policy}", + } + if authentication is not None and authentication != "authenticated": + return { + "eligible": False, + "mode": "blocked", + "reason": f"local_auth_{authentication}", + } + status = str(current.get("status", "unavailable")) + if status in HARD_BLOCKED_STATUSES: + return { + "eligible": False, + "mode": "blocked", + "reason": str(current.get("reason") or status), + } + if current.get("fresh") is True: + headroom = _number(current.get("headroom_percent")) + if headroom is None: + return {"eligible": False, "mode": "blocked", "reason": "invalid_quota"} + if not ignore_reserve and headroom <= profile.reserve_percent: + return {"eligible": False, "mode": "quota", "reason": "quota_reserve"} + return {"eligible": True, "mode": "quota", "reason": "fresh"} + if status in FALLBACK_STATUSES and current.get("verified_recent") is True: + return { + "eligible": True, + "mode": "verified-fallback", + "reason": "recent_remote_verification", + } + reason = "remote_verification_expired" if current.get("verified_at") else "remote_unverified" + return {"eligible": False, "mode": "blocked", "reason": reason} + + +def refresh_due_quotas( + registry: Registry, + profiles: list[Profile], + *, + timeout: int = 4, +) -> dict[str, str | None]: + due = [ + profile + for profile in profiles + if (age := read_quota(registry, profile.id).get("age_seconds")) is None + or int(age) > registry.settings.quota_stale_seconds + ] + if not due: + return {} + + def refresh_one(profile: Profile) -> None: + lock = DirectoryLock( + registry.settings.state_dir / "locks" / f"quota-{profile.id}.lock", + stale_seconds=max(registry.settings.lock_stale_seconds, timeout + 1), + timeout=timeout + 0.5, + ) + with lock: + age = read_quota(registry, profile.id).get("age_seconds") + if age is not None and int(age) <= registry.settings.quota_stale_seconds: + return + refresh_quota(registry, profile, timeout=timeout) + + outcomes: dict[str, str | None] = {} + with ThreadPoolExecutor(max_workers=min(8, len(due))) as executor: + futures = {executor.submit(refresh_one, profile): profile.id for profile in due} + for future in as_completed(futures): + profile_id = futures[future] + try: + future.result() + except (OSError, TimeoutError, ValueError) as exc: + outcomes[profile_id] = str(exc) + else: + outcomes[profile_id] = None + return outcomes diff --git a/tools/agent-fleet/src/agent_fleet/scheduler.py b/tools/agent-fleet/src/agent_fleet/scheduler.py new file mode 100644 index 00000000000..1556af00e5a --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/scheduler.py @@ -0,0 +1,381 @@ +from __future__ import annotations + +from collections import Counter +from datetime import UTC, datetime +from typing import Any + +from .audit import append_audit +from .cooldowns import read_cooldown +from .enrollment import recover_pending_codex_transactions +from .identity import identity_conflict, refresh_provider_identity_anchors_if_due +from .leases import active_leases, bind_lease, get_active_lease, new_lease, write_lease +from .locks import ( + provider_maintenance_active, + provider_selection_refresh_lock, + state_lock, +) +from .models import Profile, Registry +from .providers import auth_status +from .provision import profile_is_provisioned +from .quota import quota_routeability, read_quota, refresh_due_quotas +from .util import atomic_write_json, task_key + + +def _last_selected(registry: Registry) -> dict[str, str]: + path = registry.settings.state_dir / "selection.json" + try: + import json + + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return {} + return raw if isinstance(raw, dict) else {} + + +def _stamp_selection(registry: Registry, profile_id: str) -> None: + selected = _last_selected(registry) + selected[profile_id] = datetime.now(UTC).isoformat() + atomic_write_json(registry.settings.state_dir / "selection.json", selected) + + +def _candidate_record( + registry: Registry, + profile: Profile, + active_count: int, + last_selected: str, + registry_order: int, + authentication: str, +) -> dict[str, Any]: + quota = read_quota(registry, profile.id) + fresh = quota.get("fresh") is True + headroom = quota.get("headroom_percent") if fresh else None + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication, + ) + return { + "profile": profile, + "active": active_count, + "quota": quota, + "fresh": fresh, + "headroom": headroom, + "eligible": routeability["eligible"], + "selection_mode": routeability["mode"], + "block_reason": routeability["reason"], + "last_selected": last_selected, + "registry_order": registry_order, + } + + +def _choose( + records: list[dict[str, Any]], penalty: int, *, ignore_reserve: bool = False +) -> dict[str, Any]: + fresh = [ + record + for record in records + if record["fresh"] + and (record["eligible"] or (ignore_reserve and record["block_reason"] == "quota_reserve")) + ] + if fresh: + return max( + fresh, + key=lambda item: ( + float(item["headroom"]) + + (item["profile"].weight - 1) * 5 + - item["active"] * penalty, + -item["active"], + item["last_selected"] == "", + _inverse_timestamp(item["last_selected"]), + -item["registry_order"], + ), + ) + fallback = [record for record in records if record["eligible"] and not record["fresh"]] + if not fallback: + reasons = sorted({str(record["block_reason"]) for record in records}) + if reasons == ["quota_reserve"]: + raise ValueError("all eligible profiles are at or below their quota reserve") + detail = ", ".join(reasons) or "unknown" + raise ValueError(f"no remotely verified profile is routeable: {detail}") + return min( + fallback, + key=lambda item: ( + item["active"] / item["profile"].weight, + item["last_selected"], + item["registry_order"], + ), + ) + + +def _inverse_timestamp(value: str) -> float: + if not value: + return float("inf") + try: + return -datetime.fromisoformat(value).timestamp() + except ValueError: + return float("inf") + + +def _select_and_acquire( + registry: Registry, + *, + task: str, + pool: str, + provider: str | None = None, + profile_id: str | None = None, + bind_pid: int | None = None, + dry_run: bool = False, + explicit_profile: bool = False, + ignore_reserve: bool = False, + recovery_reservation: bool = False, +) -> dict[str, Any]: + if ignore_reserve and profile_id is None: + raise ValueError("ignoring quota reserve requires an explicit profile") + if recovery_reservation and not ignore_reserve: + raise ValueError("recovery reservation must ignore the new-task quota reserve") + scoped_provider_names = { + profile.provider + for profile in registry.profiles.values() + if profile.safety_policy == "worker" + and (explicit_profile or pool in profile.pools) + and (provider is None or profile.provider == provider) + and (profile_id is None or profile.id == profile_id) + } + scoped_profiles = [ + profile + for profile in registry.profiles.values() + if profile.enabled + and profile.safety_policy == "worker" + and (explicit_profile or pool in profile.pools) + and (provider is None or profile.provider == provider) + and (profile_id is None or profile.id == profile_id) + and profile_is_provisioned(profile) + ] + # Quota/identity writes share the provider maintenance interlock, but a + # selection does not hold it while committing its lease. The state-lock + # check below closes both races: maintenance either owns the marker first + # (selection refuses) or observes the newly committed lease and aborts. + for provider_name in sorted(scoped_provider_names): + if provider_maintenance_active( + registry.settings.state_dir, + provider_name, + registry.settings.lock_stale_seconds, + ): + raise ValueError( + f"provider maintenance is in progress for {provider_name}; " + "refusing to start a new Fleet lease" + ) + try: + with provider_selection_refresh_lock( + registry.settings.state_dir, + provider_name, + registry.settings.lock_stale_seconds, + ): + recover_pending_codex_transactions(registry, provider_name) + refresh_provider_identity_anchors_if_due(registry, provider_name) + refresh_due_quotas( + registry, + [profile for profile in scoped_profiles if profile.provider == provider_name], + ) + except TimeoutError as exc: + raise ValueError( + f"provider maintenance is in progress for {provider_name}; " + "refusing to start a new Fleet lease" + ) from exc + authentication = {profile.id: auth_status(registry, profile) for profile in scoped_profiles} + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + blocked_provider = next( + ( + provider_name + for provider_name in sorted(scoped_provider_names) + if provider_maintenance_active( + registry.settings.state_dir, + provider_name, + registry.settings.lock_stale_seconds, + ) + ), + None, + ) + if blocked_provider is not None: + raise ValueError( + f"provider maintenance is in progress for {blocked_provider}; " + "refusing to start a new Fleet lease" + ) + leases = active_leases(registry, prune=True) + existing = get_active_lease(registry, task) + if existing is not None: + profile = registry.require_profile(str(existing.get("profile"))) + if recovery_reservation and existing.get("state") == "running": + raise ValueError("task already has a live worker lease; refusing recovery") + if existing.get("pool") != pool: + raise ValueError(f"task already owns a lease in pool {existing.get('pool')}") + if provider is not None and profile.provider != provider: + raise ValueError(f"task is already bound to provider {profile.provider}") + if profile_id is not None and profile.id != profile_id: + raise ValueError(f"task is already bound to profile {profile.id}") + if existing.get("state") != "running": + if not profile.enabled or not profile_is_provisioned(profile): + raise ValueError("sticky profile is disabled or unprovisioned") + quota = read_quota(registry, profile.id) + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication.get(profile.id, auth_status(registry, profile)), + # The reservation already passed new-task policy. Binding it + # must recheck auth/readiness without applying reserve twice. + ignore_reserve=True, + ) + if not routeability["eligible"]: + raise ValueError(f"sticky profile is not routeable: {routeability['reason']}") + conflict = identity_conflict(registry, profile, quota) + if conflict is not None: + raise ValueError(f"sticky profile identity is not routeable: {conflict}") + else: + quota = read_quota(registry, profile.id) + if bind_pid is not None and not dry_run: + owner_pid = existing.get("pid") + if ( + existing.get("state") == "running" + and isinstance(owner_pid, int) + and owner_pid != bind_pid + ): + raise ValueError(f"task lease is already owned by live process {owner_pid}") + existing = bind_lease(registry, existing, bind_pid) + append_audit( + registry, + "lease-bound", + { + "profile": profile.id, + "provider": profile.provider, + "pool": pool, + "task_key": task_key(task), + }, + ) + return { + "schema": 1, + "task": task, + "pool": pool, + "profile": profile.id, + "provider": profile.provider, + "decision_reason": "sticky", + "quota_fresh": quota.get("fresh"), + "headroom_percent": quota.get("headroom_percent"), + "active_lease_count": sum(lease.get("profile") == profile.id for lease in leases), + "degraded": False, + "dry_run": dry_run, + "lease": existing, + } + + counts = Counter(str(lease.get("profile")) for lease in leases) + candidates = [ + profile + for profile in registry.profiles.values() + if profile.enabled + and profile.safety_policy == "worker" + and (explicit_profile or pool in profile.pools) + and (provider is None or profile.provider == provider) + and (profile_id is None or profile.id == profile_id) + and profile_is_provisioned(profile) + and read_cooldown(registry, profile.id) is None + and counts[profile.id] < profile.max_concurrent + ] + if not candidates: + constraint = profile_id or provider or pool + raise ValueError(f"no enabled, provisioned profile has capacity for {constraint}") + last_selected = _last_selected(registry) + registry_order = {profile_id: index for index, profile_id in enumerate(registry.profiles)} + records = [ + _candidate_record( + registry, + profile, + counts[profile.id], + last_selected.get(profile.id, ""), + registry_order[profile.id], + authentication.get(profile.id, "unknown"), + ) + for profile in candidates + ] + for record in records: + conflict = identity_conflict(registry, record["profile"], record["quota"]) + if conflict is not None: + record["eligible"] = False + record["selection_mode"] = "blocked" + record["block_reason"] = ( + "duplicate_provider_identity" + if conflict.startswith(("managed:", "base_", "desktop_")) + else conflict + ) + selected = _choose( + records, + registry.settings.active_lease_penalty, + ignore_reserve=ignore_reserve, + ) + profile = selected["profile"] + quota = selected["quota"] + if ignore_reserve: + reason = "sticky-resume" + else: + reason = "quota" if selected["fresh"] else "verified-fallback" + lease = None + if not dry_run: + lease = new_lease(task, profile.id, pool, pid=bind_pid) + write_lease(registry, lease) + _stamp_selection(registry, profile.id) + append_audit( + registry, + "lease-selected", + { + "profile": profile.id, + "provider": profile.provider, + "pool": pool, + "task_key": task_key(task), + "decision_reason": reason, + "lease_state": lease["state"], + }, + ) + return { + "schema": 1, + "task": task, + "pool": pool, + "profile": profile.id, + "provider": profile.provider, + "decision_reason": reason, + "quota_fresh": quota.get("fresh"), + "headroom_percent": quota.get("headroom_percent"), + "active_lease_count": counts[profile.id], + "degraded": not selected["fresh"], + "dry_run": dry_run, + "lease": lease, + } + + +def select_and_acquire( + registry: Registry, + *, + task: str, + pool: str, + provider: str | None = None, + profile_id: str | None = None, + bind_pid: int | None = None, + dry_run: bool = False, + explicit_profile: bool = False, + ignore_reserve: bool = False, + recovery_reservation: bool = False, +) -> dict[str, Any]: + return _select_and_acquire( + registry, + task=task, + pool=pool, + provider=provider, + profile_id=profile_id, + bind_pid=bind_pid, + dry_run=dry_run, + explicit_profile=explicit_profile, + ignore_reserve=ignore_reserve, + recovery_reservation=recovery_reservation, + ) diff --git a/tools/agent-fleet/src/agent_fleet/sessions.py b/tools/agent-fleet/src/agent_fleet/sessions.py new file mode 100644 index 00000000000..7712f99b2a6 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/sessions.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path +from typing import Any + +from .audit import append_audit +from .leases import get_active_lease +from .locks import state_lock +from .models import Registry +from .util import atomic_write_json, task_key, utc_now, validate_id + +SESSION_KEYS = ("session_id", "sessionId") + + +def session_path(registry: Registry, task: str) -> Path: + return registry.settings.state_dir / "sessions" / f"{task_key(task)}.json" + + +def _find_session_id(value: Any) -> str | None: + if isinstance(value, dict): + for key in SESSION_KEYS: + item = value.get(key) + if isinstance(item, str) and item: + return item + for item in value.values(): + found = _find_session_id(item) + if found: + return found + elif isinstance(value, list): + for item in value: + found = _find_session_id(item) + if found: + return found + return None + + +def record_session_from_hook(registry: Registry, payload: dict[str, Any]) -> dict[str, Any]: + task = os.environ.get("AGENT_FLEET_TASK_ID") + profile_id = os.environ.get("AGENT_FLEET_PROFILE") + provider = os.environ.get("AGENT_FLEET_PROVIDER") + if not task or not profile_id or not provider: + return {"recorded": False, "reason": "not_agent_fleet_launch"} + session_id = _find_session_id(payload) + if session_id is None: + raise ValueError("SessionStart hook payload did not contain a session id") + validate_id(session_id, "session id") + profile = registry.require_profile(profile_id) + if profile.provider != provider: + raise ValueError("hook provider does not match registered profile") + with state_lock( + registry.settings.state_dir, + registry.settings.lock_stale_seconds, + ): + lease = get_active_lease(registry, task) + if lease is None or lease.get("profile") != profile.id: + raise ValueError("hook task does not own a live lease for this profile") + mapping = { + "schema": 1, + "task": task, + "profile": profile.id, + "provider": profile.provider, + "pool": lease.get("pool"), + "session_id": session_id, + "updated_at": utc_now(), + } + atomic_write_json(session_path(registry, task), mapping) + append_audit( + registry, + "session-recorded", + { + "task_key": task_key(task), + "profile": profile.id, + "provider": profile.provider, + "pool": lease.get("pool"), + }, + ) + return { + "recorded": True, + "task": task, + "profile": profile.id, + "provider": profile.provider, + } + + +def read_hook_payload() -> dict[str, Any]: + try: + value = json.load(sys.stdin) + except json.JSONDecodeError as exc: + raise ValueError("SessionStart hook input is not valid JSON") from exc + if not isinstance(value, dict): + raise ValueError("SessionStart hook input must be a JSON object") + return value + + +def get_session(registry: Registry, task: str) -> dict[str, Any]: + path = session_path(registry, task) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise ValueError(f"no recorded provider session for task: {task}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"invalid session mapping: {path}") from exc + if not isinstance(value, dict) or value.get("task") != task: + raise ValueError(f"corrupt session mapping: {path}") + return value + + +def remove_session(registry: Registry, task: str) -> dict[str, Any]: + mapping = get_session(registry, task) + session_path(registry, task).unlink() + append_audit( + registry, + "session-removed", + { + "task_key": task_key(task), + "profile": mapping.get("profile"), + "provider": mapping.get("provider"), + "pool": mapping.get("pool"), + }, + ) + return { + "task": task, + "profile": mapping.get("profile"), + "provider": mapping.get("provider"), + "removed": True, + } diff --git a/tools/agent-fleet/src/agent_fleet/status.py b/tools/agent-fleet/src/agent_fleet/status.py new file mode 100644 index 00000000000..621059a4f1a --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/status.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from collections import Counter +from typing import Any + +from .cooldowns import read_cooldown +from .identity import identity_conflict, refresh_provider_identity_anchors_if_due +from .leases import active_leases +from .models import SUPPORTED_PROVIDERS, Registry +from .providers import auth_status +from .provision import profile_is_provisioned +from .quota import quota_routeability, read_quota, refresh_due_quotas + + +def profile_status(registry: Registry, profile_id: str) -> dict[str, Any]: + profile = registry.require_profile(profile_id) + if profile.safety_policy == "worker": + refresh_provider_identity_anchors_if_due(registry, profile.provider) + leases = active_leases(registry) + provisioned = profile_is_provisioned(profile) + authentication = auth_status(registry, profile) if provisioned else "not-provisioned" + quota = read_quota(registry, profile.id) + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication, + ) + conflict = identity_conflict(registry, profile, quota) + if conflict is not None: + routeability = { + "eligible": False, + "mode": "blocked", + "reason": conflict, + } + return { + **profile.public_dict(), + "provisioned": provisioned, + "auth_status": authentication, + "active_leases": sum(lease.get("profile") == profile.id for lease in leases), + "cooldown": read_cooldown(registry, profile.id), + "quota": quota, + "routeability": routeability, + } + + +def pool_status( + registry: Registry, + *, + pool: str, + provider: str | None = None, +) -> dict[str, Any]: + leases = active_leases(registry) + counts = Counter(str(lease.get("profile")) for lease in leases) + providers = [provider] if provider else list(SUPPORTED_PROVIDERS) + for provider_name in providers: + refresh_provider_identity_anchors_if_due(registry, provider_name) + scoped = [ + profile + for profile in registry.profiles.values() + if profile.provider in providers and pool in profile.pools and profile.enabled + ] + refresh_due_quotas(registry, scoped) + summaries: list[dict[str, Any]] = [] + for provider_name in providers: + profiles: list[dict[str, Any]] = [] + for profile in registry.profiles.values(): + if profile.provider != provider_name or pool not in profile.pools: + continue + provisioned = profile_is_provisioned(profile) + authentication = auth_status(registry, profile) if provisioned else "not-provisioned" + quota = read_quota(registry, profile.id) + fresh = quota.get("fresh") is True + headroom = quota.get("headroom_percent") if fresh else None + cooldown = read_cooldown(registry, profile.id) + capacity = counts[profile.id] < profile.max_concurrent + routeability = quota_routeability( + registry, + profile, + quota=quota, + authentication=authentication, + ) + conflict = identity_conflict(registry, profile, quota) + if conflict is not None: + routeability = { + "eligible": False, + "mode": "blocked", + "reason": conflict, + } + eligible = ( + profile.enabled + and provisioned + and capacity + and cooldown is None + and routeability["eligible"] + ) + adjusted = ( + float(headroom) + + (profile.weight - 1) * 5 + - counts[profile.id] * registry.settings.active_lease_penalty + if eligible and fresh + else None + ) + profiles.append( + { + "profile": profile.id, + "enabled": profile.enabled, + "provisioned": provisioned, + "auth_status": authentication, + "active_leases": counts[profile.id], + "max_concurrent": profile.max_concurrent, + "quota_fresh": fresh, + "quota_status": quota.get("status"), + "verified_recent": quota.get("verified_recent"), + "headroom_percent": headroom, + "adjusted_headroom_percent": adjusted, + "reserve_percent": profile.reserve_percent, + "cooldown": cooldown, + "eligible": eligible, + "routeability_reason": routeability["reason"], + "identity_fingerprint": quota.get("identity_fingerprint"), + } + ) + fresh_eligible = [item for item in profiles if item["eligible"] and item["quota_fresh"]] + fallback_eligible = [ + item for item in profiles if item["eligible"] and not item["quota_fresh"] + ] + mode = ( + "quota" + if fresh_eligible + else "verified-fallback" + if fallback_eligible + else "unavailable" + ) + eligible_profiles = fresh_eligible or fallback_eligible + summaries.append( + { + "provider": provider_name, + "available": bool(eligible_profiles), + "selection_mode": mode, + "degraded": mode == "verified-fallback", + "best_adjusted_headroom_percent": max( + (float(item["adjusted_headroom_percent"]) for item in fresh_eligible), + default=None, + ), + "eligible_profiles": len(eligible_profiles), + "active_leases": sum(counts[item["profile"]] for item in profiles), + "profiles": profiles, + } + ) + return {"schema": 1, "pool": pool, "providers": summaries} diff --git a/tools/agent-fleet/src/agent_fleet/util.py b/tools/agent-fleet/src/agent_fleet/util.py new file mode 100644 index 00000000000..2e57679d19b --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/util.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import hashlib +import json +import os +import re +import subprocess +import tempfile +from contextlib import suppress +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +SAFE_ID = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9._-]{0,127}$") + + +def validate_id(value: str, kind: str) -> str: + if not SAFE_ID.fullmatch(value): + raise ValueError(f"invalid {kind}: use 1-128 letters, numbers, dot, underscore, or hyphen") + if "@" in value: + raise ValueError(f"invalid {kind}: account emails are not allowed") + return value + + +def utc_now() -> str: + return datetime.now(UTC).isoformat() + + +def task_key(task: str) -> str: + return hashlib.sha256(task.encode("utf-8")).hexdigest() + + +def atomic_write_json(path: Path, payload: Any, mode: int = 0o600) -> None: + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + fd, temp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + os.fchmod(fd, mode) + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp_name, path) + path.chmod(mode) + except BaseException: + with suppress(FileNotFoundError): + os.unlink(temp_name) + raise + + +def atomic_write_bytes(path: Path, payload: bytes, mode: int = 0o600) -> None: + path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) + fd, temp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + os.fchmod(fd, mode) + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temp_name, path) + path.chmod(mode) + except BaseException: + with suppress(FileNotFoundError): + os.unlink(temp_name) + raise + + +def read_json(path: Path) -> Any: + with path.open("r", encoding="utf-8") as handle: + return json.load(handle) + + +def process_start_token(pid: int) -> str | None: + try: + result = subprocess.run( + ["ps", "-p", str(pid), "-o", "lstart="], + check=False, + capture_output=True, + text=True, + ) + except OSError: + return None + token = result.stdout.strip() + return token or None + + +def process_matches(pid: int, start_token: str | None) -> bool: + if pid <= 0 or start_token is None: + return False + current = process_start_token(pid) + if current is None: + return False + return current == start_token diff --git a/tools/agent-fleet/tests/conftest.py b/tools/agent-fleet/tests/conftest.py new file mode 100644 index 00000000000..5b1620185f3 --- /dev/null +++ b/tools/agent-fleet/tests/conftest.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import json +import os +import stat +from dataclasses import replace +from pathlib import Path + +import pytest + +from agent_fleet.config import initial_registry, save_registry +from agent_fleet.identity import refresh_provider_identity_anchors +from agent_fleet.models import ProviderConfig, Registry +from agent_fleet.quota import refresh_quota + + +@pytest.fixture +def fleet(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Registry, Path]: + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_ps = fake_bin / "ps" + fake_ps.write_text("#!/bin/sh\necho fixture-process-start\n", encoding="utf-8") + fake_ps.chmod(fake_ps.stat().st_mode | stat.S_IXUSR) + monkeypatch.setenv("PATH", f"{fake_bin}:{os.environ.get('PATH', '')}") + config = tmp_path / "config" / "accounts.toml" + state = tmp_path / "state" + share = tmp_path / "share" + monkeypatch.setenv("AGENT_FLEET_CONFIG", str(config)) + monkeypatch.setenv("AGENT_FLEET_STATE_DIR", str(state)) + monkeypatch.setenv("AGENT_FLEET_SHARE_DIR", str(share)) + registry = initial_registry(3, 2) + quota_binary = tmp_path / "quota-axi" + quota_binary.write_text( + """#!/usr/bin/env python3 +import json +import os +import sys +from datetime import UTC, datetime +provider = sys.argv[sys.argv.index("--provider") + 1] +profile = os.environ["AGENT_FLEET_PROFILE"] +print(json.dumps({ + "providers": [{ + "provider": provider, + "account": {"accountId": profile + "-account"}, + "state": { + "status": "fresh", + "refreshedAt": datetime.now(UTC).isoformat(), + }, + "windows": [{ + "id": "five_hour", + "kind": "session", + "percentRemaining": 80, + }], + }] +})) +""", + encoding="utf-8", + ) + quota_binary.chmod(quota_binary.stat().st_mode | stat.S_IXUSR) + registry = replace( + registry, + settings=replace(registry.settings, quota_binary=quota_binary), + ) + providers: dict[str, ProviderConfig] = {} + desktop_file = tmp_path / "desktop" / "claude-config.json" + desktop_file.parent.mkdir(parents=True) + desktop_file.write_text( + json.dumps({"lastKnownAccountUuid": "desktop-captain-account"}), + encoding="utf-8", + ) + for name, provider in registry.providers.items(): + base = tmp_path / "base" / name + base.mkdir(parents=True) + hooks = base / ("settings.json" if name == "claude" else "hooks.json") + hooks.write_text( + json.dumps( + { + "hooks": { + "SessionStart": [ + { + "matcher": "startup", + "hooks": [{"type": "command", "command": "base-hook"}], + } + ] + } + } + ), + encoding="utf-8", + ) + shared = "CLAUDE.md" if name == "claude" else "AGENTS.md" + (base / shared).write_text("workflow rules\n", encoding="utf-8") + binary = tmp_path / f"provider-{name}" + binary.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8") + binary.chmod(binary.stat().st_mode | stat.S_IXUSR) + providers[name] = replace( + provider, + binary=binary, + base_home=base, + hooks_source=hooks, + shared_entries=(shared,), + desktop_identity_file=desktop_file if name == "claude" else None, + ) + registry = replace(registry, providers=providers) + save_registry(registry, config) + for profile in registry.profiles.values(): + refresh_quota(registry, profile) + for provider in registry.providers: + refresh_provider_identity_anchors(registry, provider) + return registry, config diff --git a/tools/agent-fleet/tests/test_config_and_provision.py b/tools/agent-fleet/tests/test_config_and_provision.py new file mode 100644 index 00000000000..eb38a12b3dd --- /dev/null +++ b/tools/agent-fleet/tests/test_config_and_provision.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import json +import os +import stat +from dataclasses import replace +from pathlib import Path + +import pytest + +from agent_fleet.config import load_registry +from agent_fleet.provision import ( + profile_hook_health, + profile_is_provisioned, + provision_profile, +) + + +def test_initial_registry_is_dynamic_disabled_and_private( + fleet: tuple[object, Path], +) -> None: + _, path = fleet + registry = load_registry(path) + assert sorted(registry.profiles) == [ + "claude-1", + "claude-2", + "claude-3", + "codex-1", + "codex-2", + ] + assert not any(profile.enabled for profile in registry.profiles.values()) + assert stat.S_IMODE(path.stat().st_mode) == 0o600 + + +def test_provision_merges_hooks_and_shares_only_declared_workflow_assets( + fleet: tuple[object, Path], +) -> None: + _, path = fleet + registry = load_registry(path) + profile = registry.require_profile("codex-1") + result = provision_profile(registry, profile) + + assert result["shared_entries"] == ["AGENTS.md"] + assert profile_is_provisioned(profile) + assert (profile.home / "AGENTS.md").is_symlink() + assert not (profile.home / "auth.json").exists() + hooks = json.loads((profile.home / "hooks.json").read_text(encoding="utf-8")) + commands = [ + hook["command"] for group in hooks["hooks"]["SessionStart"] for hook in group["hooks"] + ] + assert "base-hook" in commands + assert any("agent-fleet" in command for command in commands) + config = (profile.home / "config.toml").read_text(encoding="utf-8") + assert 'cli_auth_credentials_store = "file"' in config + assert "hooks = true" in config + health = profile_hook_health(registry, profile) + assert health["agent_fleet_session_hook"] is True + assert health["inherited_workflow_hooks"] is True + assert health["herdr_session_hook"] is False + + +def test_profile_ids_reject_account_email(fleet: tuple[object, Path]) -> None: + _, path = fleet + text = path.read_text(encoding="utf-8") + text += '\n[profiles."person@example.com"]\nprovider="claude"\npools=["claude-crew"]\n' + path.write_text(text, encoding="utf-8") + with pytest.raises(ValueError, match="invalid profile id"): + load_registry(path) + + +def test_legacy_registry_without_desktop_field_migrates_to_safe_default( + fleet: tuple[object, Path], +) -> None: + _, path = fleet + lines = [ + line + for line in path.read_text(encoding="utf-8").splitlines() + if not line.startswith("desktop_identity_file =") + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + registry = load_registry(path) + assert ( + registry.require_provider("claude").desktop_identity_file + == (Path.home() / "Library/Application Support/Claude/config.json").resolve() + ) + + +def test_explicit_desktop_anchor_opt_out_remains_supported( + fleet: tuple[object, Path], +) -> None: + _, path = fleet + text = path.read_text(encoding="utf-8") + line = next(item for item in text.splitlines() if item.startswith("desktop_identity_file =")) + path.write_text(text.replace(line, "desktop_identity_file = false"), encoding="utf-8") + assert load_registry(path).require_provider("claude").desktop_identity_file is None + + +def test_provision_refuses_symlink_profile_home(fleet: tuple[object, Path], tmp_path: Path) -> None: + _, path = fleet + registry = load_registry(path) + external = tmp_path / "external" + external.mkdir() + symlink_home = tmp_path / "profile-link" + symlink_home.symlink_to(external, target_is_directory=True) + profile = replace(registry.require_profile("codex-1"), home=symlink_home) + with pytest.raises(ValueError, match="cannot be a symlink"): + provision_profile(registry, profile) + + +def test_shared_asset_install_refuses_existing_non_symlink( + fleet: tuple[object, Path], +) -> None: + _, path = fleet + registry = load_registry(path) + profile = registry.require_profile("codex-1") + profile.home.mkdir(parents=True, exist_ok=True) + os.chmod(profile.home, 0o700) + (profile.home / "AGENTS.md").write_text("attacker-owned\n", encoding="utf-8") + with pytest.raises(ValueError, match="refusing to replace"): + provision_profile(registry, profile) diff --git a/tools/agent-fleet/tests/test_contract_status.py b/tools/agent-fleet/tests/test_contract_status.py new file mode 100644 index 00000000000..39e2a051e18 --- /dev/null +++ b/tools/agent-fleet/tests/test_contract_status.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import json +import os +import stat +import subprocess +import sys +import tomllib +from dataclasses import replace +from pathlib import Path + +from agent_fleet import __version__ +from agent_fleet.config import load_registry, save_registry +from agent_fleet.providers import identity_fingerprint +from agent_fleet.provision import provision_profile +from agent_fleet.quota import quota_path +from agent_fleet.util import atomic_write_json, utc_now + + +def _auth_ok_binary(path: Path) -> None: + path.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8") + path.chmod(path.stat().st_mode | stat.S_IXUSR) + + +def test_contract_and_version_do_not_require_registry(tmp_path: Path) -> None: + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + missing = tmp_path / "missing.toml" + for command in ("contract", "version"): + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(missing), + command, + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + check=True, + ) + payload = json.loads(result.stdout) + assert payload["contract_version"] == 1 + + +def test_runtime_version_matches_package_metadata() -> None: + project_root = Path(__file__).parents[1] + metadata = tomllib.loads((project_root / "pyproject.toml").read_text(encoding="utf-8")) + assert __version__ == metadata["project"]["version"] + + +def test_pool_status_reports_provider_level_fallback( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + binary = tmp_path / "auth-ok" + _auth_ok_binary(binary) + providers = dict(registry.providers) + providers["codex"] = replace(providers["codex"], binary=binary) + profiles = dict(registry.profiles) + for profile_id in ("codex-1", "codex-2"): + profiles[profile_id] = replace(profiles[profile_id], enabled=True) + registry = replace(registry, providers=providers, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + for profile_id in ("codex-1", "codex-2"): + provision_profile(registry, registry.require_profile(profile_id)) + now = utc_now() + atomic_write_json( + quota_path(registry, profile_id), + { + "schema": 1, + "profile": profile_id, + "provider": "codex", + "status": "stale", + "headroom_percent": 50, + "windows": [], + "identity_fingerprint": identity_fingerprint("codex", f"{profile_id}-account"), + "verified_at": now, + "refreshed_at": now, + }, + ) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "pool", + "status", + "--pool", + "codex-crew", + "--provider", + "codex", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + check=True, + ) + provider = json.loads(result.stdout)["providers"][0] + assert provider["available"] is True + assert provider["selection_mode"] == "verified-fallback" + assert provider["degraded"] is True + + +def test_enroll_uses_codex_device_auth_then_verifies_while_disabled( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + log = tmp_path / "provider-argv.jsonl" + binary = tmp_path / "provider" + binary.write_text( + """#!/usr/bin/env python3 +import json +import os +import sys +if sys.argv[1:] == ["login", "--device-auth"]: + home = os.environ["CODEX_HOME"] + with open(os.path.join(home, "auth.json"), "w", encoding="utf-8") as handle: + json.dump({"tokens": "fake-test-only"}, handle) +with open(os.environ["FAKE_PROVIDER_LOG"], "a", encoding="utf-8") as handle: + handle.write(json.dumps(sys.argv[1:]) + "\\n") +""", + encoding="utf-8", + ) + binary.chmod(binary.stat().st_mode | stat.S_IXUSR) + providers = dict(registry.providers) + providers["codex"] = replace(providers["codex"], binary=binary) + registry = replace(registry, providers=providers) + save_registry(registry, config) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + env["FAKE_PROVIDER_LOG"] = str(log) + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "profile", + "enroll", + "codex-1", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + check=True, + ) + + payload = json.loads(result.stdout) + calls = [json.loads(line) for line in log.read_text(encoding="utf-8").splitlines()] + assert calls[0] == ["login", "--device-auth"] + assert ["login", "status"] in calls + assert payload["credential_verified"] is True + assert payload["enabled"] is False + assert load_registry(config).require_profile("codex-1").enabled is False + + +def test_verify_keeps_a_remotely_rejected_profile_disabled( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + provision_profile(registry, registry.require_profile("codex-1")) + fixtures = tmp_path / "quota" + fixtures.mkdir() + (fixtures / "codex-1.json").write_text( + json.dumps( + { + "providers": [ + { + "provider": "codex", + "state": {"status": "auth_required"}, + "windows": [], + } + ] + } + ), + encoding="utf-8", + ) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + env["AGENT_FLEET_QUOTA_FIXTURE_DIR"] = str(fixtures) + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "profile", + "verify", + "codex-1", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + check=True, + ) + + payload = json.loads(result.stdout) + assert payload["ready"] is False + assert payload["profiles"][0]["enabled"] is False + assert load_registry(config).require_profile("codex-1").enabled is False diff --git a/tools/agent-fleet/tests/test_exec_sessions.py b/tools/agent-fleet/tests/test_exec_sessions.py new file mode 100644 index 00000000000..4ed9ea68c5f --- /dev/null +++ b/tools/agent-fleet/tests/test_exec_sessions.py @@ -0,0 +1,291 @@ +from __future__ import annotations + +import json +import os +import stat +import subprocess +import sys +from dataclasses import replace +from pathlib import Path + +from agent_fleet.config import load_registry, save_registry +from agent_fleet.provision import provision_profile +from agent_fleet.scheduler import select_and_acquire +from agent_fleet.sessions import get_session, record_session_from_hook + + +def _fake_provider(path: Path) -> None: + path.write_text( + """#!/usr/bin/env python3 +import json +import os +import sys +print(json.dumps({ + "argv": sys.argv[1:], + "profile": os.environ.get("AGENT_FLEET_PROFILE"), + "task": os.environ.get("AGENT_FLEET_TASK_ID"), + "codex_home": os.environ.get("CODEX_HOME"), + "claude_home": os.environ.get("CLAUDE_CONFIG_DIR"), + "has_openai_key": "OPENAI_API_KEY" in os.environ, + "has_anthropic_key": "ANTHROPIC_API_KEY" in os.environ, +})) +""", + encoding="utf-8", + ) + path.chmod(path.stat().st_mode | stat.S_IXUSR) + + +def test_exec_uses_selected_home_and_clears_ambient_credentials( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + fake = tmp_path / "fake-provider" + _fake_provider(fake) + providers = dict(registry.providers) + providers["codex"] = replace(providers["codex"], binary=fake) + profiles = dict(registry.profiles) + profiles["codex-1"] = replace(profiles["codex-1"], enabled=True) + registry = replace(registry, providers=providers, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + profile = registry.require_profile("codex-1") + provision_profile(registry, profile) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + env["OPENAI_API_KEY"] = "must-not-reach-provider" + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--config", + str(config), + "exec", + "--task", + "exec-task", + "--pool", + "codex-crew", + "--profile", + "codex-1", + "--", + "example-argument", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=True, + ) + payload = json.loads(result.stdout) + assert payload["argv"] == ["example-argument"] + assert payload["profile"] == "codex-1" + assert payload["task"] == "exec-task" + assert payload["codex_home"] == str(profile.home) + assert payload["claude_home"] is None + assert payload["has_openai_key"] is False + + +def test_session_hook_persists_profile_and_provider_session( + fleet: tuple[object, Path], monkeypatch +) -> None: + _, config = fleet + registry = load_registry(config) + profiles = dict(registry.profiles) + profiles["claude-1"] = replace(profiles["claude-1"], enabled=True) + registry = replace(registry, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + profile = registry.require_profile("claude-1") + provision_profile(registry, profile) + select_and_acquire( + registry, + task="hook-task", + pool="claude-crew", + profile_id="claude-1", + bind_pid=os.getpid(), + ) + monkeypatch.setenv("AGENT_FLEET_TASK_ID", "hook-task") + monkeypatch.setenv("AGENT_FLEET_PROFILE", "claude-1") + monkeypatch.setenv("AGENT_FLEET_PROVIDER", "claude") + + result = record_session_from_hook( + registry, {"hook_event_name": "SessionStart", "session_id": "session-123"} + ) + mapping = get_session(registry, "hook-task") + assert result["recorded"] is True + assert mapping["profile"] == "claude-1" + assert mapping["provider"] == "claude" + assert mapping["session_id"] == "session-123" + assert mapping["pool"] == "claude-crew" + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + status = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "session", + "status", + "--task", + "hook-task", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=True, + ) + assert json.loads(status.stdout)["session_id"] == "session-123" + + removed = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "session", + "remove", + "--task", + "hook-task", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=True, + ) + assert json.loads(removed.stdout)["removed"] is True + + +def test_direct_resume_without_managed_task_is_refused( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + fake = tmp_path / "fake-provider" + _fake_provider(fake) + providers = dict(registry.providers) + providers["codex"] = replace(providers["codex"], binary=fake) + profiles = dict(registry.profiles) + profiles["codex-1"] = replace(profiles["codex-1"], enabled=True) + registry = replace(registry, providers=providers, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + profile = registry.require_profile("codex-1") + provision_profile(registry, profile) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--config", + str(config), + "resume", + "--profile", + "codex-1", + "--session", + "session-explicit", + "--", + "extra", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert "worker resume requires --task" in result.stderr + + +def test_direct_worker_exec_without_managed_task_is_refused( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + fake = tmp_path / "fake-provider" + _fake_provider(fake) + providers = dict(registry.providers) + providers["codex"] = replace(providers["codex"], binary=fake) + profiles = dict(registry.profiles) + profiles["codex-1"] = replace(profiles["codex-1"], enabled=True) + registry = replace(registry, providers=providers, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + provision_profile(registry, registry.require_profile("codex-1")) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--config", + str(config), + "exec", + "--profile", + "codex-1", + "--", + "example-argument", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert "worker exec requires --task" in result.stderr + + +def test_live_task_cannot_be_rebound(fleet: tuple[object, Path]) -> None: + _, config = fleet + registry = load_registry(config) + profiles = dict(registry.profiles) + profiles["claude-1"] = replace(profiles["claude-1"], enabled=True) + registry = replace(registry, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + provision_profile(registry, registry.require_profile("claude-1")) + select_and_acquire( + registry, + task="owned-task", + pool="claude-crew", + profile_id="claude-1", + bind_pid=os.getpid(), + ) + try: + select_and_acquire( + registry, + task="owned-task", + pool="claude-crew", + profile_id="claude-1", + bind_pid=os.getpid() + 1, + ) + except ValueError as exc: + assert "already owned" in str(exc) + else: + raise AssertionError("live lease was rebound") diff --git a/tools/agent-fleet/tests/test_quota_scheduler.py b/tools/agent-fleet/tests/test_quota_scheduler.py new file mode 100644 index 00000000000..bc84ba7c094 --- /dev/null +++ b/tools/agent-fleet/tests/test_quota_scheduler.py @@ -0,0 +1,498 @@ +from __future__ import annotations + +import json +import os +import stat +import subprocess +import sys +from dataclasses import replace +from datetime import UTC, datetime +from pathlib import Path + +import pytest + +from agent_fleet.config import load_registry, save_registry +from agent_fleet.cooldowns import set_cooldown +from agent_fleet.leases import active_leases, release_lease +from agent_fleet.provision import provision_profile +from agent_fleet.quota import quota_path, refresh_quota +from agent_fleet.scheduler import select_and_acquire +from agent_fleet.sessions import session_path +from agent_fleet.util import atomic_write_json + + +def _enable_and_provision(registry, config: Path, ids: list[str], *, reserve: int = 15): + profiles = dict(registry.profiles) + for profile_id in ids: + profiles[profile_id] = replace( + profiles[profile_id], enabled=True, reserve_percent=reserve, max_concurrent=50 + ) + registry = replace(registry, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + for profile_id in ids: + provision_profile(registry, registry.require_profile(profile_id)) + return registry + + +def _quota_fixture(path: Path, provider: str, remaining: int) -> None: + def write(target: Path, account: str) -> None: + target.write_text( + json.dumps( + { + "providers": [ + { + "provider": provider, + "account": {"accountId": account}, + "state": { + "status": "fresh", + "refreshedAt": datetime.now(UTC).isoformat(), + }, + "windows": [ + { + "id": "five_hour", + "kind": "session", + "percentRemaining": remaining, + } + ], + } + ] + } + ), + encoding="utf-8", + ) + + write(path, f"{path.stem}-account") + base = path.parent / f"{provider}-base-anchor.json" + if not base.exists(): + write(base, f"{provider}-base-anchor-account") + + +def test_fresh_quota_selects_best_safe_profile( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["codex-1", "codex-2"]) + fixtures = tmp_path / "quota" + fixtures.mkdir() + _quota_fixture(fixtures / "codex-1.json", "codex", 82) + _quota_fixture(fixtures / "codex-2.json", "codex", 37) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + for profile_id in ("codex-1", "codex-2"): + refresh_quota(registry, registry.require_profile(profile_id)) + + selected = select_and_acquire(registry, task="quota-task", pool="codex-crew", provider="codex") + assert selected["lease"]["profile"] == "codex-1" + assert selected["decision_reason"] == "quota" + + +def test_fresh_profile_below_reserve_is_excluded( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["claude-1", "claude-2"]) + profiles = dict(registry.profiles) + profiles["claude-1"] = replace(profiles["claude-1"], reserve_percent=90) + registry = replace(registry, profiles=profiles) + save_registry(registry, config) + fixtures = tmp_path / "quota" + fixtures.mkdir() + _quota_fixture(fixtures / "claude-1.json", "claude", 80) + _quota_fixture(fixtures / "claude-2.json", "claude", 40) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + for profile_id in ("claude-1", "claude-2"): + refresh_quota(registry, registry.require_profile(profile_id)) + + selected = select_and_acquire( + registry, task="reserve-task", pool="claude-crew", provider="claude" + ) + assert selected["lease"]["profile"] == "claude-2" + + +def test_sticky_resume_can_reacquire_its_profile_below_reserve( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["codex-1"]) + fixtures = tmp_path / "quota" + fixtures.mkdir() + _quota_fixture(fixtures / "codex-1.json", "codex", 5) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + refresh_quota(registry, registry.require_profile("codex-1")) + + with pytest.raises(ValueError, match="quota reserve"): + select_and_acquire( + registry, + task="new-task", + pool="codex-crew", + provider="codex", + profile_id="codex-1", + ) + + atomic_write_json( + session_path(registry, "resumed-task"), + { + "schema": 1, + "task": "resumed-task", + "pool": "codex-crew", + "profile": "codex-1", + "provider": "codex", + "session_id": "sticky-session", + }, + ) + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "lease", + "recover", + "--task", + "resumed-task", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=True, + ) + selected = json.loads(result.stdout) + assert selected["profile"] == "codex-1" + assert selected["decision_reason"] == "sticky-resume" + + select_and_acquire( + registry, + task="resumed-task", + pool="codex-crew", + profile_id="codex-1", + bind_pid=os.getpid(), + ) + with pytest.raises(ValueError, match="live worker lease"): + select_and_acquire( + registry, + task="resumed-task", + pool="codex-crew", + profile_id="codex-1", + ignore_reserve=True, + recovery_reservation=True, + ) + + +def test_reserved_lease_can_be_released_without_force( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["claude-1"]) + selected = select_and_acquire( + registry, + task="failed-pane-create", + pool="claude-crew", + profile_id="claude-1", + ) + assert selected["lease"]["state"] == "reserved" + + released = release_lease(registry, "failed-pane-create") + assert released["released"] is True + assert active_leases(registry) == [] + + +def test_dry_run_does_not_lease_or_change_selection( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["claude-1"]) + selected = select_and_acquire( + registry, + task="dry-run-task", + pool="claude-crew", + provider="claude", + dry_run=True, + ) + assert selected["dry_run"] is True + assert selected["lease"] is None + assert active_leases(registry) == [] + assert not (registry.settings.state_dir / "selection.json").exists() + + +def test_cooldown_excludes_profile(fleet: tuple[object, Path]) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["codex-1", "codex-2"]) + set_cooldown(registry, "codex-1", seconds=60, reason="spawn-failure") + selected = select_and_acquire( + registry, + task="cooldown-task", + pool="codex-crew", + provider="codex", + ) + assert selected["profile"] == "codex-2" + + +def test_nonzero_quota_exit_preserves_structured_provider_status( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + binary = tmp_path / "quota-axi" + binary.write_text( + """#!/usr/bin/env python3 +import json +import sys +print(json.dumps({ + "providers": [{ + "provider": "claude", + "account": {"accountId": "claude-1-account"}, + "state": {"status": "auth_required", "reason": "keychain_access_required"}, + "windows": [], + }] +})) +sys.exit(1) +""", + encoding="utf-8", + ) + binary.chmod(binary.stat().st_mode | stat.S_IXUSR) + registry = replace( + registry, + settings=replace(registry.settings, quota_binary=binary), + ) + + quota = refresh_quota(registry, registry.require_profile("claude-1")) + assert quota["status"] == "auth_required" + assert quota["reason"] == "keychain_access_required" + assert quota["headroom_percent"] is None + + +def test_stale_cache_after_remote_auth_failure_is_never_selected( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["codex-1"]) + fixtures = tmp_path / "quota" + fixtures.mkdir() + (fixtures / "codex-1.json").write_text( + json.dumps( + { + "providers": [ + { + "provider": "codex", + "state": { + "status": "stale", + "error": "Codex sign-in required", + "refreshedAt": datetime.now(UTC).isoformat(), + }, + "windows": [ + { + "id": "five_hour", + "kind": "session", + "percentRemaining": 76, + } + ], + } + ] + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + + quota = refresh_quota(registry, registry.require_profile("codex-1")) + + assert quota["reported_status"] == "stale" + assert quota["status"] == "auth_required" + assert quota["reason"] == "cached_after_auth_failure" + with pytest.raises(ValueError, match="cached_after_auth_failure"): + select_and_acquire( + registry, + task="revoked-cache", + pool="codex-crew", + profile_id="codex-1", + ) + assert active_leases(registry) == [] + + +def test_verified_stale_fallback_has_a_bounded_grace( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["claude-1"]) + fixtures = tmp_path / "quota" + fixtures.mkdir() + _quota_fixture(fixtures / "claude-1.json", "claude", 70) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + refresh_quota(registry, registry.require_profile("claude-1")) + cached = json.loads(quota_path(registry, "claude-1").read_text(encoding="utf-8")) + cached["status"] = "stale" + cached["refreshed_at"] = datetime.now(UTC).isoformat() + atomic_write_json(quota_path(registry, "claude-1"), cached) + + selected = select_and_acquire( + registry, + task="transient-outage", + pool="claude-crew", + profile_id="claude-1", + dry_run=True, + ) + assert selected["decision_reason"] == "verified-fallback" + + cached["verified_at"] = "2000-01-01T00:00:00+00:00" + atomic_write_json(quota_path(registry, "claude-1"), cached) + with pytest.raises(ValueError, match="remote_verification_expired"): + select_and_acquire( + registry, + task="expired-proof", + pool="claude-crew", + profile_id="claude-1", + dry_run=True, + ) + + +def test_duplicate_remote_identity_is_fail_closed( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["codex-1", "codex-2"]) + fixtures = tmp_path / "quota" + fixtures.mkdir() + for profile_id in ("codex-1", "codex-2"): + (fixtures / f"{profile_id}.json").write_text( + json.dumps( + { + "providers": [ + { + "provider": "codex", + "account": {"accountId": "same-account"}, + "state": { + "status": "fresh", + "refreshedAt": datetime.now(UTC).isoformat(), + }, + "windows": [ + { + "id": "five_hour", + "kind": "session", + "percentRemaining": 80, + } + ], + } + ] + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("AGENT_FLEET_QUOTA_FIXTURE_DIR", str(fixtures)) + for profile_id in ("codex-1", "codex-2"): + refresh_quota(registry, registry.require_profile(profile_id)) + + with pytest.raises(ValueError, match="duplicate_provider_identity"): + select_and_acquire( + registry, + task="duplicate-account", + pool="codex-crew", + provider="codex", + ) + + +def test_interactive_claude_verification_explicitly_allows_keychain_prompt( + fleet: tuple[object, Path], monkeypatch, tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + argv_log = tmp_path / "argv.json" + binary = tmp_path / "quota-axi" + binary.write_text( + """#!/usr/bin/env python3 +import json +import os +import sys +from datetime import UTC, datetime +with open(os.environ["ARGV_LOG"], "w", encoding="utf-8") as handle: + json.dump(sys.argv[1:], handle) +print(json.dumps({ + "providers": [{ + "provider": "claude", + "account": {"accountId": "claude-1-account"}, + "state": { + "status": "fresh", + "refreshedAt": datetime.now(UTC).isoformat(), + }, + "windows": [{ + "id": "five_hour", + "kind": "session", + "percentRemaining": 80, + }], + }] +})) +""", + encoding="utf-8", + ) + binary.chmod(binary.stat().st_mode | stat.S_IXUSR) + registry = replace( + registry, + settings=replace(registry.settings, quota_binary=binary), + ) + monkeypatch.setenv("ARGV_LOG", str(argv_log)) + + refresh_quota( + registry, + registry.require_profile("claude-1"), + allow_keychain_prompt=True, + ) + + assert "--allow-keychain-prompt" in json.loads(argv_log.read_text(encoding="utf-8")) + + +def test_concurrent_reservations_are_atomic_and_balanced( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, ["claude-1", "claude-2"]) + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + processes = [ + subprocess.Popen( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "lease", + "choose", + "--task", + f"concurrent-{index}", + "--pool", + "claude-crew", + "--provider", + "claude", + ], + cwd=project_root, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + for index in range(12) + ] + results = [process.communicate(timeout=20) for process in processes] + failures = [ + {"returncode": process.returncode, "stdout": stdout, "stderr": stderr} + for process, (stdout, stderr) in zip(processes, results, strict=True) + if process.returncode != 0 + ] + assert failures == [] + leases = active_leases(registry) + assert len(leases) == 12 + counts = { + profile_id: sum(lease["profile"] == profile_id for lease in leases) + for profile_id in ("claude-1", "claude-2") + } + assert abs(counts["claude-1"] - counts["claude-2"]) <= 1 diff --git a/tools/agent-fleet/tests/test_safety_transactions.py b/tools/agent-fleet/tests/test_safety_transactions.py new file mode 100644 index 00000000000..bd96787dd08 --- /dev/null +++ b/tools/agent-fleet/tests/test_safety_transactions.py @@ -0,0 +1,352 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from dataclasses import replace +from pathlib import Path + +import pytest + +from agent_fleet import enrollment, leases, locks +from agent_fleet.config import load_registry, save_registry +from agent_fleet.enrollment import ( + activate_codex_promotion, + create_codex_login_stage, + discard_codex_promotion, + discard_codex_stage, + finalize_codex_promotion, + prepare_codex_promotion, + recover_pending_codex_transaction, +) +from agent_fleet.leases import bind_lease, new_lease +from agent_fleet.locks import provider_enrollment_lock +from agent_fleet.providers import provider_environment +from agent_fleet.provision import provision_profile +from agent_fleet.quota import quota_path, snapshot_quota_cache +from agent_fleet.scheduler import select_and_acquire +from agent_fleet.util import atomic_write_json + + +def _enable_and_provision(registry, config: Path, profile_id: str): + profiles = dict(registry.profiles) + profiles[profile_id] = replace( + profiles[profile_id], + enabled=True, + max_concurrent=20, + ) + registry = replace(registry, profiles=profiles) + save_registry(registry, config) + registry = load_registry(config) + provision_profile(registry, registry.require_profile(profile_id)) + return registry + + +def _prepared_codex_transaction(registry): + target = registry.require_profile("codex-1") + provision_profile(registry, target) + old_auth = target.home / "auth.json" + old_auth.write_text('{"token":"old-test-token"}\n', encoding="utf-8") + old_auth.chmod(0o600) + old_quota = quota_path(registry, target.id).read_bytes() + snapshot = snapshot_quota_cache(registry, target.id) + stage = create_codex_login_stage(target) + staged_auth = stage.home / "auth.json" + staged_auth.write_text('{"token":"new-test-token"}\n', encoding="utf-8") + staged_auth.chmod(0o600) + promotion = prepare_codex_promotion(registry, target, stage) + return target, stage, promotion, snapshot, old_auth, old_quota + + +def test_codex_crash_recovery_rolls_back_auth_and_quota_before_cleanup( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = load_registry(config) + target, stage, promotion, snapshot, old_auth, old_quota = _prepared_codex_transaction(registry) + transaction = activate_codex_promotion( + registry, + target, + promotion, + snapshot, + ) + assert "new-test-token" in old_auth.read_text(encoding="utf-8") + atomic_write_json(quota_path(registry, target.id), {"schema": 1, "new": True}) + + with provider_enrollment_lock( + registry.settings.state_dir, + "codex", + registry.settings.lock_stale_seconds, + ): + assert recover_pending_codex_transaction(registry, target) is True + + assert "old-test-token" in old_auth.read_text(encoding="utf-8") + assert quota_path(registry, target.id).read_bytes() == old_quota + assert not transaction.journal_path.exists() + discard_codex_promotion(promotion, target) + discard_codex_stage(stage, target) + + +def test_codex_activation_failure_after_replace_is_recovered( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + registry = load_registry(config) + target, stage, promotion, snapshot, old_auth, old_quota = _prepared_codex_transaction(registry) + real_fsync = enrollment._fsync_directory + calls = 0 + + def fail_after_replace(path: Path) -> None: + nonlocal calls + calls += 1 + if calls == 2: + raise OSError("injected crash after auth replace") + real_fsync(path) + + monkeypatch.setattr(enrollment, "_fsync_directory", fail_after_replace) + with pytest.raises(OSError, match="injected crash"): + activate_codex_promotion(registry, target, promotion, snapshot) + monkeypatch.setattr(enrollment, "_fsync_directory", real_fsync) + + with provider_enrollment_lock( + registry.settings.state_dir, + "codex", + registry.settings.lock_stale_seconds, + ): + assert recover_pending_codex_transaction(registry, target) is True + assert "old-test-token" in old_auth.read_text(encoding="utf-8") + assert quota_path(registry, target.id).read_bytes() == old_quota + discard_codex_promotion(promotion, target) + discard_codex_stage(stage, target) + + +def test_codex_committed_crash_recovery_keeps_promoted_auth( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + registry = load_registry(config) + target, stage, promotion, snapshot, target_auth, _ = _prepared_codex_transaction(registry) + transaction = activate_codex_promotion(registry, target, promotion, snapshot) + real_unlink = Path.unlink + + def fail_journal_cleanup(path: Path, *args, **kwargs) -> None: + if path == transaction.journal_path: + raise OSError("injected crash after durable commit") + real_unlink(path, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", fail_journal_cleanup) + with pytest.raises(OSError, match="injected crash"): + finalize_codex_promotion(registry, target, transaction) + monkeypatch.setattr(Path, "unlink", real_unlink) + + with provider_enrollment_lock( + registry.settings.state_dir, + "codex", + registry.settings.lock_stale_seconds, + ): + assert recover_pending_codex_transaction(registry, target) is True + assert "new-test-token" in target_auth.read_text(encoding="utf-8") + assert not transaction.journal_path.exists() + discard_codex_promotion(promotion, target) + discard_codex_stage(stage, target) + + +def test_provider_maintenance_marker_blocks_new_lease( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, "codex-1") + with ( + provider_enrollment_lock( + registry.settings.state_dir, + "codex", + registry.settings.lock_stale_seconds, + ), + pytest.raises(ValueError, match="provider maintenance is in progress"), + ): + select_and_acquire( + registry, + task="blocked-by-maintenance", + pool="codex-crew", + profile_id="codex-1", + ) + + +def test_profile_add_obeys_provider_maintenance_lock( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = load_registry(config) + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + with provider_enrollment_lock( + registry.settings.state_dir, + "claude", + registry.settings.lock_stale_seconds, + ): + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "profile", + "add", + "claude-4", + "--provider", + "claude", + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert "provider maintenance is already in progress" in json.loads(result.stdout)["error"] + assert "claude-4" not in load_registry(config).profiles + + +def test_lease_creation_and_binding_require_process_start_tokens( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + registry = load_registry(config) + monkeypatch.setattr(leases, "process_start_token", lambda _pid: None) + with pytest.raises(ValueError, match="verified process start token"): + new_lease("no-start-token", "codex-1", "codex-crew", pid=os.getpid()) + with pytest.raises(ValueError, match="verified process start token"): + bind_lease( + registry, + new_lease("reserved", "codex-1", "codex-crew", pid=None), + os.getpid(), + ) + + +def test_lock_creation_requires_process_start_token( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + registry = load_registry(config) + monkeypatch.setattr(locks, "process_start_token", lambda _pid: None) + with pytest.raises(RuntimeError, match="verified process start token"): + provider_enrollment_lock( + registry.settings.state_dir, + "codex", + registry.settings.lock_stale_seconds, + ) + + +@pytest.mark.parametrize("profile_id", ["claude-1", "codex-1"]) +def test_enrollment_refuses_while_managed_provider_launch_is_alive( + fleet: tuple[object, Path], profile_id: str +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, profile_id) + profile = registry.require_profile(profile_id) + select_and_acquire( + registry, + task=f"live-{profile.provider}-worker", + pool=f"{profile.provider}-crew", + profile_id=profile.id, + bind_pid=os.getpid(), + ) + profiles = dict(registry.profiles) + profiles[profile.id] = replace(profile, enabled=False) + save_registry(replace(registry, profiles=profiles), config) + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "profile", + "enroll", + profile.id, + ], + cwd=project_root, + env=env, + text=True, + capture_output=True, + check=False, + ) + assert result.returncode == 2 + assert "while any same-provider Fleet lease is active" in json.loads(result.stdout)["error"] + + +def test_claude_desktop_switch_is_seen_on_the_next_selection( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, "claude-1") + selected = select_and_acquire( + registry, + task="before-desktop-switch", + pool="claude-crew", + profile_id="claude-1", + dry_run=True, + ) + assert selected["profile"] == "claude-1" + desktop_file = registry.require_provider("claude").desktop_identity_file + assert desktop_file is not None + desktop_file.write_text( + json.dumps({"lastKnownAccountUuid": "claude-1-account"}), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="duplicate_provider_identity"): + select_and_acquire( + registry, + task="after-desktop-switch", + pool="claude-crew", + profile_id="claude-1", + dry_run=True, + ) + + +def test_missing_configured_desktop_anchor_fails_closed( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = _enable_and_provision(load_registry(config), config, "claude-1") + desktop_file = registry.require_provider("claude").desktop_identity_file + assert desktop_file is not None + desktop_file.unlink() + with pytest.raises(ValueError, match="duplicate_provider_identity"): + select_and_acquire( + registry, + task="missing-desktop-anchor", + pool="claude-crew", + profile_id="claude-1", + dry_run=True, + ) + + +def test_claude_worker_environment_blocks_login_logout_and_scrubs_ambient( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + profile = load_registry(config).require_profile("claude-1") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "poison") + monkeypatch.setenv("ANTHROPIC_API_KEY", "poison") + monkeypatch.setenv("CODEX_HOME", "/tmp/poison") + environment = provider_environment(profile, "managed-task") + assert environment["CLAUDE_CONFIG_DIR"] == str(profile.home) + assert environment["DISABLE_LOGIN_COMMAND"] == "1" + assert environment["DISABLE_LOGOUT_COMMAND"] == "1" + assert "CLAUDE_CODE_OAUTH_TOKEN" not in environment + assert "ANTHROPIC_API_KEY" not in environment + assert "CODEX_HOME" not in environment diff --git a/tools/agent-fleet/uv.lock b/tools/agent-fleet/uv.lock new file mode 100644 index 00000000000..88c5f2b310a --- /dev/null +++ b/tools/agent-fleet/uv.lock @@ -0,0 +1,108 @@ +version = 1 +revision = 3 +requires-python = ">=3.11" + +[[package]] +name = "agent-fleet" +version = "0.2.0" +source = { editable = "." } + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, + { name = "ruff" }, +] + +[package.metadata] + +[package.metadata.requires-dev] +dev = [ + { name = "pytest", specifier = ">=8.3" }, + { name = "ruff", specifier = ">=0.12" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "ruff" +version = "0.15.21" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/36/6f65aa9989acdec45d417192d8f4e7921931d8a6cf87ac74bce3eed98a8e/ruff-0.15.21.tar.gz", hash = "sha256:d0cfc841c572283c36548f82664a54ce6565567f1b0d5b4cf2caac693d8b7500", size = 4769401, upload-time = "2026-07-09T20:01:34.005Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d0/c6/ede15cac6839f3dbce52565c8f5164a8210e669c7bc4decb03e5bdf47d0d/ruff-0.15.21-py3-none-linux_armv6l.whl", hash = "sha256:63ea0e965e5d73c90e95b2434beeafc70820536717f561b32ab6e777cb9bdf5d", size = 10854342, upload-time = "2026-07-09T20:00:53.998Z" }, + { url = "https://files.pythonhosted.org/packages/28/9d/d825b07ee7ea9e2d61df92a860033c94e06e7300d50a1c2653aac27d24fe/ruff-0.15.21-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:0f212c5d7d54c01bbfe6dcab02b724a39300f3e34ed7acbe995ccb320a2c58bd", size = 11139539, upload-time = "2026-07-09T20:00:57.809Z" }, + { url = "https://files.pythonhosted.org/packages/f5/de/3b107712e642f063c7a9e0887c427b22cb44097de5aab36c05f2e280670c/ruff-0.15.21-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e6312e41bc96791299614995ea3a977c5857c3b5662b1ecef6755b02b87cb646", size = 10595437, upload-time = "2026-07-09T20:01:00.006Z" }, + { url = "https://files.pythonhosted.org/packages/9a/6f/b4523cc90ba239ede441447a19d0c968846a3012e5a0b0c5b62831a3d5e3/ruff-0.15.21-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:01d65b4831c6b2a4ba8ee6faa84049d44d982b7a706e622c4094c509e51673be", size = 10990053, upload-time = "2026-07-09T20:01:02.187Z" }, + { url = "https://files.pythonhosted.org/packages/92/cc/c6a9872a5375f0628875481cf2f66b13d7d865bf3ca2e57f91c7e762d976/ruff-0.15.21-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2c5a913a589120ce67933d5d05fd6ddbcc2481c6a054980ee767f7414c72b4fd", size = 10666096, upload-time = "2026-07-09T20:01:04.299Z" }, + { url = "https://files.pythonhosted.org/packages/ab/97/c621f7a17e097f1790fa3af6374138823b330b2d03fc38337945daca212c/ruff-0.15.21-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5ef04b681d02ad4dc9620f00f83ac5c22f652d0e9a9cfe431d219b16ad5ccc41", size = 11537011, upload-time = "2026-07-09T20:01:06.771Z" }, + { url = "https://files.pythonhosted.org/packages/ea/51/d928727e476e25ccc57c6f449ffd80241a651a973ad949d39cfb2a771d28/ruff-0.15.21-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:16d090c0740916594157e75b80d666eab8e78083b39b3b0e1d698f4670a17b86", size = 12347101, upload-time = "2026-07-09T20:01:08.859Z" }, + { url = "https://files.pythonhosted.org/packages/1e/88/8cd62026802b16018ad06931d87997cf795ba2a6239ab659606c87d96bf0/ruff-0.15.21-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3a10e74757dd65004d779b73e2f3c5210156d9980b41224d50d2ebcf1db51e67", size = 11572001, upload-time = "2026-07-09T20:01:11.092Z" }, + { url = "https://files.pythonhosted.org/packages/b2/97/f63084cf55444fc110e8cb985ebfcc592af47f597d44453d778cb81bc156/ruff-0.15.21-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bab0905d2f29e0d9fbc3c373ed23db0095edaa3f71f1f4f519ec15134d9e85c8", size = 11549239, upload-time = "2026-07-09T20:01:13.27Z" }, + { url = "https://files.pythonhosted.org/packages/9d/77/f107da4a2874b7715914b03f09ba9c54424de3ff8a1cc5d015d3ee2ce0ac/ruff-0.15.21-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:00eca240af5789fec6fe7df74c088cc1f9644ed83027113468efba7c92b94075", size = 11535340, upload-time = "2026-07-09T20:01:15.206Z" }, + { url = "https://files.pythonhosted.org/packages/d5/e9/601deb322d3303a7bf212b0100ead6f2ee3f6a044d89c30f2f92bf83c731/ruff-0.15.21-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:262ab31557a75141325e32d3357f3597645a7f084e732b6b054dde428ecd9341", size = 10964048, upload-time = "2026-07-09T20:01:17.723Z" }, + { url = "https://files.pythonhosted.org/packages/ea/2e/0f2176d1e99c15192caea19c8c3a0a955246b4cb4de795042eeb616345cd/ruff-0.15.21-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:659c4e7a4212f83306045ec7c5e5a356d16d9a6ef4ae0c7a4d872914fc655d9d", size = 10667055, upload-time = "2026-07-09T20:01:19.73Z" }, + { url = "https://files.pythonhosted.org/packages/48/60/abd74a02e0c4214f12a68becfd30af7165cfdcb0e661ecdc60bbb949c09a/ruff-0.15.21-py3-none-musllinux_1_2_i686.whl", hash = "sha256:9e866eab611a5f959d36df2d10e446973a3610bc42b0c15b31dc27977d59c233", size = 11242043, upload-time = "2026-07-09T20:01:21.947Z" }, + { url = "https://files.pythonhosted.org/packages/b2/c6/583075d8ccabb4b229345edcaf1545eb3d8d6be90f686a479d7e94088bbf/ruff-0.15.21-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:e89bc93c0d3803ba870b55c29671bad9dc6d94bb1eb181b056b52eb05b52854f", size = 11648064, upload-time = "2026-07-09T20:01:24.023Z" }, + { url = "https://files.pythonhosted.org/packages/3a/3c/37d0ecb729a7cc2d393ea7dce316fc585680f35d93b8d62139d7d0a3700c/ruff-0.15.21-py3-none-win32.whl", hash = "sha256:01f8d5be84823c172b389e123174f781f9daf86d6c58719d603f941932195cdd", size = 10896555, upload-time = "2026-07-09T20:01:26.941Z" }, + { url = "https://files.pythonhosted.org/packages/c0/b8/e43466b2a6067ce91e669068f6e28d6c719a920f014b070d5c8731725de3/ruff-0.15.21-py3-none-win_amd64.whl", hash = "sha256:d4b8d9a2f0f12b816b50447f6eccb9f4bb01a6b82c86b50fb3b5354b458dc6d3", size = 12038772, upload-time = "2026-07-09T20:01:29.497Z" }, + { url = "https://files.pythonhosted.org/packages/dd/75/e90ab9aeece218a9fc5a5bc3ec97d0ee6bb3c4ff95869463c1de58e29a1c/ruff-0.15.21-py3-none-win_arm64.whl", hash = "sha256:6e83115d4b9377c1cbc13abf0e051f069fab0ef815ea0504a8a008cee24dd0a8", size = 11375265, upload-time = "2026-07-09T20:01:31.772Z" }, +] From 89e2c59f659608b96c633983a216b0c5af39064e Mon Sep 17 00:00:00 2001 From: Dongkeun Lee Date: Fri, 17 Jul 2026 23:51:23 -0400 Subject: [PATCH 02/54] no-mistakes(review): Harden Agent Fleet trust and recovery --- tools/agent-fleet/README.md | 24 +- tools/agent-fleet/src/agent_fleet/cli.py | 119 +++++-- tools/agent-fleet/src/agent_fleet/config.py | 32 ++ tools/agent-fleet/src/agent_fleet/doctor.py | 53 +++- .../agent-fleet/src/agent_fleet/enrollment.py | 138 +++++--- tools/agent-fleet/src/agent_fleet/identity.py | 11 +- tools/agent-fleet/src/agent_fleet/models.py | 5 + tools/agent-fleet/src/agent_fleet/output.py | 17 + tools/agent-fleet/src/agent_fleet/projects.py | 105 ++++++ .../agent-fleet/src/agent_fleet/providers.py | 89 +++++- .../agent-fleet/src/agent_fleet/provision.py | 300 +++++++++++++++++- tools/agent-fleet/src/agent_fleet/quota.py | 17 + .../agent-fleet/src/agent_fleet/scheduler.py | 11 +- tools/agent-fleet/tests/conftest.py | 11 + .../tests/test_config_and_provision.py | 21 ++ .../agent-fleet/tests/test_contract_status.py | 101 +++++- tools/agent-fleet/tests/test_exec_sessions.py | 90 +++++- .../tests/test_project_bootstrap.py | 188 +++++++++++ .../agent-fleet/tests/test_quota_scheduler.py | 4 +- .../tests/test_safety_transactions.py | 59 +++- 20 files changed, 1286 insertions(+), 109 deletions(-) create mode 100644 tools/agent-fleet/src/agent_fleet/projects.py create mode 100644 tools/agent-fleet/tests/test_project_bootstrap.py diff --git a/tools/agent-fleet/README.md b/tools/agent-fleet/README.md index c53bbbcd01b..c94c67ce67e 100644 --- a/tools/agent-fleet/README.md +++ b/tools/agent-fleet/README.md @@ -50,11 +50,17 @@ task-to-provider-session mappings and resumes through the original profile. ## Initial setup ```sh -agent-fleet init --claude 3 --codex 2 +agent-fleet init --claude 1 --codex 1 +agent-fleet project register --provider claude /absolute/path/to/project +agent-fleet project register --provider codex /absolute/path/to/project agent-fleet profile enroll claude-1 agent-fleet profile enroll codex-1 ``` +Register every Git project that may host a managed worker before enrollment and launch. +Registration stores the canonical worktree root, while launch authorization compares Git common directories so linked Treehouse worktrees remain eligible. +Broad, symlinked, unrelated, and unregistered working directories fail closed before a provider process or lease starts. + `profile enroll` is the login transaction: the profile must already be disabled and every same-provider Fleet lease must be drained. It provisions the isolated home and hooks, runs provider login, verifies the credential against the live @@ -64,7 +70,9 @@ after reviewing verification: ```sh agent-fleet profile verify claude-1 --allow-keychain-prompt +agent-fleet profile verify codex-1 agent-fleet profile enable claude-1 +agent-fleet profile enable codex-1 ``` Codex enrollment uses device authorization by default and a fresh staging home, @@ -138,6 +146,10 @@ while a live worker lease requires an explicit forced release after the worker has stopped. Worker `exec` and `resume` require a managed task id, and every live lease and lock requires a verified process-start token; missing tokens fail closed. A live task can never be rebound to a different process. +Raw provider auth and resume subcommands are refused by `exec` so credential maintenance and sticky session recovery cannot bypass their dedicated paths. + +Claude launch bootstrap atomically preserves opaque profile state while setting only completed onboarding and trust for the validated active and canonical project roots. +Codex launches require an exact current-version profile hook set derived from the declared hook source plus Agent Fleet SessionStart, reject project hooks, disable `plugins` and `plugin_sharing`, and apply trust only to the validated active root. Task resume is fail-closed: it requires the recorded provider session and reuses that session's exact profile. The new-task quota reserve does not block @@ -184,10 +196,8 @@ routing activation are deliberately separate phases. On macOS, **Always Allow** so later automatic checks stay non-interactive. Bare `enable` never requests Keychain access and fails closed when verification is unavailable. -`doctor` verifies private profile homes, the Agent Fleet SessionStart hook, -current Herdr session-identity hooks, inherited workflow hooks/assets, provider -auth state, pinned binaries, and (when `--workspace` is supplied) that both -Claude and Codex supervision hooks expose `PreToolUse` and `Stop` there. +`doctor` verifies private profile homes, fresh remote identity proof, the Agent Fleet SessionStart hook, current Herdr session-identity hooks, inherited workflow hooks/assets, provider auth state, trusted-project configuration, and pinned binaries. +When `--workspace` is supplied, it also verifies provider onboarding/project/hook readiness, Claude supervision hooks, and the required absence of Codex project hooks. ## Isolation and shared workflow assets @@ -196,8 +206,8 @@ force file-backed credentials inside that home. Homes and state are mode 0700; registry, lease, quota, and session files are mode 0600. Ambient provider API key and access-token variables are removed before every provider launch. -The provider registry can declare a base home, hook source, and simple shared -entries. Initial profiles share only account-neutral workflow assets: +The provider registry can declare a base home, hook source, trusted Git projects, and allowlisted shared entries. +Initial profiles share only account-neutral workflow assets: - Claude: `CLAUDE.md`, skills, plugins, and hook definitions. - Codex: `AGENTS.md`, skills, plugins, rules, and hook definitions. diff --git a/tools/agent-fleet/src/agent_fleet/cli.py b/tools/agent-fleet/src/agent_fleet/cli.py index f263e958a5e..f49f885592f 100644 --- a/tools/agent-fleet/src/agent_fleet/cli.py +++ b/tools/agent-fleet/src/agent_fleet/cli.py @@ -6,6 +6,7 @@ import sys from collections.abc import Callable, Iterator from contextlib import ExitStack, contextmanager +from dataclasses import replace from pathlib import Path from typing import Any @@ -18,6 +19,7 @@ set_profile_enabled, set_profile_safety_policy, with_profile, + with_provider, without_profile, ) from .cooldowns import clear_cooldown, set_cooldown @@ -42,18 +44,22 @@ from .leases import active_leases, release_lease from .locks import provider_enrollment_lock, state_lock from .models import PROFILE_SAFETY_POLICIES, SUPPORTED_PROVIDERS, Profile, Registry -from .output import emit +from .output import emit, preflight from .paths import default_config_path, expand_path +from .projects import invocation_workspace, lexical_path, register_trusted_project from .providers import ( auth_probe, auth_status, login_argv, - provider_argv, + managed_argv, provider_environment, resume_argv, + validate_worker_arguments, ) -from .provision import profile_is_provisioned, provision_profile +from .provision import prepare_profile_launch, profile_is_provisioned, provision_profile from .quota import ( + discard_quota_cache, + has_remote_identity_proof, probe_quota, quota_routeability, read_quota, @@ -81,7 +87,15 @@ def _parser() -> argparse.ArgumentParser: init = commands.add_parser("init", help="create a disabled profile registry") init.add_argument("--claude", type=int, default=0) init.add_argument("--codex", type=int, default=0) - init.add_argument("--force", action="store_true") + + project = commands.add_parser("project") + project_commands = project.add_subparsers(dest="project_command", required=True) + project_list = project_commands.add_parser("list") + project_list.add_argument("--provider", choices=SUPPORTED_PROVIDERS) + for name in ("register", "remove"): + item = project_commands.add_parser(name) + item.add_argument("path", type=Path) + item.add_argument("--provider", choices=SUPPORTED_PROVIDERS, required=True) profile = commands.add_parser("profile") profile_commands = profile.add_subparsers(dest="profile_command", required=True) @@ -308,21 +322,15 @@ def _contract() -> dict[str, Any]: "session_remove": "session remove --task ", "profile_enroll": "profile enroll ", "profile_verify": "profile verify |--all", + "project_register": "project register --provider ", + "project_remove": "project remove --provider ", + "project_list": "project list [--provider ]", }, } def _credential_is_remotely_verified(quota: dict[str, Any]) -> bool: - fingerprint = quota.get("identity_fingerprint") - return ( - quota.get("status") == "fresh" - and quota.get("verified_at") is not None - and quota.get("headroom_percent") is not None - and isinstance(quota.get("windows"), list) - and bool(quota["windows"]) - and isinstance(fingerprint, str) - and len(fingerprint) == 64 - ) + return has_remote_identity_proof(quota) def _cached_credential_proof_is_usable(quota: dict[str, Any]) -> bool: @@ -466,7 +474,7 @@ def _enroll_codex_profile( browser_login: bool, access_token: bool, ) -> dict[str, Any]: - stage = create_codex_login_stage(target) + stage = create_codex_login_stage(registry, target) promotion: Profile | None = None transaction: CodexAuthTransaction | None = None quota_snapshot = snapshot_quota_cache(registry, target.id) @@ -543,7 +551,7 @@ def _enroll_codex_profile( if promotion is not None and promotion.home.exists(): discard_codex_promotion(promotion, target) if stage.home.exists(): - discard_codex_stage(stage, target) + discard_codex_stage(registry, stage, target) def _run_profile_enrollment( @@ -623,6 +631,7 @@ def _run_profile_enrollment( ) else: provision_profile(registry, profile) + discard_quota_cache(registry, profile.id) completed = subprocess.run( login_argv(registry, profile), env=provider_environment(profile), @@ -689,7 +698,7 @@ def _run(args: argparse.Namespace) -> Any | None: if args.command == "contract": return _contract() if args.command == "init": - if config_path.exists() and not args.force: + if config_path.exists(): raise ValueError(f"registry already exists: {config_path}") registry = initial_registry(args.claude, args.codex) save_registry(registry, config_path) @@ -700,6 +709,48 @@ def _run(args: argparse.Namespace) -> Any | None: } registry = load_registry(config_path) + if args.command == "project": + if args.project_command == "list": + providers = [args.provider] if args.provider else list(SUPPORTED_PROVIDERS) + return { + "providers": [ + { + "provider": provider, + "trusted_projects": [ + str(path) + for path in registry.require_provider(provider).trusted_projects + ], + } + for provider in providers + ] + } + project_root = register_trusted_project(args.path) + with _provider_maintenance(registry, config_path, {args.provider}) as current: + if _provider_has_active_lease(current, args.provider): + raise ValueError( + f"refusing {args.provider} project maintenance while a Fleet lease is active" + ) + provider = current.require_provider(args.provider) + projects = set(provider.trusted_projects) + if args.project_command == "register": + projects.add(project_root) + else: + projects.discard(project_root) + updated = _mutate( + current, + lambda item: with_provider( + item, + replace(provider, trusted_projects=tuple(sorted(projects, key=str))), + ), + config_path, + ) + return { + "provider": args.provider, + "trusted_projects": [ + str(path) for path in updated.require_provider(args.provider).trusted_projects + ], + } + if args.command == "profile": if args.profile_command == "list": return { @@ -973,6 +1024,26 @@ def _run(args: argparse.Namespace) -> Any | None: ) if args.command == "exec": + provider_args = _strip_separator(args.provider_args) + if args.profile: + validate_worker_arguments(registry.require_profile(args.profile), provider_args) + elif args.provider: + for profile in registry.profiles.values(): + if profile.provider == args.provider: + validate_worker_arguments(profile, provider_args) + break + else: + for provider in SUPPORTED_PROVIDERS: + candidate = next( + ( + profile + for profile in registry.profiles.values() + if profile.provider == provider + ), + None, + ) + if candidate is not None: + validate_worker_arguments(candidate, provider_args) task = _task(args.task) if args.task else None if task is None: raise ValueError("worker exec requires --task so its provider lease is tracked") @@ -989,15 +1060,19 @@ def _run(args: argparse.Namespace) -> Any | None: explicit_profile=pool == "explicit", ) profile = registry.require_profile(str(selected["profile"])) - argv = provider_argv(registry, profile, _strip_separator(args.provider_args)) + validate_worker_arguments(profile, provider_args) + project = prepare_profile_launch(registry, profile, invocation_workspace()) + argv = managed_argv(registry, profile, project.active_root, provider_args) os.execvpe(argv[0], argv, provider_environment(profile, task)) if args.command == "resume": + provider_args = _strip_separator(args.provider_args) task = _task(args.task) if args.task else None if task is None: raise ValueError("worker resume requires --task so its provider lease is tracked") mapping = get_session(registry, task) profile = registry.require_profile(str(mapping.get("profile"))) + validate_worker_arguments(profile, provider_args) if args.profile and args.profile != profile.id: raise ValueError("explicit profile does not match task session mapping") session_id = mapping.get("session_id") @@ -1022,7 +1097,10 @@ def _run(args: argparse.Namespace) -> Any | None: registry, profile, session_id, - _strip_separator(args.provider_args), + provider_args, + active_root=prepare_profile_launch( + registry, profile, invocation_workspace() + ).active_root, ) os.execvpe(argv[0], argv, provider_environment(profile, task)) @@ -1043,7 +1121,7 @@ def _run(args: argparse.Namespace) -> Any | None: return run_doctor( registry, config_path, - workspace=expand_path(args.workspace) if args.workspace else None, + workspace=lexical_path(args.workspace) if args.workspace else None, ) if args.command == "status": with _provider_maintenance( @@ -1067,6 +1145,7 @@ def main(argv: list[str] | None = None) -> int: parser = _parser() args = parser.parse_args(argv) try: + preflight(args.format) payload = _run(args) if payload is not None: emit(payload, args.format) diff --git a/tools/agent-fleet/src/agent_fleet/config.py b/tools/agent-fleet/src/agent_fleet/config.py index 1b7e2d262d3..f516f222d5a 100644 --- a/tools/agent-fleet/src/agent_fleet/config.py +++ b/tools/agent-fleet/src/agent_fleet/config.py @@ -9,6 +9,7 @@ from .models import ( PROFILE_SAFETY_POLICIES, + SHARED_WORKFLOW_ENTRIES, SUPPORTED_PROVIDERS, Profile, ProviderConfig, @@ -149,6 +150,7 @@ def load_registry(path: Path | None = None) -> Registry: "~/Library/Application Support/Claude/config.json" if provider == "claude" else None, ) shared_raw = item.get("shared_entries", []) + trusted_projects_raw = item.get("trusted_projects", []) if base_home_raw is not None and not isinstance(base_home_raw, str): raise ValueError(f"providers.{provider}.base_home must be a path string") if hooks_source_raw is not None and not isinstance(hooks_source_raw, str): @@ -173,6 +175,16 @@ def load_registry(path: Path | None = None) -> Registry: for entry in shared_raw ): raise ValueError(f"providers.{provider}.shared_entries must contain simple file names") + disallowed_shared = sorted(set(shared_raw) - SHARED_WORKFLOW_ENTRIES[provider]) + if disallowed_shared: + raise ValueError( + f"providers.{provider}.shared_entries contains non-workflow assets: " + + ", ".join(disallowed_shared) + ) + if not isinstance(trusted_projects_raw, list) or not all( + isinstance(entry, str) and entry for entry in trusted_projects_raw + ): + raise ValueError(f"providers.{provider}.trusted_projects must contain path strings") providers[provider] = ProviderConfig( provider, expand_path(binary_raw), @@ -184,6 +196,7 @@ def load_registry(path: Path | None = None) -> Registry: if isinstance(desktop_identity_file_raw, str) and desktop_identity_file_raw else None ), + tuple(expand_path(entry) for entry in trusted_projects_raw), ) profiles_raw = raw.get("profiles", {}) @@ -214,6 +227,7 @@ def initial_registry(claude_count: int, codex_count: int) -> Registry: expand_path("~/.claude/settings.json"), ("CLAUDE.md", "skills", "plugins"), expand_path("~/Library/Application Support/Claude/config.json"), + (), ), "codex": ProviderConfig( "codex", @@ -226,6 +240,8 @@ def initial_registry(claude_count: int, codex_count: int) -> Registry: expand_path("~/.codex"), expand_path("~/.codex/hooks.json"), ("AGENTS.md", "skills", "plugins", "rules"), + None, + (), ), } profiles: dict[str, Profile] = {} @@ -264,6 +280,12 @@ def with_profile(registry: Registry, profile: Profile) -> Registry: return updated +def with_provider(registry: Registry, provider: ProviderConfig) -> Registry: + providers = dict(registry.providers) + providers[provider.name] = provider + return replace(registry, providers=providers) + + def without_profile(registry: Registry, profile_id: str) -> Registry: registry.require_profile(profile_id) profiles = dict(registry.profiles) @@ -276,6 +298,13 @@ def _paths_overlap(first: Path, second: Path) -> bool: def _validate_profile_invariants(registry: Registry) -> None: + for provider_name, provider in registry.providers.items(): + disallowed = sorted(set(provider.shared_entries) - SHARED_WORKFLOW_ENTRIES[provider_name]) + if disallowed: + raise ValueError( + f"providers.{provider_name}.shared_entries contains non-workflow assets: " + + ", ".join(disallowed) + ) profiles = list(registry.profiles.values()) for profile in profiles: if profile.safety_policy not in PROFILE_SAFETY_POLICIES: @@ -402,6 +431,9 @@ def save_registry(registry: Registry, path: Path | None = None) -> Path: "shared_entries = [" + ", ".join(_toml_string(entry) for entry in provider_config.shared_entries) + "]", + "trusted_projects = [" + + ", ".join(_toml_string(entry) for entry in provider_config.trusted_projects) + + "]", ] ) for profile_id in sorted(registry.profiles): diff --git a/tools/agent-fleet/src/agent_fleet/doctor.py b/tools/agent-fleet/src/agent_fleet/doctor.py index 1e79cdf3f51..174c47dbb37 100644 --- a/tools/agent-fleet/src/agent_fleet/doctor.py +++ b/tools/agent-fleet/src/agent_fleet/doctor.py @@ -8,12 +8,15 @@ from typing import Any from .models import Registry +from .projects import canonical_git_project from .providers import auth_status from .provision import ( profile_hook_health, profile_is_provisioned, + profile_launch_ready, profile_shared_assets_healthy, ) +from .quota import has_remote_identity_proof, read_quota def _mode(path: Path) -> str | None: @@ -62,6 +65,18 @@ def add(name: str, ok: bool, detail: str, *, required: bool = True) -> None: binary_ok, f"{config.binary}", ) + projects_ready = bool(config.trusted_projects) + if projects_ready: + try: + for project in config.trusted_projects: + canonical_git_project(project) + except ValueError: + projects_ready = False + add( + f"trusted-projects:{provider}", + projects_ready, + ",".join(str(path) for path in config.trusted_projects) or "none", + ) for profile in sorted(registry.profiles.values(), key=lambda item: item.id): provisioned = profile_is_provisioned(profile) add( @@ -75,6 +90,18 @@ def add(name: str, ok: bool, detail: str, *, required: bool = True) -> None: status == "authenticated" or not profile.enabled, status, ) + quota = read_quota(registry, profile.id) + remote_verified = ( + provisioned + and status == "authenticated" + and quota.get("fresh") is True + and has_remote_identity_proof(quota) + ) + add( + f"profile:{profile.id}:remote-identity-proof", + remote_verified, + "fresh" if remote_verified else str(quota.get("reason") or quota.get("status")), + ) if provisioned: home_mode = _mode(profile.home) add( @@ -96,15 +123,27 @@ def add(name: str, ok: bool, detail: str, *, required: bool = True) -> None: "healthy" if shared_healthy else "missing or redirected link", ) if workspace is not None: - workspace = workspace.resolve() claude_events = _workspace_hook_events(workspace / ".claude" / "settings.json") - codex_events = _workspace_hook_events(workspace / ".codex" / "hooks.json") - for provider, events in (("claude", claude_events), ("codex", codex_events)): - required = {"PreToolUse", "Stop"} + required = {"PreToolUse", "Stop"} + add( + "workspace:claude:supervision-hooks", + required.issubset(claude_events), + f"{workspace}: events={','.join(sorted(claude_events)) or 'none'}", + ) + codex_project_hooks = workspace / ".codex" / "hooks.json" + add( + "workspace:codex:project-hooks-absent", + not (codex_project_hooks.exists() or codex_project_hooks.is_symlink()), + str(codex_project_hooks), + ) + for profile in sorted(registry.profiles.values(), key=lambda item: item.id): + ready = profile_is_provisioned(profile) and profile_launch_ready( + registry, profile, workspace + ) add( - f"workspace:{provider}:supervision-hooks", - required.issubset(events), - f"{workspace}: events={','.join(sorted(events)) or 'none'}", + f"workspace:{profile.id}:provider-bootstrap", + ready, + "ready" if ready else "project, onboarding, or hook readiness failed", ) required_failures = [check for check in checks if check["required"] and not check["ok"]] return { diff --git a/tools/agent-fleet/src/agent_fleet/enrollment.py b/tools/agent-fleet/src/agent_fleet/enrollment.py index 49262ef6231..0c41cc0de77 100644 --- a/tools/agent-fleet/src/agent_fleet/enrollment.py +++ b/tools/agent-fleet/src/agent_fleet/enrollment.py @@ -174,13 +174,18 @@ def _copy_regular_file(source: Path, destination: Path) -> None: os.close(source_fd) -def create_codex_login_stage(profile: Profile) -> Profile: - ensure_private_dir(profile.home.parent) - _private_directory(profile.home.parent, "managed Codex accounts directory") +def _codex_stage_root(registry: Registry) -> Path: + return registry.settings.state_dir / "staging" / "codex" + + +def create_codex_login_stage(registry: Registry, profile: Profile) -> Profile: + stage_root = _codex_stage_root(registry) + ensure_private_dir(stage_root) + _private_directory(stage_root, "managed Codex staging directory") stage = Path( tempfile.mkdtemp( prefix=f".{profile.home.name}.login-", - dir=profile.home.parent, + dir=stage_root, ) ) stage.chmod(0o700) @@ -197,9 +202,11 @@ def _remove_private_tree(path: Path) -> None: shutil.rmtree(path) -def discard_codex_stage(stage: Profile, target: Profile) -> None: +def discard_codex_stage(registry: Registry, stage: Profile, target: Profile) -> None: expected_prefix = f".{target.home.name}.login-" - if stage.home.parent != target.home.parent or not stage.home.name.startswith(expected_prefix): + if stage.home.parent != _codex_stage_root(registry) or not stage.home.name.startswith( + expected_prefix + ): raise ValueError(f"refusing to discard unrecognized Codex stage: {stage.home}") _remove_private_tree(stage.home) @@ -378,24 +385,8 @@ def rollback_codex_promotion( transaction: CodexAuthTransaction, ) -> None: _private_directory(target.home, "managed Codex profile home") - _validate_transaction(registry, target, transaction) - backup = transaction.backup_auth - if backup is None: - transaction.target_auth.unlink() - else: - expected_prefix = f".{CODEX_AUTH_FILE}.backup-" - if backup.parent != target.home or not backup.name.startswith(expected_prefix): - raise ValueError("Codex auth backup is outside the target profile") - _regular_file_stat(backup, "Codex auth backup") - os.replace(backup, transaction.target_auth) - transaction.target_auth.chmod(0o600) - _fsync_directory(target.home) - journal = _read_journal(transaction.journal_path) - restore_quota_cache( - registry, - target.id, - _snapshot_from_payload(journal.get("quota_snapshot")), - ) + journal = _validate_transaction(registry, target, transaction) + _recover_rollback(registry, target, journal) transaction.journal_path.unlink() _fsync_directory(transaction.journal_path.parent) @@ -433,11 +424,16 @@ def recover_pending_codex_transaction(registry: Registry, target: Profile) -> bo journal = _read_journal(journal_path) target_auth, backup, temporary = _validate_journal_paths(registry, target, journal) installed_stat = _decoded_stat(journal.get("installed_stat"), "installed stat") - original_stat = _decoded_stat(journal.get("original_stat"), "original stat") if installed_stat is None: raise ValueError("Codex auth transaction has no installed stat") phase = journal.get("phase") - if phase not in {"prepared", "committed"}: + if phase not in { + "prepared", + "rolling_back", + "auth_restored", + "quota_restored", + "committed", + }: raise ValueError("Codex auth transaction has an invalid phase") _private_directory(target.home, "managed Codex profile home") @@ -451,31 +447,95 @@ def recover_pending_codex_transaction(registry: Registry, target: Profile) -> bo if current_stat != installed_stat: raise ValueError("committed Codex auth changed before crash recovery") else: + _recover_rollback(registry, target, journal) + + # Cleanup happens only after rollback+quota restore or a durable commit. + temporary.unlink(missing_ok=True) + if backup is not None: + backup.unlink(missing_ok=True) + journal_path.unlink() + _fsync_directory(target.home) + _fsync_directory(journal_path.parent) + return True + + +def _write_journal_phase(path: Path, journal: dict[str, Any], phase: str) -> None: + journal["phase"] = phase + atomic_write_json(path, journal) + _fsync_directory(path.parent) + + +def _recover_rollback( + registry: Registry, + target: Profile, + journal: dict[str, Any], +) -> None: + target_auth, backup, _ = _validate_journal_paths(registry, target, journal) + journal_path = _journal_path(registry, target) + installed_stat = _decoded_stat(journal.get("installed_stat"), "installed stat") + original_stat = _decoded_stat(journal.get("original_stat"), "original stat") + if installed_stat is None: + raise ValueError("Codex auth transaction has no installed stat") + phase = journal.get("phase") + current_stat = ( + _stat_identity(_regular_file_stat(target_auth, "managed Codex auth")) + if target_auth.exists() or target_auth.is_symlink() + else None + ) + if phase == "prepared": + if current_stat == installed_stat: + rollback_stat = ( + _stat_identity(_regular_file_stat(backup, "Codex auth backup")) + if backup is not None + else None + ) + elif current_stat == original_stat: + rollback_stat = original_stat + else: + raise ValueError("Codex auth changed outside the interrupted transaction") + journal["rollback_stat"] = _encoded_stat(rollback_stat) + _write_journal_phase(journal_path, journal, "rolling_back") + phase = "rolling_back" + rollback_stat = _decoded_stat(journal.get("rollback_stat"), "rollback stat") + if phase != "prepared" and "rollback_stat" not in journal: + raise ValueError("Codex auth transaction has no rollback stat") + if phase == "rolling_back": + current_stat = ( + _stat_identity(_regular_file_stat(target_auth, "managed Codex auth")) + if target_auth.exists() or target_auth.is_symlink() + else None + ) if current_stat == installed_stat: if backup is None: target_auth.unlink() else: - _regular_file_stat(backup, "Codex auth backup") + backup_stat = _stat_identity(_regular_file_stat(backup, "Codex auth backup")) + if backup_stat != rollback_stat: + raise ValueError("Codex auth backup changed during rollback") os.replace(backup, target_auth) target_auth.chmod(0o600) _fsync_directory(target.home) - elif current_stat != original_stat: - raise ValueError("Codex auth changed outside the interrupted transaction") - # Restore normalized quota evidence before deleting any recovery artifact. + elif current_stat != rollback_stat: + raise ValueError("Codex auth changed outside the interrupted rollback") + _write_journal_phase(journal_path, journal, "auth_restored") + phase = "auth_restored" + if phase == "auth_restored": + current_stat = ( + _stat_identity(_regular_file_stat(target_auth, "managed Codex auth")) + if target_auth.exists() or target_auth.is_symlink() + else None + ) + if current_stat != rollback_stat: + raise ValueError("restored Codex auth changed before quota recovery") restore_quota_cache( registry, target.id, _snapshot_from_payload(journal.get("quota_snapshot")), ) - - # Cleanup happens only after rollback+quota restore or a durable commit. - temporary.unlink(missing_ok=True) - if backup is not None: - backup.unlink(missing_ok=True) - journal_path.unlink() - _fsync_directory(target.home) - _fsync_directory(journal_path.parent) - return True + _write_journal_phase(journal_path, journal, "quota_restored") + phase = "quota_restored" + if phase != "quota_restored": + raise ValueError("Codex auth transaction has an invalid rollback phase") def recover_pending_codex_transactions(registry: Registry, provider: str) -> list[str]: diff --git a/tools/agent-fleet/src/agent_fleet/identity.py b/tools/agent-fleet/src/agent_fleet/identity.py index 3f3fe8891bc..eb5315f29f4 100644 --- a/tools/agent-fleet/src/agent_fleet/identity.py +++ b/tools/agent-fleet/src/agent_fleet/identity.py @@ -7,7 +7,7 @@ from .models import Profile, Registry from .providers import identity_fingerprint -from .quota import probe_quota, read_quota +from .quota import has_remote_identity_proof, probe_quota, read_quota from .util import atomic_write_json, utc_now @@ -43,14 +43,7 @@ def _read_anchor(registry: Registry, provider: str, kind: str) -> dict[str, Any] def _quota_identity_is_verified(quota: dict[str, Any]) -> bool: - fingerprint = quota.get("identity_fingerprint") - return ( - quota.get("status") == "fresh" - and quota.get("verified_at") is not None - and quota.get("headroom_percent") is not None - and isinstance(fingerprint, str) - and len(fingerprint) == 64 - ) + return has_remote_identity_proof(quota) def _managed_identity_has_recent_proof(quota: dict[str, Any]) -> bool: diff --git a/tools/agent-fleet/src/agent_fleet/models.py b/tools/agent-fleet/src/agent_fleet/models.py index 8742288cb2c..94429a5f514 100644 --- a/tools/agent-fleet/src/agent_fleet/models.py +++ b/tools/agent-fleet/src/agent_fleet/models.py @@ -6,6 +6,10 @@ SUPPORTED_PROVIDERS = ("claude", "codex") PROFILE_SAFETY_POLICIES = ("worker", "manual_only", "desktop_shared") +SHARED_WORKFLOW_ENTRIES = { + "claude": frozenset({"CLAUDE.md", "skills", "plugins"}), + "codex": frozenset({"AGENTS.md", "skills", "plugins", "rules"}), +} @dataclass(frozen=True) @@ -16,6 +20,7 @@ class ProviderConfig: hooks_source: Path | None = None shared_entries: tuple[str, ...] = () desktop_identity_file: Path | None = None + trusted_projects: tuple[Path, ...] = () @dataclass(frozen=True) diff --git a/tools/agent-fleet/src/agent_fleet/output.py b/tools/agent-fleet/src/agent_fleet/output.py index 3100a583b1f..28fb26c1297 100644 --- a/tools/agent-fleet/src/agent_fleet/output.py +++ b/tools/agent-fleet/src/agent_fleet/output.py @@ -7,6 +7,23 @@ from typing import Any +def preflight(output_format: str) -> None: + if output_format != "toon": + return + toon = shutil.which("toon") + if toon is None: + raise ValueError("TOON output requested but `toon` is not on PATH; use --format json") + result = subprocess.run( + [toon], + input="{}", + text=True, + capture_output=True, + check=False, + ) + if result.returncode != 0: + raise ValueError(f"TOON encoder preflight failed: {result.stderr.strip()}") + + def emit(payload: Any, output_format: str) -> None: if output_format == "json": json.dump(payload, sys.stdout, sort_keys=True, separators=(",", ":")) diff --git a/tools/agent-fleet/src/agent_fleet/projects.py b/tools/agent-fleet/src/agent_fleet/projects.py new file mode 100644 index 00000000000..850a05d3cb8 --- /dev/null +++ b/tools/agent-fleet/src/agent_fleet/projects.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +import os +import stat +import subprocess +from dataclasses import dataclass +from pathlib import Path + +from .models import Registry + + +@dataclass(frozen=True) +class TrustedProject: + active_root: Path + canonical_root: Path + common_dir: Path + + +def lexical_path(path: Path) -> Path: + return Path(os.path.abspath(os.path.expandvars(os.path.expanduser(str(path))))) + + +def invocation_workspace() -> Path: + current = Path.cwd() + raw = os.environ.get("PWD") + if not raw: + return current + candidate = Path(raw) + try: + if candidate.resolve() == current.resolve(): + return candidate + except OSError: + pass + return current + + +def _owned_directory(path: Path, label: str) -> None: + try: + current = path.lstat() + except FileNotFoundError as exc: + raise ValueError(f"{label} is missing: {path}") from exc + if not stat.S_ISDIR(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"{label} must be a current-user directory: {path}") + + +def _git_path(path: Path, argument: str) -> Path: + environment = {name: value for name, value in os.environ.items() if not name.startswith("GIT_")} + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "--path-format=absolute", argument], + env=environment, + capture_output=True, + text=True, + check=False, + ) + except OSError as exc: + raise ValueError("Git is required to validate trusted projects") from exc + value = result.stdout.strip() + if result.returncode != 0 or not value: + raise ValueError(f"trusted project must be inside a Git worktree: {path}") + return Path(value) + + +def canonical_git_project(path: Path) -> tuple[Path, Path]: + expanded = lexical_path(path) + if expanded.is_symlink() or expanded.resolve() != expanded: + raise ValueError(f"trusted project path must not be symlinked: {path}") + _owned_directory(expanded, "trusted project path") + root = _git_path(expanded, "--show-toplevel") + common_dir = _git_path(expanded, "--git-common-dir") + if root.is_symlink() or root.resolve() != root: + raise ValueError(f"Git worktree root must not be symlinked: {root}") + _owned_directory(root, "Git worktree root") + _owned_directory(common_dir, "Git common directory") + if root == Path(root.anchor) or root == Path.home().resolve(): + raise ValueError(f"trusted project root is too broad: {root}") + try: + expanded.relative_to(root) + except ValueError as exc: + raise ValueError(f"trusted project path escaped its Git worktree: {path}") from exc + return root, common_dir.resolve() + + +def register_trusted_project(path: Path) -> Path: + root, _ = canonical_git_project(path) + return root + + +def resolve_trusted_project(registry: Registry, provider: str, workspace: Path) -> TrustedProject: + active_root, active_common_dir = canonical_git_project(workspace) + configured = registry.require_provider(provider).trusted_projects + matches: list[Path] = [] + for configured_path in configured: + try: + canonical_root, common_dir = canonical_git_project(configured_path) + except ValueError: + continue + if common_dir == active_common_dir: + matches.append(canonical_root) + if not matches: + raise ValueError( + f"workspace is not registered for {provider}: {active_root}; " + f"run `agent-fleet project register --provider {provider} {active_root}`" + ) + return TrustedProject(active_root, sorted(set(matches), key=str)[0], active_common_dir) diff --git a/tools/agent-fleet/src/agent_fleet/providers.py b/tools/agent-fleet/src/agent_fleet/providers.py index b1c821d28ba..1d9dbf3cc57 100644 --- a/tools/agent-fleet/src/agent_fleet/providers.py +++ b/tools/agent-fleet/src/agent_fleet/providers.py @@ -1,11 +1,13 @@ from __future__ import annotations +import json import os import shlex import shutil import subprocess from collections.abc import Iterable from hashlib import sha256 +from pathlib import Path from .models import Profile, Registry @@ -41,6 +43,85 @@ def provider_argv(registry: Registry, profile: Profile, command: Iterable[str] = return [str(binary), *suffix] +def validate_worker_arguments(profile: Profile, arguments: list[str]) -> None: + blocked = { + "login", + "logout", + "resume", + "fork", + "auth", + "-r", + "--resume", + "--continue", + } + if any( + argument in blocked + or argument.startswith("--resume=") + or argument.startswith("--continue=") + for argument in arguments + ): + raise ValueError("worker exec refuses provider auth and resume commands") + if any( + argument in {"-C", "--cd", "--cwd", "--directory", "--add-dir"} + or argument.startswith(("-C", "--cd=", "--cwd=", "--directory=", "--add-dir=")) + for argument in arguments + ): + raise ValueError("worker exec refuses provider working-directory overrides") + if profile.provider != "codex": + return + if any(argument in {"plugin", "features", "mcp"} for argument in arguments): + raise ValueError("worker exec refuses Codex plugin and configuration administration") + for index, argument in enumerate(arguments): + if argument in {"-p", "--profile", "--remote"} or argument.startswith( + ("-p=", "--profile=", "--remote=") + ): + raise ValueError("worker exec refuses alternate Codex config and runtime profiles") + if argument == "--dangerously-bypass-hook-trust": + raise ValueError("worker exec owns the Codex hook-trust override") + if ( + argument == "--enable" + and index + 1 < len(arguments) + and arguments[index + 1] in {"plugins", "plugin_sharing"} + ): + raise ValueError("managed Codex launches keep plugins disabled") + if argument.startswith("--enable=") and argument.split("=", 1)[1] in { + "plugins", + "plugin_sharing", + }: + raise ValueError("managed Codex launches keep plugins disabled") + if argument in {"-c", "--config"} and index + 1 < len(arguments): + value = arguments[index + 1].lower() + if any(token in value for token in ("project", "trust", "hook", "plugin")): + raise ValueError("worker exec refuses managed Codex config overrides") + if argument.startswith("--config=") and any( + token in argument.lower() for token in ("project", "trust", "hook", "plugin") + ): + raise ValueError("worker exec refuses managed Codex config overrides") + + +def codex_launch_prefix(active_root: Path) -> list[str]: + trust_override = f'projects.{json.dumps(str(active_root))}.trust_level="trusted"' + return [ + "--disable", + "plugins", + "--disable", + "plugin_sharing", + "-c", + trust_override, + "--dangerously-bypass-hook-trust", + ] + + +def managed_argv( + registry: Registry, + profile: Profile, + active_root: Path, + extra: list[str], +) -> list[str]: + prefix = codex_launch_prefix(active_root) if profile.provider == "codex" else [] + return provider_argv(registry, profile, [*prefix, *extra]) + + def login_argv( registry: Registry, profile: Profile, @@ -104,7 +185,13 @@ def resume_argv( profile: Profile, session_id: str, extra: list[str], + *, + active_root: Path, ) -> list[str]: if profile.provider == "claude": return provider_argv(registry, profile, ["--resume", session_id, *extra]) - return provider_argv(registry, profile, ["resume", session_id, *extra]) + return provider_argv( + registry, + profile, + ["resume", *codex_launch_prefix(active_root), session_id, *extra], + ) diff --git a/tools/agent-fleet/src/agent_fleet/provision.py b/tools/agent-fleet/src/agent_fleet/provision.py index d61943e6fd5..44800389583 100644 --- a/tools/agent-fleet/src/agent_fleet/provision.py +++ b/tools/agent-fleet/src/agent_fleet/provision.py @@ -1,5 +1,7 @@ from __future__ import annotations +import copy +import hashlib import json import os import stat @@ -7,12 +9,15 @@ from pathlib import Path from typing import Any -from .models import Profile, ProviderConfig, Registry +from . import __version__ +from .models import SHARED_WORKFLOW_ENTRIES, Profile, ProviderConfig, Registry from .paths import ensure_private_dir +from .projects import TrustedProject, resolve_trusted_project from .providers import session_hook_command from .util import atomic_write_json HOOK_MARKER = " hook session-start" +CODEX_HOOK_MARKER_FILE = ".agent-fleet-hooks.json" def _read_json_object(path: Path) -> dict[str, Any]: @@ -27,6 +32,20 @@ def _read_json_object(path: Path) -> dict[str, Any]: return value +def _read_owned_json_object(path: Path) -> dict[str, Any]: + if not path.exists() and not path.is_symlink(): + return {} + try: + current = path.lstat() + except FileNotFoundError: + return {} + if not stat.S_ISREG(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"managed JSON must be a current-user regular file: {path}") + if stat.S_IMODE(current.st_mode) & 0o077: + raise ValueError(f"managed JSON must not grant group/world access: {path}") + return _read_json_object(path) + + def _merge_source_hooks(payload: dict[str, Any], source: Path | None) -> None: if source is None or not source.exists(): return @@ -82,15 +101,227 @@ def _install_session_hook(path: Path, source: Path | None) -> None: atomic_write_json(path, payload) +def _source_hash(source: Path | None) -> str | None: + if source is None: + return None + if not source.exists() and not source.is_symlink(): + return None + current = source.lstat() + if not stat.S_ISREG(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"hook source must be a current-user regular file: {source}") + return hashlib.sha256(source.read_bytes()).hexdigest() + + +def _hook_payload_hash(payload: dict[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _agent_fleet_session_group(command: str) -> dict[str, Any]: + return { + "matcher": "startup|resume|clear|compact", + "hooks": [ + { + "type": "command", + "command": command, + "statusMessage": "Recording Agent Fleet session identity", + } + ], + } + + +def _codex_hook_payload(source: Path | None, command: str) -> dict[str, Any]: + source_payload = _read_json_object(source) if source is not None and source.exists() else {} + source_hooks = source_payload.get("hooks", {}) + if not isinstance(source_hooks, dict): + raise ValueError(f"hooks must be an object: {source}") + hooks = copy.deepcopy(source_hooks) + for event, groups in list(hooks.items()): + if not isinstance(groups, list): + raise ValueError(f"hooks.{event} must be an array: {source}") + canonical_groups: list[Any] = [] + for group in groups: + if not isinstance(group, dict): + raise ValueError(f"hooks.{event} entries must be objects: {source}") + entries = group.get("hooks", []) + if not isinstance(entries, list): + raise ValueError(f"hooks.{event} hook entries must be an array: {source}") + canonical_entries = [ + entry + for entry in entries + if not (isinstance(entry, dict) and HOOK_MARKER in str(entry.get("command", ""))) + ] + if canonical_entries: + group["hooks"] = canonical_entries + canonical_groups.append(group) + hooks[event] = canonical_groups + groups = hooks.setdefault("SessionStart", []) + groups.append(_agent_fleet_session_group(command)) + return {"hooks": hooks} + + +def _install_codex_hooks(profile: Profile, provider: ProviderConfig) -> None: + command = session_hook_command() + source_hash = _source_hash(provider.hooks_source) + payload = _codex_hook_payload(provider.hooks_source, command) + if source_hash != _source_hash(provider.hooks_source): + raise ValueError("Codex hook source changed during provisioning") + path = profile.home / "hooks.json" + atomic_write_json(path, payload) + atomic_write_json( + profile.home / CODEX_HOOK_MARKER_FILE, + { + "schema": 1, + "agent_fleet_version": __version__, + "profile": profile.id, + "provider": profile.provider, + "source": str(provider.hooks_source) if provider.hooks_source is not None else None, + "source_hash": source_hash, + "session_command": command, + "hooks_hash": _hook_payload_hash(payload), + }, + ) + + +def codex_hooks_ready(registry: Registry, profile: Profile) -> bool: + provider = registry.require_provider(profile.provider) + try: + marker = _read_owned_json_object(profile.home / CODEX_HOOK_MARKER_FILE) + if ( + set(marker) + != { + "schema", + "agent_fleet_version", + "profile", + "provider", + "source", + "source_hash", + "session_command", + "hooks_hash", + } + or marker.get("schema") != 1 + or marker.get("agent_fleet_version") != __version__ + or marker.get("profile") != profile.id + or marker.get("provider") != "codex" + or marker.get("source") + != (str(provider.hooks_source) if provider.hooks_source is not None else None) + or marker.get("source_hash") != _source_hash(provider.hooks_source) + or not isinstance(marker.get("session_command"), str) + ): + return False + payload = _read_owned_json_object(profile.home / "hooks.json") + expected = _codex_hook_payload(provider.hooks_source, marker["session_command"]) + except (OSError, ValueError): + return False + expected_hash = _hook_payload_hash(expected) + return ( + payload == expected + and marker.get("hooks_hash") == expected_hash + and _hook_payload_hash(payload) == expected_hash + ) + + +def _merge_claude_project_trust(profile: Profile, project: TrustedProject) -> None: + path = profile.home / ".claude.json" + payload = _read_owned_json_object(path) + projects = payload.setdefault("projects", {}) + if not isinstance(projects, dict): + raise ValueError(f"Claude projects state must be an object: {path}") + changed = payload.get("hasCompletedOnboarding") is not True + payload["hasCompletedOnboarding"] = True + for root in {project.active_root, project.canonical_root}: + key = str(root) + existing = projects.setdefault(key, {}) + if not isinstance(existing, dict): + raise ValueError(f"Claude project state must be an object: {key}") + if existing.get("hasTrustDialogAccepted") is not True: + changed = True + existing["hasTrustDialogAccepted"] = True + if changed: + atomic_write_json(path, payload) + + +def claude_project_ready(profile: Profile, project: TrustedProject) -> bool: + try: + payload = _read_owned_json_object(profile.home / ".claude.json") + except ValueError: + return False + projects = payload.get("projects") + return ( + payload.get("hasCompletedOnboarding") is True + and isinstance(projects, dict) + and all( + isinstance(projects.get(str(root)), dict) + and projects[str(root)].get("hasTrustDialogAccepted") is True + for root in {project.active_root, project.canonical_root} + ) + ) + + +def prepare_profile_launch( + registry: Registry, + profile: Profile, + workspace: Path, +) -> TrustedProject: + if not profile_is_provisioned(profile): + raise ValueError(f"profile is not provisioned: {profile.id}") + project = resolve_trusted_project(registry, profile.provider, workspace) + if profile.provider == "claude": + _merge_claude_project_trust(profile, project) + if not claude_project_ready(profile, project): + raise ValueError(f"Claude project trust bootstrap failed for {profile.id}") + hook_health = profile_hook_health(registry, profile) + if not ( + hook_health["agent_fleet_session_hook"] and hook_health["inherited_workflow_hooks"] + ): + raise ValueError(f"managed Claude hook set is not ready for {profile.id}") + else: + project_hook = project.active_root / ".codex" / "hooks.json" + if project_hook.exists() or project_hook.is_symlink(): + raise ValueError(f"managed Codex launch refuses project hooks: {project_hook}") + if not _codex_config_ready(profile.home): + raise ValueError(f"managed Codex config is not ready for {profile.id}") + if not codex_hooks_ready(registry, profile): + raise ValueError(f"managed Codex hook set is not ready for {profile.id}") + return project + + +def profile_launch_ready( + registry: Registry, + profile: Profile, + workspace: Path, +) -> bool: + try: + project = resolve_trusted_project(registry, profile.provider, workspace) + except ValueError: + return False + if profile.provider == "claude": + hook_health = profile_hook_health(registry, profile) + return ( + claude_project_ready(profile, project) + and hook_health["agent_fleet_session_hook"] + and hook_health["inherited_workflow_hooks"] + ) + project_hook = project.active_root / ".codex" / "hooks.json" + return ( + not (project_hook.exists() or project_hook.is_symlink()) + and _codex_config_ready(profile.home) + and codex_hooks_ready(registry, profile) + ) + + def _ensure_codex_config(home: Path) -> None: path = home / "config.toml" - if not path.exists(): + if not path.exists() and not path.is_symlink(): path.write_text( 'cli_auth_credentials_store = "file"\n\n[features]\nhooks = true\n', encoding="utf-8", ) path.chmod(0o600) return + current = path.lstat() + if not stat.S_ISREG(current.st_mode) or current.st_uid != os.getuid(): + raise ValueError(f"managed Codex config must be a current-user regular file: {path}") try: raw = tomllib.loads(path.read_text(encoding="utf-8")) except tomllib.TOMLDecodeError as exc: @@ -103,10 +334,43 @@ def _ensure_codex_config(home: Path) -> None: features = raw.get("features", {}) if not isinstance(features, dict) or features.get("hooks") is not True: raise ValueError(f"managed Codex profile requires [features] hooks=true: {path}") + projects = raw.get("projects", {}) + if not isinstance(projects, dict) or projects: + raise ValueError(f"managed Codex profile cannot persist project trust: {path}") path.chmod(0o600) +def _codex_config_ready(home: Path) -> bool: + path = home / "config.toml" + try: + current = path.lstat() + if ( + not stat.S_ISREG(current.st_mode) + or current.st_uid != os.getuid() + or stat.S_IMODE(current.st_mode) & 0o077 + ): + return False + raw = tomllib.loads(path.read_text(encoding="utf-8")) + except (OSError, tomllib.TOMLDecodeError): + return False + features = raw.get("features", {}) + projects = raw.get("projects", {}) + return ( + raw.get("cli_auth_credentials_store") == "file" + and isinstance(features, dict) + and features.get("hooks") is True + and isinstance(projects, dict) + and not projects + ) + + def _share_workflow_entries(profile: Profile, provider: ProviderConfig) -> list[str]: + disallowed = sorted(set(provider.shared_entries) - SHARED_WORKFLOW_ENTRIES[provider.name]) + if disallowed: + raise ValueError( + f"providers.{provider.name}.shared_entries contains non-workflow assets: " + + ", ".join(disallowed) + ) if provider.base_home is None: return [] shared: list[str] = [] @@ -172,11 +436,16 @@ def provision_profile(registry: Registry, profile: Profile) -> dict[str, Any]: elif profile.provider == "codex": ensure_private_dir(profile.home / "hooks") _ensure_codex_config(profile.home) - _install_session_hook(profile.home / "hooks.json", provider.hooks_source) + _install_codex_hooks(profile, provider) marker = profile.home / ".agent-fleet-profile.json" atomic_write_json( marker, - {"schema": 1, "profile": profile.id, "provider": profile.provider}, + { + "schema": 2, + "agent_fleet_version": __version__, + "profile": profile.id, + "provider": profile.provider, + }, ) os.chmod(marker, 0o600) return { @@ -190,13 +459,25 @@ def provision_profile(registry: Registry, profile: Profile) -> dict[str, Any]: def profile_is_provisioned(profile: Profile) -> bool: marker = profile.home / ".agent-fleet-profile.json" - if not profile.home.is_dir() or not marker.is_file(): + if not profile.home.is_dir(): return False try: + current = marker.lstat() + if ( + not stat.S_ISREG(current.st_mode) + or current.st_uid != os.getuid() + or stat.S_IMODE(current.st_mode) != 0o600 + ): + return False raw = json.loads(marker.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return False - return raw.get("profile") == profile.id and raw.get("provider") == profile.provider + return raw == { + "schema": 2, + "agent_fleet_version": __version__, + "profile": profile.id, + "provider": profile.provider, + } def profile_hook_health(registry: Registry, profile: Profile) -> dict[str, bool]: @@ -250,11 +531,16 @@ def profile_hook_health(registry: Registry, profile: Profile) -> dict[str, bool] ): source_ok = False break - return { + if profile.provider == "codex": + source_ok = codex_hooks_ready(registry, profile) + health = { "agent_fleet_session_hook": any(HOOK_MARKER in command for command in commands), "herdr_session_hook": any("herdr-agent-state" in command for command in commands), "inherited_workflow_hooks": source_ok, } + if profile.provider == "codex": + health["closed_profile_hooks"] = codex_hooks_ready(registry, profile) + return health def profile_shared_assets_healthy(registry: Registry, profile: Profile) -> bool: diff --git a/tools/agent-fleet/src/agent_fleet/quota.py b/tools/agent-fleet/src/agent_fleet/quota.py index a6bb6b98c03..bddcc970b7a 100644 --- a/tools/agent-fleet/src/agent_fleet/quota.py +++ b/tools/agent-fleet/src/agent_fleet/quota.py @@ -70,6 +70,23 @@ def restore_quota_cache( path.unlink() +def discard_quota_cache(registry: Registry, profile_id: str) -> None: + restore_quota_cache(registry, profile_id, QuotaCacheSnapshot(False)) + + +def has_remote_identity_proof(quota: dict[str, Any]) -> bool: + fingerprint = quota.get("identity_fingerprint") + return ( + quota.get("status") == "fresh" + and quota.get("verified_at") is not None + and quota.get("headroom_percent") is not None + and isinstance(quota.get("windows"), list) + and bool(quota["windows"]) + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + ) + + def _number(value: Any) -> float | None: if isinstance(value, (int, float)) and not isinstance(value, bool): return max(0.0, min(100.0, float(value))) diff --git a/tools/agent-fleet/src/agent_fleet/scheduler.py b/tools/agent-fleet/src/agent_fleet/scheduler.py index 1556af00e5a..bb9c6cf3401 100644 --- a/tools/agent-fleet/src/agent_fleet/scheduler.py +++ b/tools/agent-fleet/src/agent_fleet/scheduler.py @@ -2,6 +2,7 @@ from collections import Counter from datetime import UTC, datetime +from pathlib import Path from typing import Any from .audit import append_audit @@ -15,8 +16,9 @@ state_lock, ) from .models import Profile, Registry +from .projects import invocation_workspace from .providers import auth_status -from .provision import profile_is_provisioned +from .provision import prepare_profile_launch, profile_is_provisioned from .quota import quota_routeability, read_quota, refresh_due_quotas from .util import atomic_write_json, task_key @@ -129,6 +131,7 @@ def _select_and_acquire( explicit_profile: bool = False, ignore_reserve: bool = False, recovery_reservation: bool = False, + workspace: Path | None = None, ) -> dict[str, Any]: if ignore_reserve and profile_id is None: raise ValueError("ignoring quota reserve requires an explicit profile") @@ -152,6 +155,7 @@ def _select_and_acquire( and (profile_id is None or profile.id == profile_id) and profile_is_provisioned(profile) ] + active_workspace = workspace or invocation_workspace() # Quota/identity writes share the provider maintenance interlock, but a # selection does not hold it while committing its lease. The state-lock # check below closes both races: maintenance either owns the marker first @@ -173,6 +177,9 @@ def _select_and_acquire( registry.settings.lock_stale_seconds, ): recover_pending_codex_transactions(registry, provider_name) + for profile in scoped_profiles: + if profile.provider == provider_name: + prepare_profile_launch(registry, profile, active_workspace) refresh_provider_identity_anchors_if_due(registry, provider_name) refresh_due_quotas( registry, @@ -366,6 +373,7 @@ def select_and_acquire( explicit_profile: bool = False, ignore_reserve: bool = False, recovery_reservation: bool = False, + workspace: Path | None = None, ) -> dict[str, Any]: return _select_and_acquire( registry, @@ -378,4 +386,5 @@ def select_and_acquire( explicit_profile=explicit_profile, ignore_reserve=ignore_reserve, recovery_reservation=recovery_reservation, + workspace=workspace, ) diff --git a/tools/agent-fleet/tests/conftest.py b/tools/agent-fleet/tests/conftest.py index 5b1620185f3..bc9b54822e9 100644 --- a/tools/agent-fleet/tests/conftest.py +++ b/tools/agent-fleet/tests/conftest.py @@ -3,6 +3,7 @@ import json import os import stat +import subprocess from dataclasses import replace from pathlib import Path @@ -28,6 +29,15 @@ def fleet(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Registry, Pa monkeypatch.setenv("AGENT_FLEET_CONFIG", str(config)) monkeypatch.setenv("AGENT_FLEET_STATE_DIR", str(state)) monkeypatch.setenv("AGENT_FLEET_SHARE_DIR", str(share)) + trusted_project = tmp_path / "trusted-project" + trusted_project.mkdir() + subprocess.run( + ["git", "init", "-q", str(trusted_project)], + check=True, + capture_output=True, + text=True, + ) + monkeypatch.chdir(trusted_project) registry = initial_registry(3, 2) quota_binary = tmp_path / "quota-axi" quota_binary.write_text( @@ -99,6 +109,7 @@ def fleet(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Registry, Pa hooks_source=hooks, shared_entries=(shared,), desktop_identity_file=desktop_file if name == "claude" else None, + trusted_projects=(trusted_project.resolve(),), ) registry = replace(registry, providers=providers) save_registry(registry, config) diff --git a/tools/agent-fleet/tests/test_config_and_provision.py b/tools/agent-fleet/tests/test_config_and_provision.py index eb38a12b3dd..9f78584f976 100644 --- a/tools/agent-fleet/tests/test_config_and_provision.py +++ b/tools/agent-fleet/tests/test_config_and_provision.py @@ -118,3 +118,24 @@ def test_shared_asset_install_refuses_existing_non_symlink( (profile.home / "AGENTS.md").write_text("attacker-owned\n", encoding="utf-8") with pytest.raises(ValueError, match="refusing to replace"): provision_profile(registry, profile) + + +@pytest.mark.parametrize( + ("provider", "entry"), + [("claude", ".claude.json"), ("codex", "auth.json")], +) +def test_registry_rejects_non_workflow_shared_entries( + fleet: tuple[object, Path], provider: str, entry: str +) -> None: + _, path = fleet + text = path.read_text(encoding="utf-8") + marker = f"[providers.{provider}]" + before, after = text.split(marker, 1) + section, remainder = after.split("\n[", 1) + lines = [ + f'shared_entries = ["{entry}"]' if line.startswith("shared_entries =") else line + for line in section.splitlines() + ] + path.write_text(before + marker + "\n".join(lines) + "\n[" + remainder, encoding="utf-8") + with pytest.raises(ValueError, match="non-workflow assets"): + load_registry(path) diff --git a/tools/agent-fleet/tests/test_contract_status.py b/tools/agent-fleet/tests/test_contract_status.py index 39e2a051e18..bd2da0caa18 100644 --- a/tools/agent-fleet/tests/test_contract_status.py +++ b/tools/agent-fleet/tests/test_contract_status.py @@ -39,7 +39,7 @@ def test_contract_and_version_do_not_require_registry(tmp_path: Path) -> None: str(missing), command, ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -107,7 +107,7 @@ def test_pool_status_reports_provider_level_fallback( "--provider", "codex", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -163,7 +163,7 @@ def test_enroll_uses_codex_device_auth_then_verifies_while_disabled( "enroll", "codex-1", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -219,7 +219,7 @@ def test_verify_keeps_a_remotely_rejected_profile_disabled( "verify", "codex-1", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -230,3 +230,96 @@ def test_verify_keeps_a_remotely_rejected_profile_disabled( assert payload["ready"] is False assert payload["profiles"][0]["enabled"] is False assert load_registry(config).require_profile("codex-1").enabled is False + + +def test_claude_enrollment_discards_prelogin_identity_proof( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + old_quota = json.loads(quota_path(registry, "claude-1").read_text(encoding="utf-8")) + assert old_quota.get("identity_fingerprint") + fixtures = tmp_path / "quota" + fixtures.mkdir() + keychain_required = { + "providers": [ + { + "provider": "claude", + "state": {"status": "auth_required", "reason": "keychain_access_required"}, + "windows": [], + } + ] + } + (fixtures / "claude-1.json").write_text(json.dumps(keychain_required), encoding="utf-8") + base = { + "providers": [ + { + "provider": "claude", + "account": {"accountId": "base-account"}, + "state": {"status": "fresh", "refreshedAt": utc_now()}, + "windows": [{"id": "five_hour", "kind": "session", "percentRemaining": 80}], + } + ] + } + (fixtures / "claude-base-anchor.json").write_text(json.dumps(base), encoding="utf-8") + + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + env["AGENT_FLEET_QUOTA_FIXTURE_DIR"] = str(fixtures) + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "profile", + "enroll", + "claude-1", + ], + cwd=Path.cwd(), + env=env, + text=True, + capture_output=True, + timeout=20, + check=True, + ) + + payload = json.loads(result.stdout) + stored = json.loads(quota_path(registry, "claude-1").read_text(encoding="utf-8")) + assert payload["verification_pending"] is True + assert payload["credential_verified"] is False + assert stored.get("identity_fingerprint") is None + assert stored["reason"] == "keychain_access_required" + + +def test_toon_preflight_precedes_registry_mutation(tmp_path: Path) -> None: + project_root = Path(__file__).parents[1] + config = tmp_path / "accounts.toml" + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + env["PATH"] = str(tmp_path) + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--config", + str(config), + "init", + "--claude", + "1", + ], + cwd=tmp_path, + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert "TOON output requested" in result.stderr + assert not config.exists() diff --git a/tools/agent-fleet/tests/test_exec_sessions.py b/tools/agent-fleet/tests/test_exec_sessions.py index 4ed9ea68c5f..4c7f788c948 100644 --- a/tools/agent-fleet/tests/test_exec_sessions.py +++ b/tools/agent-fleet/tests/test_exec_sessions.py @@ -73,7 +73,7 @@ def test_exec_uses_selected_home_and_clears_ambient_credentials( "--", "example-argument", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -81,7 +81,13 @@ def test_exec_uses_selected_home_and_clears_ambient_credentials( check=True, ) payload = json.loads(result.stdout) - assert payload["argv"] == ["example-argument"] + assert payload["argv"][-1:] == ["example-argument"] + assert payload["argv"][:4] == [ + "--disable", + "plugins", + "--disable", + "plugin_sharing", + ] assert payload["profile"] == "codex-1" assert payload["task"] == "exec-task" assert payload["codex_home"] == str(profile.home) @@ -139,7 +145,7 @@ def test_session_hook_persists_profile_and_provider_session( "--task", "hook-task", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -162,7 +168,7 @@ def test_session_hook_persists_profile_and_provider_session( "--task", "hook-task", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -207,7 +213,7 @@ def test_direct_resume_without_managed_task_is_refused( "--", "extra", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -250,7 +256,7 @@ def test_direct_worker_exec_without_managed_task_is_refused( "--", "example-argument", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -261,6 +267,78 @@ def test_direct_worker_exec_without_managed_task_is_refused( assert "worker exec requires --task" in result.stderr +def test_worker_exec_refuses_auth_resume_and_plugin_overrides( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + project_root = Path(__file__).parents[1] + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + for provider_args in ( + ["login"], + ["resume", "foreign-session"], + ["--enable", "plugins"], + ["-c", 'projects."/tmp".trust_level="trusted"'], + ): + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--format", + "json", + "--config", + str(config), + "exec", + "--task", + "guarded-exec", + "--profile", + "codex-1", + "--", + *provider_args, + ], + cwd=Path.cwd(), + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert ( + "refuses" in json.loads(result.stdout)["error"] + or "disabled" in json.loads(result.stdout)["error"] + ) + + +def test_init_has_no_force_replacement_path(tmp_path: Path) -> None: + project_root = Path(__file__).parents[1] + config = tmp_path / "accounts.toml" + env = dict(os.environ) + env["PYTHONPATH"] = str(project_root / "src") + result = subprocess.run( + [ + sys.executable, + "-m", + "agent_fleet", + "--config", + str(config), + "init", + "--claude", + "1", + "--force", + ], + cwd=tmp_path, + env=env, + text=True, + capture_output=True, + timeout=20, + check=False, + ) + assert result.returncode == 2 + assert not config.exists() + + def test_live_task_cannot_be_rebound(fleet: tuple[object, Path]) -> None: _, config = fleet registry = load_registry(config) diff --git a/tools/agent-fleet/tests/test_project_bootstrap.py b/tools/agent-fleet/tests/test_project_bootstrap.py new file mode 100644 index 00000000000..675ca0a99d2 --- /dev/null +++ b/tools/agent-fleet/tests/test_project_bootstrap.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest + +from agent_fleet.config import load_registry +from agent_fleet.projects import resolve_trusted_project +from agent_fleet.providers import managed_argv, resume_argv, validate_worker_arguments +from agent_fleet.provision import ( + codex_hooks_ready, + prepare_profile_launch, + profile_launch_ready, + provision_profile, +) + + +def test_claude_bootstrap_preserves_opaque_state_and_ignores_base_state( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = load_registry(config) + profile = registry.require_profile("claude-1") + provision_profile(registry, profile) + base_home = registry.require_provider("claude").base_home + assert base_home is not None + (base_home / ".claude.json").write_text( + json.dumps({"oauthAccount": {"secret": "must-not-copy"}}), encoding="utf-8" + ) + state_path = profile.home / ".claude.json" + state_path.write_text( + json.dumps( + { + "oauthAccount": {"accountUuid": "opaque-worker"}, + "projects": {"/unrelated": {"opaque": True}}, + } + ), + encoding="utf-8", + ) + state_path.chmod(0o600) + + project = prepare_profile_launch(registry, profile, Path.cwd()) + state = json.loads(state_path.read_text(encoding="utf-8")) + + assert state["oauthAccount"] == {"accountUuid": "opaque-worker"} + assert state["projects"]["/unrelated"] == {"opaque": True} + assert "secret" not in json.dumps(state) + assert state["hasCompletedOnboarding"] is True + for root in {project.active_root, project.canonical_root}: + assert state["projects"][str(root)]["hasTrustDialogAccepted"] is True + assert profile_launch_ready(registry, profile, Path.cwd()) is True + + +def test_linked_worktree_matches_registered_git_common_dir( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + registered = Path.cwd() + (registered / "tracked.txt").write_text("tracked\n", encoding="utf-8") + subprocess.run(["git", "add", "tracked.txt"], cwd=registered, check=True) + subprocess.run( + [ + "git", + "-c", + "user.name=Agent Fleet Tests", + "-c", + "user.email=agent-fleet@example.invalid", + "commit", + "-qm", + "fixture", + ], + cwd=registered, + check=True, + ) + linked = tmp_path / "linked-worktree" + subprocess.run( + ["git", "worktree", "add", "-q", "-b", "linked-test", str(linked)], + cwd=registered, + check=True, + ) + + project = resolve_trusted_project(registry, "codex", linked) + + assert project.active_root == linked + assert project.canonical_root == registered + + +def test_unrelated_and_symlinked_workspaces_fail_closed( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + unrelated = tmp_path / "unrelated" + unrelated.mkdir() + subprocess.run(["git", "init", "-q", str(unrelated)], check=True) + with pytest.raises(ValueError, match="not registered"): + resolve_trusted_project(registry, "claude", unrelated) + linked = tmp_path / "project-link" + linked.symlink_to(Path.cwd(), target_is_directory=True) + with pytest.raises(ValueError, match="must not be symlinked"): + resolve_trusted_project(registry, "claude", linked) + + +def test_codex_launch_uses_exact_managed_prefix_for_new_and_resume( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = load_registry(config) + profile = registry.require_profile("codex-1") + provision_profile(registry, profile) + project = prepare_profile_launch(registry, profile, Path.cwd()) + trust = f'projects.{json.dumps(str(project.active_root))}.trust_level="trusted"' + prefix = [ + "--disable", + "plugins", + "--disable", + "plugin_sharing", + "-c", + trust, + "--dangerously-bypass-hook-trust", + ] + + assert managed_argv(registry, profile, project.active_root, ["--full-auto"])[1:] == [ + *prefix, + "--full-auto", + ] + assert resume_argv( + registry, + profile, + "session-1", + ["--full-auto"], + active_root=project.active_root, + )[1:] == ["resume", *prefix, "session-1", "--full-auto"] + + +def test_codex_launch_refuses_changed_hooks_markers_sources_and_project_hooks( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + registry = load_registry(config) + profile = registry.require_profile("codex-1") + provision_profile(registry, profile) + assert codex_hooks_ready(registry, profile) is True + + hooks_path = profile.home / "hooks.json" + hooks = json.loads(hooks_path.read_text(encoding="utf-8")) + hooks["hooks"]["SessionStart"].append({"matcher": "", "hooks": []}) + hooks_path.write_text(json.dumps(hooks), encoding="utf-8") + assert codex_hooks_ready(registry, profile) is False + + provision_profile(registry, profile) + marker = profile.home / ".agent-fleet-hooks.json" + marker_payload = json.loads(marker.read_text(encoding="utf-8")) + marker_payload["agent_fleet_version"] = "0.0.0" + marker.write_text(json.dumps(marker_payload), encoding="utf-8") + assert codex_hooks_ready(registry, profile) is False + + provision_profile(registry, profile) + source = registry.require_provider("codex").hooks_source + assert source is not None + source.write_text(source.read_text(encoding="utf-8") + "\n", encoding="utf-8") + assert codex_hooks_ready(registry, profile) is False + + provision_profile(registry, profile) + project_hooks = Path.cwd() / ".codex" / "hooks.json" + project_hooks.parent.mkdir() + project_hooks.write_text("{}\n", encoding="utf-8") + with pytest.raises(ValueError, match="refuses project hooks"): + prepare_profile_launch(registry, profile, Path.cwd()) + + +def test_codex_worker_arguments_cannot_reenable_managed_surfaces( + fleet: tuple[object, Path], +) -> None: + _, config = fleet + profile = load_registry(config).require_profile("codex-1") + for arguments in ( + ["login"], + ["resume", "other"], + ["--enable=plugin_sharing"], + ["--config", 'projects."/tmp".trust_level="trusted"'], + ["-C", "/tmp"], + ): + with pytest.raises(ValueError): + validate_worker_arguments(profile, arguments) diff --git a/tools/agent-fleet/tests/test_quota_scheduler.py b/tools/agent-fleet/tests/test_quota_scheduler.py index bc84ba7c094..0551dc39c84 100644 --- a/tools/agent-fleet/tests/test_quota_scheduler.py +++ b/tools/agent-fleet/tests/test_quota_scheduler.py @@ -157,7 +157,7 @@ def test_sticky_resume_can_reacquire_its_profile_below_reserve( "--task", "resumed-task", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -474,7 +474,7 @@ def test_concurrent_reservations_are_atomic_and_balanced( "--provider", "claude", ], - cwd=project_root, + cwd=Path.cwd(), env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, diff --git a/tools/agent-fleet/tests/test_safety_transactions.py b/tools/agent-fleet/tests/test_safety_transactions.py index bd96787dd08..b9e263df5b9 100644 --- a/tools/agent-fleet/tests/test_safety_transactions.py +++ b/tools/agent-fleet/tests/test_safety_transactions.py @@ -2,6 +2,7 @@ import json import os +import stat import subprocess import sys from dataclasses import replace @@ -19,6 +20,7 @@ finalize_codex_promotion, prepare_codex_promotion, recover_pending_codex_transaction, + rollback_codex_promotion, ) from agent_fleet.leases import bind_lease, new_lease from agent_fleet.locks import provider_enrollment_lock @@ -51,7 +53,7 @@ def _prepared_codex_transaction(registry): old_auth.chmod(0o600) old_quota = quota_path(registry, target.id).read_bytes() snapshot = snapshot_quota_cache(registry, target.id) - stage = create_codex_login_stage(target) + stage = create_codex_login_stage(registry, target) staged_auth = stage.home / "auth.json" staged_auth.write_text('{"token":"new-test-token"}\n', encoding="utf-8") staged_auth.chmod(0o600) @@ -85,7 +87,7 @@ def test_codex_crash_recovery_rolls_back_auth_and_quota_before_cleanup( assert quota_path(registry, target.id).read_bytes() == old_quota assert not transaction.journal_path.exists() discard_codex_promotion(promotion, target) - discard_codex_stage(stage, target) + discard_codex_stage(registry, stage, target) def test_codex_activation_failure_after_replace_is_recovered( @@ -118,7 +120,7 @@ def fail_after_replace(path: Path) -> None: assert "old-test-token" in old_auth.read_text(encoding="utf-8") assert quota_path(registry, target.id).read_bytes() == old_quota discard_codex_promotion(promotion, target) - discard_codex_stage(stage, target) + discard_codex_stage(registry, stage, target) def test_codex_committed_crash_recovery_keeps_promoted_auth( @@ -149,7 +151,52 @@ def fail_journal_cleanup(path: Path, *args, **kwargs) -> None: assert "new-test-token" in target_auth.read_text(encoding="utf-8") assert not transaction.journal_path.exists() discard_codex_promotion(promotion, target) - discard_codex_stage(stage, target) + discard_codex_stage(registry, stage, target) + + +def test_codex_rollback_recovery_is_idempotent_after_auth_restore( + fleet: tuple[object, Path], monkeypatch: pytest.MonkeyPatch +) -> None: + _, config = fleet + registry = load_registry(config) + target, stage, promotion, snapshot, target_auth, old_quota = _prepared_codex_transaction( + registry + ) + transaction = activate_codex_promotion(registry, target, promotion, snapshot) + real_write_phase = enrollment._write_journal_phase + + def fail_after_auth_restore(path: Path, journal: dict, phase: str) -> None: + if phase == "auth_restored": + raise OSError("injected crash after auth restore") + real_write_phase(path, journal, phase) + + monkeypatch.setattr(enrollment, "_write_journal_phase", fail_after_auth_restore) + with pytest.raises(OSError, match="injected crash"): + rollback_codex_promotion(registry, target, transaction) + monkeypatch.setattr(enrollment, "_write_journal_phase", real_write_phase) + + assert recover_pending_codex_transaction(registry, target) is True + assert "old-test-token" in target_auth.read_text(encoding="utf-8") + assert quota_path(registry, target.id).read_bytes() == old_quota + assert not transaction.journal_path.exists() + discard_codex_promotion(promotion, target) + discard_codex_stage(registry, stage, target) + + +def test_codex_stage_does_not_modify_custom_existing_parent( + fleet: tuple[object, Path], tmp_path: Path +) -> None: + _, config = fleet + registry = load_registry(config) + parent = tmp_path / "shared-parent" + parent.mkdir(mode=0o755) + profile = replace(registry.require_profile("codex-1"), home=parent / "codex-home") + + stage = create_codex_login_stage(registry, profile) + + assert stat.S_IMODE(parent.stat().st_mode) == 0o755 + assert stage.home.parent == registry.settings.state_dir / "staging" / "codex" + discard_codex_stage(registry, stage, profile) def test_provider_maintenance_marker_blocks_new_lease( @@ -201,7 +248,7 @@ def test_profile_add_obeys_provider_maintenance_lock( "--provider", "claude", ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, @@ -277,7 +324,7 @@ def test_enrollment_refuses_while_managed_provider_launch_is_alive( "enroll", profile.id, ], - cwd=project_root, + cwd=Path.cwd(), env=env, text=True, capture_output=True, From b7b0df5e796d6c691c10e327412a9b07ca85ac20 Mon Sep 17 00:00:00 2001 From: Dongkeun Lee Date: Sat, 18 Jul 2026 00:30:36 -0400 Subject: [PATCH 03/54] no-mistakes(review): Harden Agent Fleet routing and trust boundaries --- bin/fm-account-routing-lib.sh | 43 +-- bin/fm-spawn.sh | 17 +- docs/configuration.md | 4 +- tests/fm-account-routing.test.sh | 21 +- tools/agent-fleet/README.md | 25 +- tools/agent-fleet/src/agent_fleet/cli.py | 134 ++++++--- tools/agent-fleet/src/agent_fleet/config.py | 10 +- tools/agent-fleet/src/agent_fleet/doctor.py | 26 +- tools/agent-fleet/src/agent_fleet/identity.py | 117 ++++---- tools/agent-fleet/src/agent_fleet/paths.py | 4 + tools/agent-fleet/src/agent_fleet/projects.py | 74 +++-- .../agent-fleet/src/agent_fleet/providers.py | 126 +++++++-- .../agent-fleet/src/agent_fleet/provision.py | 69 ++++- tools/agent-fleet/src/agent_fleet/quota.py | 93 +++---- .../agent-fleet/src/agent_fleet/scheduler.py | 127 +++++---- tools/agent-fleet/src/agent_fleet/sessions.py | 4 +- tools/agent-fleet/tests/conftest.py | 7 + .../agent-fleet/tests/test_contract_status.py | 46 +++- tools/agent-fleet/tests/test_exec_sessions.py | 70 ++++- .../tests/test_project_bootstrap.py | 256 +++++++++++++++++- .../agent-fleet/tests/test_quota_scheduler.py | 153 ++++++++++- .../tests/test_safety_transactions.py | 32 +++ 22 files changed, 1135 insertions(+), 323 deletions(-) diff --git a/bin/fm-account-routing-lib.sh b/bin/fm-account-routing-lib.sh index 4d0565a431c..e3e21f74d5e 100644 --- a/bin/fm-account-routing-lib.sh +++ b/bin/fm-account-routing-lib.sh @@ -867,10 +867,10 @@ fm_account_json_field() { #