From 303e26d2154142fa03781583b1ca045baa6663f1 Mon Sep 17 00:00:00 2001 From: moneytools Date: Fri, 9 Oct 2026 21:00:55 -0500 Subject: [PATCH 1/5] feat(budget): aegis budget status|check|reset, docs/budget.md, latency - status: today's project total and its sessions, rebuilt from the logs (--json); check: exit 0 under / 3 over / 65 unusable / 66 none, with an optional --agent --session check; reset --day: a marker signed with the policy key acknowledging an incomplete day, honoured by the hook - evaluate refactored to share the project-only check - price table parse cached by content hash (~15 ms per hook call saved); scripts/budget_latency.py: budget adds ~9 ms per hook call - docs/budget.md, signed budget.example.yaml from aegis init, README/docs index links; tests keep the real ~/.cache and ~/.local/state untouched --- CHANGELOG.md | 8 + README.md | 1 + data/budget.example.yaml | 21 ++ data/budget.example.yaml.sig | 2 + docs/budget.md | 154 +++++++++++++++ docs/index.md | 1 + scripts/budget_latency.py | 113 +++++++++++ src/aegis_core/_examples/budget.example.yaml | 21 ++ .../_examples/budget.example.yaml.sig | 2 + src/aegis_core/budget/evaluate.py | 182 ++++++++++------- src/aegis_core/budget/gate.py | 8 +- src/aegis_core/budget/pricing.py | 27 ++- src/aegis_core/budget/report.py | 187 ++++++++++++++++++ src/aegis_core/cli.py | 147 ++++++++++++-- src/aegis_core/hook.py | 4 +- tests/conftest.py | 14 ++ tests/test_budget_cli.py | 171 ++++++++++++++++ tests/test_budget_policy.py | 33 +++- tests/test_signing.py | 3 +- 19 files changed, 1008 insertions(+), 91 deletions(-) create mode 100644 data/budget.example.yaml create mode 100644 data/budget.example.yaml.sig create mode 100644 docs/budget.md create mode 100644 scripts/budget_latency.py create mode 100644 src/aegis_core/_examples/budget.example.yaml create mode 100644 src/aegis_core/_examples/budget.example.yaml.sig create mode 100644 src/aegis_core/budget/report.py create mode 100644 tests/test_budget_cli.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b19dd17..f576fb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,14 @@ test on each agent are in. can block them (Claude Code and Codex `UserPromptSubmit`, Gemini CLI `BeforeAgent`); a warning is shown once per session per threshold (`systemMessage`; an OpenCode toast; stderr for Copilot CLI). A `budget.yaml` that cannot be used (bad signature, shape, authority) fails closed. +- `aegis budget status|check|reset` (design §7): today's project total and its sessions, rebuilt + from the logs (`--json`); `check` exits 0 under, 3 over, 65 for an unusable `budget.yaml`, 66 + for none; `reset --day` acknowledges an incomplete day's history with a marker signed by the + policy key, so `on_unknown_log: deny` stops denying it (an agent without the key cannot). +- `docs/budget.md`; `aegis init` writes a signed `budget.example.yaml`. +- `scripts/budget_latency.py`: the budget adds ~9 ms to a hook call (74 vs 65 ms median, + `results/budget-latency.json`), under the 20 ms target. The parsed price table is cached by + content hash; parsing it on every call had cost ~15 ms. - `aegis install` adds the budget hooks when the policy directory that applies holds a `budget.yaml` (`--budget` / `--no-budget` to choose): the tool hook sees every tool call and the prompt hook is added where the agent has one. The Claude Code plugin and the Gemini diff --git a/README.md b/README.md index 0bfc9d1..07943fb 100644 --- a/README.md +++ b/README.md @@ -298,6 +298,7 @@ of anything you care about. | [`docs/constraints.md`](docs/constraints.md) | Writing constraints, metadata vocabulary, authority policy, environment mapping | | [`docs/configuration.md`](docs/configuration.md) | Configuration/config-dir discovery, signing, source verification, the identity model (`agents.yaml`), rate limits & ledger | | [`docs/server-side.md`](docs/server-side.md) | Preview: compiling the policy to AWS Service Control Policies scoped to agent identities | +| [`docs/budget.md`](docs/budget.md) | The session budget cap: per-session and per-project-per-day limits in estimated USD or tokens, across Claude Code, Codex, Gemini CLI, OpenCode and Copilot CLI | | [`docs/benchmark.md`](docs/benchmark.md) | Benchmark methodology, full results table, real LLM baselines, agent-harness baselines, corpus, adversarial suite | | [`docs/CONTRIBUTING.md`](docs/CONTRIBUTING.md) | Adding a new parser or tool | | [`docs/dev/PLAN.md`](docs/dev/PLAN.md) | Design plan and open gaps | diff --git a/data/budget.example.yaml b/data/budget.example.yaml new file mode 100644 index 0000000..748c2cc --- /dev/null +++ b/data/budget.example.yaml @@ -0,0 +1,21 @@ +# The session budget cap (docs/budget.md). Copy to budget.yaml, set your own limits, and sign +# it with `aegis sign`; its principal must hold the 'budget' class in authority.yaml. Until +# it is named budget.yaml it does nothing. Then reinstall your agents' hooks (`aegis install +# `) so they see every tool call, and check with: aegis budget status +version: 1 +principal: admin # must hold the 'budget' class in authority.yaml +unit: usd # usd (estimated from list prices) | tokens; required +session: + limit: 20 # per agent session + warn_at: 0.8 # warn once at 80% +project_day: + limit: 100 # every measured session in this project, per day + warn_at: 0.8 + tz: UTC # the day boundary +agents: [claude, codex, gemini, opencode, copilot] # only listed agents are measured +on_unknown_log: allow # allow (warn) | deny: usage Aegis cannot verify +on_unknown_price: estimate # estimate (warn) | deny: models with no price +pricing: {} # signed overrides, USD per million tokens: +# claude-opus-5-5: {input: 4, output: 20, cache_read: 0.2, cache_write: 8} +copilot: + premium_requests: 50 # Copilot CLI reports premium requests, not tokens diff --git a/data/budget.example.yaml.sig b/data/budget.example.yaml.sig new file mode 100644 index 0000000..4fec410 --- /dev/null +++ b/data/budget.example.yaml.sig @@ -0,0 +1,2 @@ +aegis-sig-v1 +d8ad9e56f8548c60028c05c17ce2f46ca190127cbafde00669faef484f3eaea4 diff --git a/docs/budget.md b/docs/budget.md new file mode 100644 index 0000000..237e25c --- /dev/null +++ b/docs/budget.md @@ -0,0 +1,154 @@ +# Session budget cap + +A coding agent stuck in a loop, or fanning out sub-agents, spends money (API keys) or quota +(subscriptions) with nothing to stop it. Aegis can cap that spend per agent session and per +project per day, in estimated dollars or in tokens, across Claude Code, Codex, Gemini CLI, +OpenCode and Copilot CLI (premium requests). The cap is written into a **signed** policy +file, so an agent cannot raise its own limit by editing a file. + +← back to the [README](https://github.com/moneytool/aegis-devops#readme) · design: [`dev/DESIGN-v1.0-budget-cap.md`](https://github.com/moneytool/aegis-devops/blob/main/docs/dev/DESIGN-v1.0-budget-cap.md) + +## What it guarantees, and what it does not + +> Aegis denies covered tool calls and prompts when the usage available at that hook reaches +> the limit. It does not impose a finite bound on additional model spending: in-flight +> requests, delayed usage records, retries after a denial, and tool-free activity can add +> cost. + +Hooks run before tool calls and prompts, never before a model request, and an agent writes a +request's usage to its log only after the request has been paid for. So a session can end +**above** its limit — by the request that crossed it, by the model's replies to the denials, +by sub-agents that were already running, and, in OpenCode and Copilot CLI (which have no prompt +hook that can block), by chat that uses no tools. `aegis budget status` shows any usage over +a limit as **overshoot**. If you must never exceed an amount, set the limit below it, and put a +spend limit in front of the model requests themselves as well (a gateway, your vendor's +spending limits); those are separate layers whose guarantees you need to check. + +The budget is a guard against loops and mistakes. An agent runs as you and can rewrite its own +session log; Aegis does not defend against an agent that does that on purpose. + +## Set it up + +```bash +aegis init .aegis # once; writes budget.example.yaml among the examples +cp .aegis/budget.example.yaml .aegis/budget.yaml +$EDITOR .aegis/budget.yaml # your limits +aegis sign --key file:KEY .aegis/budget.yaml +aegis install claude # and codex, gemini, opencode, copilot: reinstall +aegis budget status +``` + +Reinstalling matters: with a `budget.yaml` in the policy directory, `aegis install` widens each +agent's hook from shell commands to **every tool call** and adds the prompt hook where the +agent has one (`--budget` / `--no-budget` choose explicitly). The Claude Code plugin and the +Gemini extension already carry those entries; they start no extra process in projects without +a budget. + +## `budget.yaml` + +```yaml +version: 1 +principal: admin # must hold the 'budget' class in authority.yaml +unit: usd # usd | tokens — required +session: {limit: 20, warn_at: 0.8} # per agent session +project_day: {limit: 100, warn_at: 0.8, tz: America/Chicago} # all sessions in the project, per day +agents: [claude, codex, gemini, opencode, copilot] # only listed agents are measured +on_unknown_log: allow # allow (warn) | deny +on_unknown_price: estimate # estimate (warn) | deny — usd only +pricing: # signed overrides, USD per million tokens + claude-opus-5-5: {input: 4, output: 20, cache_read: 0.2, cache_write: 8} +copilot: {premium_requests: 50, warn_at: 0.8} # per Copilot CLI session +``` + +- **Who may change it.** The file is signed like every policy file, and its `principal` must + hold the `budget` class (the example `authority.yaml` gives it to `admin`). A file that was + edited after signing, or whose principal lacks the class, is refused — and inside a project + with a budget the hook then **denies** every covered call with a message saying so, rather + than run without the budget you asked for. +- **Errors, never guesses.** Unknown fields, a missing `unit`, a limit that is not a finite + positive number, `warn_at` outside (0, 1), an unknown time zone or agent, a price override + missing a category: all load errors. +- **`tokens`** counts input, output, cache-read and cache-write tokens together, and ignores + prices. +- **Copilot CLI** writes token counts only at shutdown, so it is measured in premium requests + and only when `copilot.premium_requests` is set. +- **Cursor and VS Code** expose no token data and are not measured. + +## What happens near and at the limit + +| | Tool calls | New prompts | You see | +|---|---|---|---| +| Below `warn_at` | allowed | allowed | nothing | +| Past `warn_at` | allowed | allowed | one warning per session per threshold | +| At the limit | **denied**, every tool | **denied** where the agent can block them | what was spent, which limit, how to continue | + +The infrastructure policy still decides first: a command the policy blocks stays blocked +whatever the budget says. To continue past a limit: start a new session (session limit), wait +for the next day in `tz` (project limit), or raise the limit in `budget.yaml` and re-sign it. + +| Agent | Tool hook | Prompt hook | Warning shows as | +|---|---|---|---| +| Claude Code | `PreToolUse`, every tool | `UserPromptSubmit` | a message in the transcript | +| Codex CLI | `PreToolUse`, every tool | `UserPromptSubmit` | a UI warning | +| Gemini CLI | `BeforeTool`, every tool | `BeforeAgent` | a message in the terminal | +| OpenCode | plugin, every tool | — | a toast | +| Copilot CLI | `preToolUse`, every tool (premium requests) | — (it cannot block) | stderr only | + +## Where the numbers come from + +No hook payload carries token usage, so Aegis reads each agent's own session log, incrementally: +Claude Code's transcript (and its sub-agents'), Codex rollouts, Gemini CLI chats, OpenCode's +database (read-only; only its session and message tables), and Copilot CLI's events. A +sub-agent is counted inside its parent session. Nothing is sent anywhere. + +**Dollars are estimates** from a dated table of list prices (`aegis_core/budget/prices.yaml`; +Anthropic, OpenAI and Google). On a subscription they are an API-equivalent figure, not your +bill; `unit: tokens` may suit you better there. A model the table does not know is priced, per +category, at the highest rate of its vendor (or of the whole table) — a guess that can still be +low — and flagged; with `on_unknown_price: deny` its sessions are denied until you add a signed +`pricing:` entry for it. + +**The day** is the calendar day in `project_day.tz`, from the time each log record carries. + +**What Aegis keeps** is derived from the logs and can be rebuilt: a cache under +`~/.cache/aegis/budget/`, and a small durable state under `~/.local/state/aegis/budget/` (a +per-session revision counter, locks, and a list of the sessions seen each day, so a rebuild can +tell when history is missing). Deleting the cache resets nothing while the logs exist. + +### When usage cannot be verified + +Some usage can only be counted as a lower bound: a log Aegis cannot find or read, records it +does not recognise, usage with no recorded time, a project day rebuilt after sessions' logs +were deleted, or the first day of a budget when other sessions had already run that day. +`on_unknown_log` decides what that means: + +- `allow` (default): a warning, once; the lower bound is used. +- `deny`: covered calls are denied. A session's own unreadable log stays denied (start a new + session). For an incomplete **project day**, `aegis budget reset --day` acknowledges that + day's history as it stands — it is signed with the policy key, so an agent cannot do it — and + the day is no longer denied for that reason. Otherwise it clears at the next day. + +## Commands + +```bash +aegis budget status [--project DIR] [--agent A] [--json] +aegis budget check [--project DIR] [--agent A --session ID] +aegis budget reset --day [YYYY-MM-DD] [--project DIR] [--key KEY] +``` + +- `status` rebuilds today's project total from the logs, then shows it against the limit, each + session that contributed (today's usage, its session total against the session limit, and + whether it is a lower bound), estimated prices and the policy's warnings. +- `check` exits **0** under budget (a warning included), **3** over, **65** if `budget.yaml` + cannot be used, **66** if there is none: for CI and scripts, including headless agent runs. + With `--agent` and `--session` it also checks that session's own limit. +- `reset --day` writes the signed acknowledgement described above (default: today in the + policy's `tz`); it needs the signing key and refuses `--insecure`. + +## Cost + +Measured with [`scripts/budget_latency.py`](https://github.com/moneytool/aegis-devops/blob/main/scripts/budget_latency.py) +([`results/budget-latency.json`](https://github.com/moneytool/aegis-devops/blob/main/results/budget-latency.json)): +the budget adds about **9 ms** to one hook call (Apple M4 Pro, Python 3.12: 74 ms median with a +budget, 65 ms without), against a design target of 20 ms. Projects without a `budget.yaml` pay +nothing. diff --git a/docs/index.md b/docs/index.md index c004fca..f42b405 100644 --- a/docs/index.md +++ b/docs/index.md @@ -63,6 +63,7 @@ ticket or chat message (prompt injection, context poisoning) cannot become polic - [Writing constraints](constraints.md) — rules, scopes, authority, environments - [Configuration](configuration.md) — config directory, signing, Git sources, `agents.yaml` - [Server-side enforcement](server-side.md) — AWS SCPs and Kubernetes admission policies +- [Session budget cap](budget.md) — per-session and per-day spend limits across agents - [Benchmark](benchmark.md) — Aegis against OPA and LLM self-checks on poisoned rules ## Project diff --git a/scripts/budget_latency.py b/scripts/budget_latency.py new file mode 100644 index 0000000..649079b --- /dev/null +++ b/scripts/budget_latency.py @@ -0,0 +1,113 @@ +"""What the budget cap adds to one ``aegis hook`` call (design v1.0 §4: target +under 20 ms). + +Times ``python -m aegis_core.cli hook claude`` as the agent runs it (a new +process per tool call) on a non-shell tool call, in two otherwise identical +projects made with ``aegis init``: one with a signed ``budget.yaml`` and a +Claude Code transcript to read, one without. Runs are interleaved so both see +the same machine load, after a warm-up run, and with bytecode caching on (a +pip install ships compiled bytecode; ``PYTHONDONTWRITEBYTECODE`` would make +every run recompile aegis and inflate both numbers). + + python scripts/budget_latency.py [--runs 40] [--out results/budget-latency.json] +""" + +from __future__ import annotations + +import argparse +import json +import os +import platform +import statistics +import subprocess +import sys +import tempfile +import time +from datetime import UTC, datetime +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent + + +def _setup(base: Path) -> tuple[dict[str, str], dict[str, str]]: + sys.path.insert(0, str(ROOT / "src")) + from aegis_core.config import init_config_dir + from aegis_core.signing import load_key, sign_file + + home = base / "home" + projects = {"budget": base / "with-budget", "none": base / "without-budget"} + payloads = {} + now = datetime.now(UTC).isoformat().replace("+00:00", "Z") + for name, project in projects.items(): + init_config_dir(project / ".aegis") + if name == "budget": + budget = project / ".aegis" / "budget.yaml" + budget.write_text("principal: admin\nunit: usd\nsession: {limit: 20}\n" + "project_day: {limit: 100}\nagents: [claude]\n") + sign_file(budget, load_key(f"file:{project / '.aegis' / 'example-signing.key'}")) + log = home / ".claude" / "projects" / f"-{name}" / "s.jsonl" + log.parent.mkdir(parents=True, exist_ok=True) + recs = [{"type": "user", "sessionId": "s", "cwd": str(project), "timestamp": now}] + recs += [{"type": "assistant", "sessionId": "s", "requestId": f"r{i}", "timestamp": now, + "message": {"id": f"m{i}", "model": "claude-opus-5-5", + "usage": {"input_tokens": 1000, "output_tokens": 100}}} + for i in range(200)] + log.write_text("".join(json.dumps(r) + "\n" for r in recs)) + payloads[name] = json.dumps({"session_id": "s", "transcript_path": str(log), + "cwd": str(project), "tool_name": "Edit", + "tool_input": {"file_path": "x"}}) + env = {k: v for k, v in os.environ.items() if k not in ( + "PYTHONDONTWRITEBYTECODE", "AEGIS_CONFIG_DIR", "XDG_CACHE_HOME", "XDG_STATE_HOME")} + env.update({"HOME": str(home), "PYTHONPATH": str(ROOT / "src")}) + return payloads, {**env, "_budget": str(projects["budget"]), + "_none": str(projects["none"])} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--runs", type=int, default=40) + ap.add_argument("--out", default=str(ROOT / "results" / "budget-latency.json")) + args = ap.parse_args() + with tempfile.TemporaryDirectory() as tmp: + payloads, env = _setup(Path(tmp)) + cwd = {"budget": env.pop("_budget"), "none": env.pop("_none")} + cmd = [sys.executable, "-m", "aegis_core.cli", "hook", "claude"] + + def run(name: str) -> float: + start = time.perf_counter() + proc = subprocess.run(cmd, input=payloads[name], text=True, capture_output=True, + env=env, cwd=cwd[name], check=False) + if proc.returncode != 0: + raise SystemExit(f"hook failed ({name}): {proc.stderr}") + return (time.perf_counter() - start) * 1000 + + for name in payloads: # warm-up: bytecode, price-table cache, session record + run(name) + samples: dict[str, list[float]] = {"budget": [], "none": []} + for _ in range(args.runs): + for name in samples: + samples[name].append(run(name)) + + def summary(xs: list[float]) -> dict[str, float]: + ordered = sorted(xs) + return {"median_ms": round(statistics.median(xs), 1), + "p95_ms": round(ordered[max(0, int(len(xs) * 0.95) - 1)], 1)} + + with_b, without = summary(samples["budget"]), summary(samples["none"]) + result = { + "measured": datetime.now(UTC).date().isoformat(), + "runs": args.runs, + "python": platform.python_version(), + "platform": f"{platform.system()} {platform.machine()}", + "hook_without_budget": without, + "hook_with_budget": with_b, + "added_median_ms": round(with_b["median_ms"] - without["median_ms"], 1), + "target_ms": 20, + } + Path(args.out).write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/aegis_core/_examples/budget.example.yaml b/src/aegis_core/_examples/budget.example.yaml new file mode 100644 index 0000000..748c2cc --- /dev/null +++ b/src/aegis_core/_examples/budget.example.yaml @@ -0,0 +1,21 @@ +# The session budget cap (docs/budget.md). Copy to budget.yaml, set your own limits, and sign +# it with `aegis sign`; its principal must hold the 'budget' class in authority.yaml. Until +# it is named budget.yaml it does nothing. Then reinstall your agents' hooks (`aegis install +# `) so they see every tool call, and check with: aegis budget status +version: 1 +principal: admin # must hold the 'budget' class in authority.yaml +unit: usd # usd (estimated from list prices) | tokens; required +session: + limit: 20 # per agent session + warn_at: 0.8 # warn once at 80% +project_day: + limit: 100 # every measured session in this project, per day + warn_at: 0.8 + tz: UTC # the day boundary +agents: [claude, codex, gemini, opencode, copilot] # only listed agents are measured +on_unknown_log: allow # allow (warn) | deny: usage Aegis cannot verify +on_unknown_price: estimate # estimate (warn) | deny: models with no price +pricing: {} # signed overrides, USD per million tokens: +# claude-opus-5-5: {input: 4, output: 20, cache_read: 0.2, cache_write: 8} +copilot: + premium_requests: 50 # Copilot CLI reports premium requests, not tokens diff --git a/src/aegis_core/_examples/budget.example.yaml.sig b/src/aegis_core/_examples/budget.example.yaml.sig new file mode 100644 index 0000000..4fec410 --- /dev/null +++ b/src/aegis_core/_examples/budget.example.yaml.sig @@ -0,0 +1,2 @@ +aegis-sig-v1 +d8ad9e56f8548c60028c05c17ce2f46ca190127cbafde00669faef484f3eaea4 diff --git a/src/aegis_core/budget/evaluate.py b/src/aegis_core/budget/evaluate.py index 5bcfe94..821f594 100644 --- a/src/aegis_core/budget/evaluate.py +++ b/src/aegis_core/budget/evaluate.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from typing import Any -from aegis_core.budget.accounting import UNKNOWN_DAY, Snapshot +from aegis_core.budget.accounting import UNKNOWN_DAY, ProjectDay, Snapshot from aegis_core.budget.policy import BudgetPolicy, Limit from aegis_core.budget.pricing import CATEGORIES, PriceTable @@ -77,91 +77,129 @@ def _fmt(unit: str, amount: float) -> str: return f"{amount:,.2f} premium requests" -def evaluate(policy: BudgetPolicy, prices: PriceTable, snap: Snapshot) -> Verdict: - agent = snap.record.agent - if not policy.measures(agent): - return Verdict(ALLOW) - verdict = Verdict(ALLOW) +class _Builder: + """Collects measures, price flags and unverifiable usage into a Verdict.""" - def amount(usage: Mapping[str, Mapping[str, float]]) -> float: - if policy.unit == "tokens": + def __init__(self, policy: BudgetPolicy, prices: PriceTable, providers: Mapping[str, str]): + self.policy, self.prices, self.providers = policy, prices, providers + self.verdict = Verdict(ALLOW) + self.levels: list[str] = [] + self.unknown: list[str] = [] + + def amount(self, usage: Mapping[str, Mapping[str, float]]) -> float: + if self.policy.unit == "tokens": return tokens_of(usage) - cost = prices.cost(usage, snap.providers) + cost = self.prices.cost(usage, self.providers) for model, cats in cost.estimated.items(): - verdict.estimated[model] = sorted(set(verdict.estimated.get(model, [])) | set(cats)) + merged = set(self.verdict.estimated.get(model, [])) | set(cats) + self.verdict.estimated[model] = sorted(merged) return cost.dollars - def measure(name: str, unit: str, used: float, lim: Limit) -> None: - verdict.measures.append(Measure(name, unit, used, lim.limit, lim.warn_at)) + def measure(self, name: str, unit: str, used: float, lim: Limit) -> None: + self.verdict.measures.append(Measure(name, unit, used, lim.limit, lim.warn_at)) + def project_history(self, pd: ProjectDay, exclude: str | None, acknowledged: bool) -> None: + """A project day whose history is incomplete, or whose other + sessions are only a lower bound, is unverifiable -- unless a signed + ``aegis budget reset --day`` acknowledged that day.""" + if acknowledged: + return + if pd.completeness != "complete": + detail = (f"{pd.missing_sessions} session(s) without a log" if pd.history == "ok" + else f"session inventory {pd.history}") + self.unknown.append(f"today's project total is {pd.completeness} ({detail})") + others = pd.unverified(exclude=exclude) + if others: + reasons = sorted({r for rs in others.values() for r in rs}) + who = "other session(s)" if exclude else "session(s)" + self.unknown.append(f"{len(others)} {who} in the project are a lower bound " + f"({', '.join(reasons)})") + + def finish(self) -> Verdict: + verdict, policy = self.verdict, self.policy + for m in verdict.measures: + self.levels.append(m.level) + if m.level == DENY: + verdict.reasons.append(f"{m.name}-over") + verdict.messages.append( + f"budget: {m.name.replace('_', ' ')} limit reached: {_fmt(m.unit, m.used)} " + f"of {_fmt(m.unit, m.limit)}" + ( + f" (overshoot {_fmt(m.unit, m.overshoot)})" if m.overshoot else "")) + elif m.level == WARN: + verdict.reasons.append(f"{m.name}-warn") + verdict.messages.append( + f"budget: {m.name.replace('_', ' ')} at {m.fraction:.0%}: " + f"{_fmt(m.unit, m.used)} of {_fmt(m.unit, m.limit)}, " + f"{_fmt(m.unit, m.limit - m.used)} left") + # Unpriced models: estimated in 'estimate' mode, refused in 'deny' mode. + if policy.unit == "usd" and verdict.estimated: + names = ", ".join(sorted(verdict.estimated)) + if policy.on_unknown_price == "deny": + self.levels.append(DENY) + verdict.reasons.append("unpriced-model") + verdict.messages.append(f"budget: no explicit price for {names}; add it under " + "'pricing:' in budget.yaml and sign it") + else: + verdict.reasons.append("estimated-price") + verdict.messages.append(f"budget: estimated, unpriced model(s): {names}") + # Usage Aegis could not read or place: an unknown log. + if self.unknown: + if policy.on_unknown_log == "deny": + self.levels.append(DENY) + verdict.reasons.append("unknown-log") + verdict.messages.append("budget: usage cannot be verified (" + + "; ".join(self.unknown) + "); on_unknown_log: deny") + else: + verdict.reasons.append("unknown-log-warn") + verdict.messages.append("budget: usage may be under-counted: " + + "; ".join(self.unknown)) + if DENY in self.levels: + verdict.decision = DENY + elif WARN in self.levels or verdict.reasons: + verdict.decision = WARN + return verdict + + +def evaluate(policy: BudgetPolicy, prices: PriceTable, snap: Snapshot, *, + acknowledged: bool = False) -> Verdict: + """The decision for one session at one hook. ``acknowledged``: a signed + ``aegis budget reset --day`` covers today, so an incomplete project + history no longer counts as unverifiable (the session's own log problems + still do).""" + agent = snap.record.agent + if not policy.measures(agent): + return Verdict(ALLOW) + b = _Builder(policy, prices, snap.providers) if agent == "copilot": assert policy.copilot_premium_requests used = sum(u.get("requests", 0) for u in snap.session_usage.values()) - measure("copilot_premium_requests", "requests", used, policy.copilot_premium_requests) + b.measure("copilot_premium_requests", "requests", used, policy.copilot_premium_requests) else: if policy.session: - measure("session", policy.unit, amount(snap.session_usage), policy.session) + b.measure("session", policy.unit, b.amount(snap.session_usage), policy.session) if policy.project_day and snap.project is not None: - measure("project_day", policy.unit, amount(snap.project_usage), policy.project_day) - - levels = [m.level for m in verdict.measures] - for m in verdict.measures: - if m.level == DENY: - verdict.reasons.append(f"{m.name}-over") - verdict.messages.append( - f"budget: {m.name.replace('_', ' ')} limit reached: {_fmt(m.unit, m.used)} of " - f"{_fmt(m.unit, m.limit)}" + (f" (overshoot {_fmt(m.unit, m.overshoot)})" - if m.overshoot else "")) - elif m.level == WARN: - verdict.reasons.append(f"{m.name}-warn") - verdict.messages.append( - f"budget: {m.name.replace('_', ' ')} at {m.fraction:.0%}: {_fmt(m.unit, m.used)} " - f"of {_fmt(m.unit, m.limit)}, {_fmt(m.unit, m.limit - m.used)} left") - - # Unpriced models: estimated in 'estimate' mode, refused in 'deny' mode. - if policy.unit == "usd" and verdict.estimated: - names = ", ".join(sorted(verdict.estimated)) - if policy.on_unknown_price == "deny": - levels.append(DENY) - verdict.reasons.append("unpriced-model") - verdict.messages.append(f"budget: no explicit price for {names}; add it under " - "'pricing:' in budget.yaml and sign it") - else: - verdict.reasons.append("estimated-price") - verdict.messages.append(f"budget: estimated, unpriced model(s): {names}") - - # Usage Aegis could not read or place: an unknown log. - unknown: list[str] = [] + b.measure("project_day", policy.unit, b.amount(snap.project_usage), + policy.project_day) rec = snap.record if rec.log_missing: - unknown.append("the session's log is missing") + b.unknown.append("the session's log is missing") if rec.problem_total: - unknown.append(f"{rec.problem_total} unreadable log record(s)") + b.unknown.append(f"{rec.problem_total} unreadable log record(s)") if rec.buckets().get(UNKNOWN_DAY): - unknown.append("usage with no recorded time (counted for the session only)") + b.unknown.append("usage with no recorded time (counted for the session only)") if snap.project is not None: - pd = snap.project - if pd.completeness != "complete": - detail = (f"{pd.missing_sessions} session(s) without a log" if pd.history == "ok" - else f"session inventory {pd.history}") - unknown.append(f"today's project total is {pd.completeness} ({detail})") - others = pd.unverified(exclude=rec.key) - if others: - reasons = sorted({r for rs in others.values() for r in rs}) - unknown.append(f"{len(others)} other session(s) in the project are a lower bound " - f"({', '.join(reasons)})") - if unknown: - if policy.on_unknown_log == "deny": - levels.append(DENY) - verdict.reasons.append("unknown-log") - verdict.messages.append("budget: usage cannot be verified (" + "; ".join(unknown) - + "); on_unknown_log: deny") - else: - verdict.reasons.append("unknown-log-warn") - verdict.messages.append("budget: usage may be under-counted: " + "; ".join(unknown)) - - if DENY in levels: - verdict.decision = DENY - elif WARN in levels or verdict.reasons: - verdict.decision = WARN - return verdict + b.project_history(snap.project, rec.key, acknowledged) + return b.finish() + + +def evaluate_project(policy: BudgetPolicy, prices: PriceTable, pd: ProjectDay, *, + acknowledged: bool = False) -> Verdict: + """The project/day limit alone, for ``aegis budget check`` without a + session: the day's total against ``project_day``, with every + contributing session's verification status.""" + b = _Builder(policy, prices, pd.providers()) + if policy.project_day: + token_usage = {m: u for m, u in pd.usage().items() if "requests" not in u} + b.measure("project_day", policy.unit, b.amount(token_usage), policy.project_day) + b.project_history(pd, None, acknowledged) + return b.finish() diff --git a/src/aegis_core/budget/gate.py b/src/aegis_core/budget/gate.py index 83e71cf..56029eb 100644 --- a/src/aegis_core/budget/gate.py +++ b/src/aegis_core/budget/gate.py @@ -26,6 +26,7 @@ from aegis_core.budget.logs import AgentLog, default_logs from aegis_core.budget.policy import BudgetPolicy from aegis_core.budget.pricing import PriceTable, load_builtin_table +from aegis_core.budget.report import reset_acknowledged # Hook events that come before a new prompt rather than a tool call. PROMPT_EVENTS = frozenset({"UserPromptSubmit", "BeforeAgent"}) @@ -73,7 +74,8 @@ def _first_time(store: BudgetStore, agent: str, sid: str, keys: list[str]) -> li def check(agent: str, payload: Mapping[str, Any], cwd: str | None, policy: BudgetPolicy, *, store: BudgetStore | None = None, logs: Mapping[str, AgentLog] | None = None, - now: float | None = None) -> Outcome: + now: float | None = None, key: bytes | None = None, + insecure: bool = False) -> Outcome: """The budget decision for one hook call. Raises on internal errors; the caller fails closed inside an opted-in project.""" if not policy.measures(agent): @@ -96,7 +98,9 @@ def check(agent: str, payload: Mapping[str, Any], cwd: str | None, policy: Budge snap = refresh(store, logs, loc, policy.zone, now=time.time() if now is None else now, fallback_cwd=cwd) prices = PriceTable(load_builtin_table(), policy.pricing) - verdict = evaluate(policy, prices, snap) + acknowledged = bool(snap.record.project_root) and reset_acknowledged( + store, snap.record.project_root, snap.day, key, insecure) + verdict = evaluate(policy, prices, snap, acknowledged=acknowledged) if verdict.decision == DENY: over = [m.name for m in verdict.measures if m.level == DENY] ways = sorted({_HOW_TO_CONTINUE[n] for n in over}) diff --git a/src/aegis_core/budget/pricing.py b/src/aegis_core/budget/pricing.py index 9561dee..213b2c9 100644 --- a/src/aegis_core/budget/pricing.py +++ b/src/aegis_core/budget/pricing.py @@ -14,11 +14,16 @@ from __future__ import annotations +import contextlib +import hashlib +import json import math +import os import re from collections.abc import Mapping from dataclasses import dataclass, field from importlib import resources +from pathlib import Path from typing import Any import yaml @@ -161,9 +166,29 @@ def cost(self, usage_by_model: Mapping[str, Mapping[str, int]], return total +def _table_cache_dir() -> Path: + home = Path(os.environ.get("HOME") or Path.home()) + return Path(os.environ.get("XDG_CACHE_HOME") or home / ".cache") / "aegis" / "budget" + + def load_builtin_table() -> dict[str, Any]: + """The built-in price table. Parsing it with PyYAML's pure-Python loader + costs ~15 ms, paid on every hook call, so the parsed table is cached as + JSON under the cache directory, keyed by a hash of the YAML text: an + edited table is re-parsed, and a damaged or missing cache only costs a + re-parse.""" text = resources.files("aegis_core.budget").joinpath("prices.yaml").read_text() - table = yaml.safe_load(text) + digest = hashlib.sha256(text.encode()).hexdigest()[:16] + cache = _table_cache_dir() / f"prices-{digest}.json" + try: + table = json.loads(cache.read_text()) + except (OSError, ValueError): + table = yaml.safe_load(text) + with contextlib.suppress(OSError): + cache.parent.mkdir(parents=True, exist_ok=True) + tmp = cache.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text(json.dumps(table)) + os.replace(tmp, cache) if not isinstance(table, dict) or table.get("version") != 1: raise ValueError("aegis_core/budget/prices.yaml: unsupported shape or version") return table diff --git a/src/aegis_core/budget/report.py b/src/aegis_core/budget/report.py new file mode 100644 index 0000000..28d0856 --- /dev/null +++ b/src/aegis_core/budget/report.py @@ -0,0 +1,187 @@ +"""``aegis budget status|check|reset`` (design v1.0 §7, §4.1). + +``status`` and ``check`` rebuild today's project total from the logs before +reporting, so they show the usage the logs hold now, not only what hooks +have seen. ``reset --day`` is the strict-recovery path of §4.1: it writes a +marker **signed with the policy key** that acknowledges one day's history +as it stands, so an incomplete or unverifiable project total stops denying +under ``on_unknown_log: deny``. An agent without the key cannot write one. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import time +from collections.abc import Mapping +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +from aegis_core.budget.accounting import BudgetStore, day_of, rebuild +from aegis_core.budget.evaluate import ALLOW, DENY, Measure, evaluate_project, tokens_of +from aegis_core.budget.logs import AgentLog +from aegis_core.budget.policy import BudgetPolicy +from aegis_core.budget.pricing import PriceTable + +RESET_HEADER = "aegis-budget-reset-v1" + + +# --- signed reset markers ------------------------------------------------------------ + + +def _mac(data: bytes, key: bytes) -> str: + return hashlib.blake2b(data, key=key, digest_size=32).hexdigest() + + +def _canonical(fields: Mapping[str, Any]) -> bytes: + return json.dumps(fields, sort_keys=True, separators=(",", ":")).encode() + + +def reset_path(store: BudgetStore, root: str, day: str) -> Path: + return store.project_dir(root) / f"reset-{day}.json" + + +def write_reset(store: BudgetStore, root: str, day: str, key: bytes, + now: float | None = None) -> Path: + """Acknowledges ``day``'s project history as it stands, signed with + ``key`` (the policy signing key).""" + fields = {"header": RESET_HEADER, "root": root, "day": day, + "at": datetime.fromtimestamp(time.time() if now is None else now, + tz=UTC).isoformat()} + path = reset_path(store, root, day) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({**fields, "mac": _mac(_canonical(fields), key)}, indent=2)) + return path + + +def reset_acknowledged(store: BudgetStore, root: str, day: str, key: bytes | None, + insecure: bool = False) -> bool: + """Whether a valid reset marker covers ``day`` in project ``root``. + Without a key (``--insecure``) nothing is verified, so a marker's + presence is taken as it is, like every other policy file then.""" + try: + doc = json.loads(reset_path(store, root, day).read_text()) + except (OSError, ValueError): + return False + if not isinstance(doc, dict): + return False + fields = {k: doc.get(k) for k in ("header", "root", "day", "at")} + if fields["header"] != RESET_HEADER or fields["root"] != root or fields["day"] != day: + return False + if key is None: + return insecure + mac = doc.get("mac") + return isinstance(mac, str) and hmac.compare_digest(mac, _mac(_canonical(fields), key)) + + +# --- status / check -------------------------------------------------------------------- + + +def _amount(policy: BudgetPolicy, prices: PriceTable, usage: Mapping[str, Mapping[str, float]], + providers: Mapping[str, str]) -> tuple[float, dict[str, list[str]]]: + if policy.unit == "tokens": + return tokens_of(usage), {} + cost = prices.cost(usage, providers) + return cost.dollars, cost.estimated + + +def project_status(policy: BudgetPolicy, prices: PriceTable, store: BudgetStore, + logs: Mapping[str, AgentLog], root: str, *, now: float | None = None, + key: bytes | None = None, insecure: bool = False, + agent: str | None = None) -> dict[str, Any]: + """Today's project total (rebuilt from the logs) and each session that + contributed to it, as a JSON-ready mapping.""" + now = time.time() if now is None else now + zone = policy.zone + day = day_of(datetime.fromtimestamp(now, tz=UTC), zone) + measured = {a: lg for a, lg in logs.items() if policy.measures(a)} + pd = rebuild(store, measured, root, day, zone) + acknowledged = reset_acknowledged(store, root, day, key, insecure) + verdict = evaluate_project(policy, prices, pd, acknowledged=acknowledged) + + sessions = [] + for entry_key, entry in sorted(pd.entries.items()): + a, _, sid = entry_key.partition("/") + if agent and a != agent: + continue + rec = store.load_record(a, sid) + total_usage = rec.usage_total() if rec else entry.get("usage", {}) + providers = {**entry.get("providers", {}), **(rec.providers if rec else {})} + if a == "copilot": + unit = "requests" + today = sum(u.get("requests", 0) for u in entry.get("usage", {}).values()) + total = sum(u.get("requests", 0) for u in total_usage.values()) + limit = policy.copilot_premium_requests + estimated: dict[str, list[str]] = {} + else: + unit = policy.unit + today, _ = _amount(policy, prices, entry.get("usage", {}), providers) + total, estimated = _amount(policy, prices, total_usage, providers) + limit = policy.session + m = Measure("session", unit, total, limit.limit, limit.warn_at) if limit else None + sessions.append({ + "agent": a, "session_id": sid, "unit": unit, "today": today, "total": total, + "limit": limit.limit if limit else None, + "level": m.level if m else ALLOW, "overshoot": m.overshoot if m else 0.0, + "unverified": entry.get("unverified", []), "estimated": estimated, + }) + project = next((m.to_dict() for m in verdict.measures if m.name == "project_day"), None) + return { + "project": root, "day": day, "tz": policy.tz, "unit": policy.unit, + "project_day": project, "completeness": pd.completeness, "history": pd.history, + "missing_sessions": pd.missing_sessions, "acknowledged": acknowledged, + "decision": verdict.decision, "messages": verdict.messages, + "estimated": verdict.estimated, "sessions": sessions, + "warnings": list(policy.warnings), + } + + +def _fmt(unit: str, amount: float | None) -> str: + if amount is None: + return "—" + if unit == "usd": + return f"${amount:,.2f}" + if unit == "tokens": + return f"{amount:,.0f} tokens" + return f"{amount:,.2f} premium requests" + + +def render_status(status: Mapping[str, Any]) -> str: + """``aegis budget status`` for people.""" + unit = status["unit"] + lines = [f"budget: {status['project']} — {status['day']} ({status['tz']}), unit {unit}"] + pdm = status["project_day"] + if pdm: + lines.append(f" project today: {_fmt(unit, pdm['used'])} of {_fmt(unit, pdm['limit'])} " + f"({pdm['fraction']:.0%}, {pdm['level']})" + + (f", overshoot {_fmt(unit, pdm['overshoot'])}" if pdm["overshoot"] + else "")) + else: + lines.append(" project today: no project_day limit") + if status["completeness"] != "complete": + note = " — acknowledged by a signed reset" if status["acknowledged"] else "" + lines.append(f" history: {status['completeness']} (inventory {status['history']}, " + f"{status['missing_sessions']} session(s) without a log){note}") + if not status["sessions"]: + lines.append(" no sessions today") + for s in status["sessions"]: + flag = f" [{s['level']}]" if s["level"] != ALLOW else "" + lower = f" lower bound: {', '.join(s['unverified'])}" if s["unverified"] else "" + lines.append(f" {s['agent']:9} {s['session_id']}: today {_fmt(s['unit'], s['today'])}, " + f"session {_fmt(s['unit'], s['total'])} of {_fmt(s['unit'], s['limit'])}" + f"{flag}{lower}") + if status["estimated"]: + lines.append(" estimated prices for: " + ", ".join(sorted(status["estimated"]))) + for w in status["warnings"]: + lines.append(f" warning: {w}") + return "\n".join(lines) + + +EXIT_UNDER, EXIT_OVER = 0, 3 + + +def check_exit(decision: str) -> int: + """``aegis budget check``: 0 under budget (a warning included), 3 over.""" + return EXIT_OVER if decision == DENY else EXIT_UNDER diff --git a/src/aegis_core/cli.py b/src/aegis_core/cli.py index 3d54f71..12bc12a 100644 --- a/src/aegis_core/cli.py +++ b/src/aegis_core/cli.py @@ -82,6 +82,7 @@ import logging import os import sys +import time from dataclasses import asdict from datetime import UTC, datetime, timedelta from pathlib import Path @@ -106,6 +107,7 @@ load_repos, load_signers, ) +from aegis_core.hook import BUDGET_AGENTS from aegis_core.identity import AGENTS_FILE, load_identity_model from aegis_core.intent import InfrastructureIntent from aegis_core.interceptor import AegisInterceptor, Decision @@ -639,6 +641,38 @@ def _build_parser() -> argparse.ArgumentParser: "directory that applies holds a budget.yaml)", ) + budget_parser = subparsers.add_parser( + "budget", help="The session budget cap: status, check, and the signed day reset " + "(docs/budget.md)") + budget_sub = budget_parser.add_subparsers(dest="budget_command", required=True) + + def _budget_common(p): + p.add_argument("--project", default=".", + help="project directory (default: .); its policy directory is found as " + "the hook finds it") + p.add_argument("--key", help="signing key (file:PATH, env:VAR or hex)") + p.add_argument("--insecure", action="store_true", + help="do not verify signatures (not accepted by reset)") + + status_p = budget_sub.add_parser( + "status", help="today's project total and each session in it, rebuilt from the logs") + _budget_common(status_p) + status_p.add_argument("--agent", choices=sorted(BUDGET_AGENTS)) + status_p.add_argument("--json", action="store_true", help="machine-readable output") + check_p = budget_sub.add_parser( + "check", help="exit 0 under budget, 3 over (65: budget.yaml unusable, 66: no budget)") + _budget_common(check_p) + check_p.add_argument("--agent", choices=sorted(BUDGET_AGENTS), + help="with --session: also check that session's limit") + check_p.add_argument("--session", help="a session id of --agent") + reset_p = budget_sub.add_parser( + "reset", help="acknowledge one day's incomplete history (signed with the policy key), " + "so on_unknown_log: deny stops denying for it") + _budget_common(reset_p) + reset_p.add_argument("--day", nargs="?", const="today", required=True, + help="the day to acknowledge, YYYY-MM-DD (default: today in the " + "policy's tz)") + return parser @@ -1348,19 +1382,16 @@ def _load_identity(args: argparse.Namespace): return model, [*key_warnings, *authority_map.warnings, *model.warnings] -def load_budget_for_hook(config_dir: Path, extra_args: list[str] | None = None): - """``budget.yaml`` in an opted-in policy directory, verified like every - other policy file (signature, and its principal must hold ``budget``), - for ``aegis hook``. ``None`` when the directory has no budget. The key - is resolved as for ``aegis check`` (``--key``/``--insecure`` among the - hook's extra arguments, ``$AEGIS_SIGNING_KEY``, the example key next to - the constraints). Raises when the budget cannot be used; the hook fails - closed on that.""" +def load_budget_with_key(config_dir: Path, extra_args: list[str] | None = None): + """``(policy, key, insecure)`` for ``budget.yaml`` in a policy directory, + verified like every other policy file (signature, and its principal must + hold ``budget``); ``(None, key, insecure)`` when there is no budget. The + key is resolved as for ``aegis check`` (``--key``/``--insecure`` among + ``extra_args``, ``$AEGIS_SIGNING_KEY``, the example key next to the + constraints). Raises when the budget cannot be used.""" from aegis_core.budget.policy import BUDGET_FILE, load_budget_policy path = Path(config_dir) / BUDGET_FILE - if not path.exists(): - return None opts = argparse.ArgumentParser(add_help=False) opts.add_argument("--key") opts.add_argument("--insecure", action="store_true") @@ -1370,8 +1401,100 @@ def load_budget_for_hook(config_dir: Path, extra_args: list[str] | None = None): insecure=known.insecure) _resolve_config_paths(args) key, insecure = _resolve_key(args, []) + if not path.exists(): + return None, key, insecure authority_map = load_authority_map(args.authority, key=key, insecure=insecure) - return load_budget_policy(path, authority_map=authority_map, key=key, insecure=insecure) + policy = load_budget_policy(path, authority_map=authority_map, key=key, insecure=insecure) + return policy, key, insecure + + +def _run_budget(args: argparse.Namespace) -> int: + """``aegis budget status|check|reset`` (docs/budget.md).""" + from aegis_core.budget import BudgetStore, default_logs, evaluate, price_table, refresh + from aegis_core.budget.accounting import day_of, find_project_root + from aegis_core.budget.report import ( + check_exit, + project_status, + render_status, + reset_acknowledged, + write_reset, + ) + + project = Path(args.project).resolve() + config_dir = hook_module.find_opt_in_config(project) + if config_dir is None: + raise ConfigNotFoundError(f"no policy directory for {project} (.aegis/ here or above, " + "$AEGIS_CONFIG_DIR, ~/.config/aegis)") + extra = [*(["--key", args.key] if args.key else []), *(["--insecure"] if args.insecure + else [])] + policy, key, insecure = load_budget_with_key(config_dir, extra) + if policy is None: + raise ConfigNotFoundError(f"no budget: {config_dir / 'budget.yaml'} does not exist") + root = find_project_root(str(project)) + store, logs, prices = BudgetStore(), default_logs(), price_table(policy) + now = time.time() + today = day_of(datetime.fromtimestamp(now, tz=UTC), policy.zone) + + if args.budget_command == "reset": + if key is None: + raise UsageError("budget reset must be signed with the policy key; --insecure is " + "not accepted") + if root is None: + raise DataError(f"{project} is not inside a project with .aegis/: project/day " + "totals (and their history) exist only there") + day = today if args.day == "today" else args.day + try: + datetime.strptime(day, "%Y-%m-%d") + except ValueError: + raise UsageError(f"--day {day!r} is not YYYY-MM-DD") from None + path = write_reset(store, root, day, key, now) + print(f"aegis budget: history for {day} acknowledged for {root} ({path})") + return 0 + + if args.budget_command == "status": + if root is None: + print(f"aegis budget: {project} is not inside a project with .aegis/; only session " + "limits apply here") + return 0 + status = project_status(policy, prices, store, logs, root, now=now, key=key, + insecure=insecure, agent=args.agent) + print(json.dumps(status, indent=2, sort_keys=True) if args.json + else render_status(status)) + return 0 + + # check + if bool(args.agent) != bool(args.session): + raise UsageError("--agent and --session go together") + if args.session: + if not policy.measures(args.agent): + print(f"aegis budget: {args.agent} is not measured by this budget") + return 0 + loc = logs[args.agent].locate(args.session) + if loc is None: + raise DataError(f"cannot find the log of {args.agent} session {args.session}") + measured = {a: lg for a, lg in logs.items() if policy.measures(a)} + snap = refresh(store, measured, loc, policy.zone, now=now, fallback_cwd=str(project)) + acknowledged = bool(snap.record.project_root) and reset_acknowledged( + store, snap.record.project_root, snap.day, key, insecure) + verdict = evaluate(policy, prices, snap, acknowledged=acknowledged) + decision, messages = verdict.decision, verdict.messages + elif root is None: + print(f"aegis budget: {project} is not inside a project with .aegis/; pass --agent " + "and --session to check a session") + return 0 + else: + status = project_status(policy, prices, store, logs, root, now=now, key=key, + insecure=insecure) + decision, messages = status["decision"], status["messages"] + for m in messages: + print(f"aegis {m}") + print(f"aegis budget: {'over' if decision == 'deny' else 'under'} budget") + return check_exit(decision) + + +def load_budget_for_hook(config_dir: Path, extra_args: list[str] | None = None): + """The budget policy alone (see :func:`load_budget_with_key`).""" + return load_budget_with_key(config_dir, extra_args)[0] def _run_audit_aws(args: argparse.Namespace) -> int: @@ -1675,6 +1798,8 @@ def _run(argv: list[str]) -> int: escalate_as=args.escalate_as, extra_args=_strip_leading_separator(args.check_args), ) + if args.command == "budget": + return _run_budget(args) if args.command == "install": return hook_module.main_install( args.agent, user=args.user, project=args.project, remove=args.remove, diff --git a/src/aegis_core/hook.py b/src/aegis_core/hook.py index 0cc2d7e..9db9ed8 100644 --- a/src/aegis_core/hook.py +++ b/src/aegis_core/hook.py @@ -369,14 +369,14 @@ def _budget_outcome(agent: str, payload: dict, request: HookRequest, config_dir: from aegis_core.budget.gate import Outcome, check try: - policy = cli.load_budget_for_hook(config_dir, extra_args) + policy, key, insecure = cli.load_budget_with_key(config_dir, extra_args) except Exception as exc: # noqa: BLE001 - any failure to load is reported, not raised return Outcome("deny", f"aegis budget: {config_dir / BUDGET_FILE} cannot be used " f"({exc.__class__.__name__}: {exc}); fix and re-sign it, or " "remove it") if policy is None: return Outcome("allow") - return check(agent, payload, request.cwd, policy) + return check(agent, payload, request.cwd, policy, key=key, insecure=insecure) def run_hook(agent: str, stdin_text: str, *, escalate_as: str = "ask", diff --git a/tests/conftest.py b/tests/conftest.py index e354497..afce315 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,3 +15,17 @@ REPO_ROOT = Path(__file__).resolve().parent.parent os.chdir(REPO_ROOT) + + +import pytest # noqa: E402 + + +@pytest.fixture(autouse=True) +def _private_budget_dirs(tmp_path_factory, monkeypatch): + """The budget cap writes a cache and a state directory (default under + ~/.cache and ~/.local/state). Point both at a per-test temporary + directory so no test touches the real ones; tests that set HOME or the + XDG variables themselves still override this.""" + base = tmp_path_factory.mktemp("xdg") + monkeypatch.setenv("XDG_CACHE_HOME", str(base / "cache")) + monkeypatch.setenv("XDG_STATE_HOME", str(base / "state")) diff --git a/tests/test_budget_cli.py b/tests/test_budget_cli.py new file mode 100644 index 0000000..553b8e4 --- /dev/null +++ b/tests/test_budget_cli.py @@ -0,0 +1,171 @@ +"""``aegis budget status|check|reset`` (design v1.0 §7, §4.1).""" + +import json +from datetime import UTC, datetime +from pathlib import Path + +import pytest +from test_budget_hook import BUDGET, env, write_budget # noqa: F401 (fixture) + +from aegis_core.budget.accounting import BudgetStore +from aegis_core.budget.report import reset_acknowledged, reset_path +from aegis_core.cli import main +from aegis_core.hook import run_hook + +NOW = datetime.now(UTC) +TODAY = NOW.date().isoformat() + + +def transcript(env, sid: str, input_tokens: int, *, malformed: bool = False) -> Path: # noqa: F811 + """A Claude Code session in the project, with usage timestamped now.""" + path = env["home"] / ".claude" / "projects" / "-project" / f"{sid}.jsonl" + path.parent.mkdir(parents=True, exist_ok=True) + ts = NOW.isoformat().replace("+00:00", "Z") + recs = [{"type": "user", "sessionId": sid, "cwd": str(env["project"]), "timestamp": ts}, + {"type": "assistant", "sessionId": sid, "requestId": f"r-{sid}", "timestamp": ts, + "message": {"id": f"m-{sid}", "model": "claude-opus-5-5", + "usage": {"input_tokens": input_tokens, "output_tokens": 0}}}] + text = "".join(json.dumps(r) + "\n" for r in recs) + path.write_text(text + ("{broken\n" if malformed else "")) + return path + + +def hook_call(env, sid: str): # noqa: F811 + payload = {"session_id": sid, "cwd": str(env["project"]), "tool_name": "Edit", + "tool_input": {}, "transcript_path": str( + env["home"] / ".claude" / "projects" / "-project" / f"{sid}.jsonl")} + return run_hook("claude", json.dumps(payload)) + + +DAY_BUDGET = "principal: admin\nunit: usd\nsession: {limit: 1.0}\nproject_day: {limit: 2.0}\n" + + +def _status(env, capsys, *extra) -> dict: # noqa: F811 + assert main(["budget", "status", "--project", str(env["project"]), "--json", *extra]) == 0 + return json.loads(capsys.readouterr().out) + + +# --- status -------------------------------------------------------------------------------- + + +def test_status_rebuilds_today_from_the_logs(env, capsys): # noqa: F811 + write_budget(env, DAY_BUDGET) + transcript(env, "a", 125_000) # $0.50 + transcript(env, "b", 375_000) # $1.50: over its session limit + st = _status(env, capsys) + assert st["day"] == TODAY and st["unit"] == "usd" + assert st["project_day"]["used"] == pytest.approx(2.0) + assert st["project_day"]["level"] == "deny" + by_id = {s["session_id"]: s for s in st["sessions"]} + assert by_id["a"]["today"] == pytest.approx(0.5) and by_id["a"]["level"] == "allow" + assert by_id["b"]["level"] == "deny" and by_id["b"]["overshoot"] == pytest.approx(0.5) + assert st["decision"] == "deny" + + +def test_status_for_people_and_agent_filter(env, capsys): # noqa: F811 + write_budget(env, DAY_BUDGET) + transcript(env, "a", 125_000) + assert main(["budget", "status", "--project", str(env["project"])]) == 0 + out = capsys.readouterr().out + assert "project today: $0.50 of $2.00 (25%, allow)" in out + assert "claude a: today $0.50, session $0.50 of $1.00" in out + assert _status(env, capsys, "--agent", "codex")["sessions"] == [] + + +def test_status_shows_lower_bounds(env, capsys): # noqa: F811 + write_budget(env, DAY_BUDGET) + transcript(env, "a", 1_000, malformed=True) + st = _status(env, capsys) + assert st["sessions"][0]["unverified"] == ["unreadable-records"] + + +# --- check --------------------------------------------------------------------------------- + + +def _check(env, *extra) -> int: # noqa: F811 + return main(["budget", "check", "--project", str(env["project"]), *extra]) + + +def test_check_exit_codes(env, capsys, tmp_path): # noqa: F811 + assert _check(env) == 66 # no budget.yaml + write_budget(env, DAY_BUDGET) + transcript(env, "a", 125_000) + assert _check(env) == 0 + transcript(env, "b", 500_000) # project now $2.50 of $2.00 + assert _check(env) == 3 + assert "project day limit reached" in capsys.readouterr().out + assert _check(env, "--agent", "claude", "--session", "a") == 3 # project limit applies + assert _check(env, "--agent", "claude") == 64 # --session missing + (env["project"] / ".aegis" / "budget.yaml").write_text(DAY_BUDGET.replace("2.0", "20")) + assert _check(env) == 65 # edited after signing + + +def test_check_a_session_against_its_own_limit(env): # noqa: F811 + write_budget(env, "principal: admin\nunit: usd\nsession: {limit: 1.0}\n") + transcript(env, "a", 125_000) + transcript(env, "b", 300_000) + assert _check(env, "--agent", "claude", "--session", "a") == 0 + assert _check(env, "--agent", "claude", "--session", "b") == 3 + assert _check(env, "--agent", "claude", "--session", "nobody") == 65 + + +# --- reset --------------------------------------------------------------------------------- + + +STRICT = "principal: admin\nunit: usd\nsession: {limit: 100}\nproject_day: {limit: 100}\n" \ + "on_unknown_log: deny\n" + + +def test_a_signed_reset_ends_a_strict_unknown_history_day(env, capsys): # noqa: F811 + """First use with other sessions already in today's logs: the history is + unknown, so strict mode denies -- until a signed reset acknowledges it.""" + write_budget(env, STRICT) + transcript(env, "earlier", 1_000) # ran before the budget existed + transcript(env, "now", 1_000) + code, _, err = hook_call(env, "now") + assert code == 2 and "unknown" in err and "on_unknown_log: deny" in err + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + assert TODAY in capsys.readouterr().out + assert hook_call(env, "now")[0] == 0 + st = _status(env, capsys) + assert st["acknowledged"] is True and st["decision"] == "allow" + + +def test_reset_does_not_excuse_the_sessions_own_problems(env): # noqa: F811 + write_budget(env, STRICT) + transcript(env, "now", 1_000, malformed=True) + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + code, _, err = hook_call(env, "now") + assert code == 2 and "unreadable log record" in err + + +def test_reset_markers_are_verified(env, capsys): # noqa: F811 + write_budget(env, STRICT) + transcript(env, "earlier", 1_000) + transcript(env, "now", 1_000) + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + store = BudgetStore() + root = str(env["project"].resolve()) + marker = reset_path(store, root, TODAY) + key = bytes.fromhex("".join( + line.strip() for line in (env["project"] / ".aegis" / "example-signing.key") + .read_text().splitlines() if not line.lstrip().startswith("#"))) + assert reset_acknowledged(store, root, TODAY, key) + doc = json.loads(marker.read_text()) + marker.write_text(json.dumps({**doc, "day": "2001-01-01"})) # moved to another day + assert not reset_acknowledged(store, root, TODAY, key) + marker.write_text(json.dumps({**doc, "mac": "0" * 64})) # forged without the key + assert not reset_acknowledged(store, root, TODAY, key) + assert hook_call(env, "now")[0] == 2 + assert not reset_acknowledged(store, root, TODAY, None) # no key: not verified + assert reset_acknowledged(store, root, TODAY, None, insecure=True) is True + + +def test_reset_needs_the_key_and_a_project(env, capsys, tmp_path): # noqa: F811 + write_budget(env, STRICT) + assert main(["budget", "reset", "--day", "--insecure", + "--project", str(env["project"])]) == 64 + assert main(["budget", "reset", "--day", "2026-13-40", + "--project", str(env["project"])]) == 64 + assert main(["budget", "reset", "--day", "2026-10-01", + "--project", str(env["project"])]) == 0 diff --git a/tests/test_budget_policy.py b/tests/test_budget_policy.py index 20a14d7..e2d7a32 100644 --- a/tests/test_budget_policy.py +++ b/tests/test_budget_policy.py @@ -1,5 +1,6 @@ """budget.yaml loading (design v1.0 §3) and the price table (§5).""" +import json from pathlib import Path import pytest @@ -166,8 +167,8 @@ def test_unsigned_without_key_is_recorded(tmp_path): # --- prices --------------------------------------------------------------------------- -@pytest.fixture(scope="module") -def table(): +@pytest.fixture +def table(): # function scope: the cache dir is redirected per test (conftest) return load_builtin_table() @@ -264,3 +265,31 @@ def test_an_override_reprices_earlier_estimated_usage(table): after = PriceTable(table, {"claude-new-1": {"input": 1, "output": 2, "cache_read": 0, "cache_write": 0}}).cost(usage) assert after.estimated == {} and after.dollars == pytest.approx(4.0) + + +def test_the_parsed_price_table_is_cached_by_content(tmp_path, monkeypatch): + """Parsing prices.yaml costs ~15 ms on every hook call; the parsed table + is cached as JSON keyed by a hash of the YAML text.""" + import aegis_core.budget.pricing as pricing + + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + first = load_builtin_table() + [cache] = (tmp_path / "aegis" / "budget").glob("prices-*.json") + assert json.loads(cache.read_text()) == first + cache.write_text("{damaged") # a damaged cache re-parses + assert load_builtin_table() == first + assert json.loads(cache.read_text()) == first + real = pricing.resources.files + + class Edited: + def __init__(self, pkg): + self.pkg = pkg + + def joinpath(self, name): + text = real(self.pkg).joinpath(name).read_text().replace("as_of: \"2026-10-07\"", + "as_of: \"2027-01-01\"") + return type("T", (), {"read_text": lambda self: text})() + + monkeypatch.setattr(pricing.resources, "files", Edited) + assert load_builtin_table()["as_of"] == "2027-01-01" # an edited table is re-read + assert len(list((tmp_path / "aegis" / "budget").glob("prices-*.json"))) == 2 diff --git a/tests/test_signing.py b/tests/test_signing.py index 5b19020..1dc1302 100644 --- a/tests/test_signing.py +++ b/tests/test_signing.py @@ -178,6 +178,7 @@ def test_shipped_example_and_corpus_files_verify_under_the_example_key(): "data/repos.example.yaml", "data/signers.example.yaml", "data/agents.example.yaml", + "data/budget.example.yaml", "data/sources/PRINCIPALS.yaml", "data/sources/jira-1001.json", "data/corpus/constraints.yaml", @@ -191,4 +192,4 @@ def test_shipped_example_and_corpus_files_verify_under_the_example_key(): ): assert verify_file(path, key), path assert not any(p.name.endswith(".json.sig") for p in Path("data").rglob("*.sig")) - assert len(list(Path("data").rglob("*.sig"))) <= 12 + assert len(list(Path("data").rglob("*.sig"))) <= 13 From ff412e2b727d6db777c42830126b32b77a996794 Mon Sep 17 00:00:00 2001 From: moneytools Date: Fri, 9 Oct 2026 21:01:12 -0500 Subject: [PATCH 2/5] chore: publish results/budget-latency.json (cited by docs/budget.md) --- .gitignore | 1 + results/budget-latency.json | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) create mode 100644 results/budget-latency.json diff --git a/.gitignore b/.gitignore index b71a750..30677d4 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ results/* !results/benchmark.json !results/latency.md !results/latency.json +!results/budget-latency.json !results/llm-external.md !results/llm-cache-naive.jsonl !results/llm-cache-aware.jsonl diff --git a/results/budget-latency.json b/results/budget-latency.json new file mode 100644 index 0000000..870947c --- /dev/null +++ b/results/budget-latency.json @@ -0,0 +1,16 @@ +{ + "measured": "2026-10-10", + "runs": 40, + "python": "3.12.12", + "platform": "Darwin arm64", + "hook_without_budget": { + "median_ms": 65.3, + "p95_ms": 68.0 + }, + "hook_with_budget": { + "median_ms": 74.1, + "p95_ms": 77.6 + }, + "added_median_ms": 8.8, + "target_ms": 20 +} From 0518c271420ee99211b5f46180b1fb6488397226 Mon Sep 17 00:00:00 2001 From: moneytools Date: Fri, 9 Oct 2026 21:15:51 -0500 Subject: [PATCH 3/5] fix(budget): prices only from the package; session check rebuilds the day; resets bind to a baseline - the price table ships as prices.json generated from prices.yaml (scripts/sync_prices.py, test keeps them equal); no writable cache can act as an unsigned price override - check --agent/--session rebuilds the project day from every measured agent before evaluating, as status and an unscoped check do - reset markers sign the acknowledged verification state; only that is excused, later unverifiable usage counts again Review of #40. --- CHANGELOG.md | 10 +- docs/budget.md | 15 +- pyproject.toml | 2 +- scripts/sync_prices.py | 34 ++++ src/aegis_core/budget/evaluate.py | 64 +++++-- src/aegis_core/budget/gate.py | 6 +- src/aegis_core/budget/prices.json | 282 ++++++++++++++++++++++++++++++ src/aegis_core/budget/pricing.py | 41 ++--- src/aegis_core/budget/report.py | 56 ++++-- src/aegis_core/cli.py | 21 ++- tests/test_budget_cli.py | 50 +++++- tests/test_budget_policy.py | 57 +++--- 12 files changed, 529 insertions(+), 109 deletions(-) create mode 100644 scripts/sync_prices.py create mode 100644 src/aegis_core/budget/prices.json diff --git a/CHANGELOG.md b/CHANGELOG.md index f576fb5..c25c6b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -40,12 +40,14 @@ test on each agent are in. CLI). A `budget.yaml` that cannot be used (bad signature, shape, authority) fails closed. - `aegis budget status|check|reset` (design §7): today's project total and its sessions, rebuilt from the logs (`--json`); `check` exits 0 under, 3 over, 65 for an unusable `budget.yaml`, 66 - for none; `reset --day` acknowledges an incomplete day's history with a marker signed by the - policy key, so `on_unknown_log: deny` stops denying it (an agent without the key cannot). + for none; `reset --day` acknowledges an incomplete day's history as it stands, with a marker + signed by the policy key that records the acknowledged state, so `on_unknown_log: deny` stops + denying for exactly that (later problems count again; an agent without the key cannot reset). - `docs/budget.md`; `aegis init` writes a signed `budget.example.yaml`. - `scripts/budget_latency.py`: the budget adds ~9 ms to a hook call (74 vs 65 ms median, - `results/budget-latency.json`), under the 20 ms target. The parsed price table is cached by - content hash; parsing it on every call had cost ~15 ms. + `results/budget-latency.json`), under the 20 ms target. The price table ships as + `prices.json` generated from `prices.yaml` (`scripts/sync_prices.py`; a test keeps them equal) + and is read only from the package: parsing the YAML on every call had cost ~15 ms. - `aegis install` adds the budget hooks when the policy directory that applies holds a `budget.yaml` (`--budget` / `--no-budget` to choose): the tool hook sees every tool call and the prompt hook is added where the agent has one. The Claude Code plugin and the Gemini diff --git a/docs/budget.md b/docs/budget.md index 237e25c..77e3ded 100644 --- a/docs/budget.md +++ b/docs/budget.md @@ -125,8 +125,11 @@ were deleted, or the first day of a budget when other sessions had already run t - `allow` (default): a warning, once; the lower bound is used. - `deny`: covered calls are denied. A session's own unreadable log stays denied (start a new session). For an incomplete **project day**, `aegis budget reset --day` acknowledges that - day's history as it stands — it is signed with the policy key, so an agent cannot do it — and - the day is no longer denied for that reason. Otherwise it clears at the next day. + day's history **as it stands**: it records the day's verification state at that moment (how + complete the history is, which sessions are lower bounds and why) and is signed with the + policy key, so an agent cannot write one. Exactly that is excused; anything that becomes + unverifiable afterwards — a new session with an unreadable log, more sessions without one — + counts again. Otherwise it clears at the next day. ## Commands @@ -141,9 +144,11 @@ aegis budget reset --day [YYYY-MM-DD] [--project DIR] [--key KEY] whether it is a lower bound), estimated prices and the policy's warnings. - `check` exits **0** under budget (a warning included), **3** over, **65** if `budget.yaml` cannot be used, **66** if there is none: for CI and scripts, including headless agent runs. - With `--agent` and `--session` it also checks that session's own limit. -- `reset --day` writes the signed acknowledgement described above (default: today in the - policy's `tz`); it needs the signing key and refuses `--insecure`. + It rebuilds the project day from the logs too; with `--agent` and `--session` it also checks + that session's own limit. +- `reset --day` rebuilds the day, then writes the signed acknowledgement of its current state + described above (default: today in the policy's `tz`); it needs the signing key and refuses + `--insecure`. ## Cost diff --git a/pyproject.toml b/pyproject.toml index 3545b50..08b6f0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,7 @@ version = {attr = "aegis_core.__version__"} where = ["src"] [tool.setuptools.package-data] -aegis_core = ["_examples/**/*", "integrations/*", "compile/actions/*", "budget/*.yaml"] +aegis_core = ["_examples/**/*", "integrations/*", "compile/actions/*", "budget/*.yaml", "budget/*.json"] [tool.ruff] line-length = 100 diff --git a/scripts/sync_prices.py b/scripts/sync_prices.py new file mode 100644 index 0000000..cbc2043 --- /dev/null +++ b/scripts/sync_prices.py @@ -0,0 +1,34 @@ +"""Regenerates src/aegis_core/budget/prices.json from prices.yaml. + +prices.yaml is the table people edit (it carries the sources and the +reasoning in comments); the package loads prices.json, which is faster to read +on every hook call. tests/test_budget_policy.py fails if the two differ. + + python scripts/sync_prices.py [--check] +""" + +import json +import sys +from pathlib import Path + +import yaml + +BUDGET = Path(__file__).resolve().parent.parent / "src" / "aegis_core" / "budget" + + +def main() -> int: + table = yaml.safe_load((BUDGET / "prices.yaml").read_text()) + text = json.dumps(table, indent=1, sort_keys=True) + "\n" + target = BUDGET / "prices.json" + if "--check" in sys.argv[1:]: + if not target.exists() or target.read_text() != text: + print(f"{target} is stale: run python scripts/sync_prices.py", file=sys.stderr) + return 1 + return 0 + target.write_text(text) + print(f"wrote {target}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/aegis_core/budget/evaluate.py b/src/aegis_core/budget/evaluate.py index 821f594..b4ac4c1 100644 --- a/src/aegis_core/budget/evaluate.py +++ b/src/aegis_core/budget/evaluate.py @@ -77,6 +77,40 @@ def _fmt(unit: str, amount: float) -> str: return f"{amount:,.2f} premium requests" +@dataclass(frozen=True) +class Acknowledgement: + """What a signed ``aegis budget reset --day`` acknowledged: the project + day's verification state at reset time. Only that is excused; anything + that becomes unverifiable later (a new session's unreadable log, more + missing sessions) counts again.""" + + completeness: str + history: str + missing_sessions: int + unverified: Mapping[str, list[str]] + + @classmethod + def of(cls, pd: ProjectDay) -> Acknowledgement: + return cls(pd.completeness, pd.history, pd.missing_sessions, pd.unverified()) + + def to_dict(self) -> dict[str, Any]: + return {"completeness": self.completeness, "history": self.history, + "missing_sessions": self.missing_sessions, + "unverified": {k: sorted(v) for k, v in sorted(self.unverified.items())}} + + @classmethod + def from_dict(cls, d: Mapping[str, Any]) -> Acknowledgement: + return cls(str(d["completeness"]), str(d["history"]), int(d["missing_sessions"]), + {str(k): list(v) for k, v in dict(d["unverified"]).items()}) + + def covers_history(self, pd: ProjectDay) -> bool: + return (pd.completeness == self.completeness and pd.history == self.history + and pd.missing_sessions <= self.missing_sessions) + + def covers(self, key: str, reasons: list[str]) -> bool: + return key in self.unverified and set(reasons) <= set(self.unverified[key]) + + class _Builder: """Collects measures, price flags and unverifiable usage into a Verdict.""" @@ -98,17 +132,17 @@ def amount(self, usage: Mapping[str, Mapping[str, float]]) -> float: def measure(self, name: str, unit: str, used: float, lim: Limit) -> None: self.verdict.measures.append(Measure(name, unit, used, lim.limit, lim.warn_at)) - def project_history(self, pd: ProjectDay, exclude: str | None, acknowledged: bool) -> None: + def project_history(self, pd: ProjectDay, exclude: str | None, + ack: Acknowledgement | None) -> None: """A project day whose history is incomplete, or whose other - sessions are only a lower bound, is unverifiable -- unless a signed - ``aegis budget reset --day`` acknowledged that day.""" - if acknowledged: - return - if pd.completeness != "complete": + sessions are only a lower bound, is unverifiable -- except for what + a signed ``aegis budget reset --day`` acknowledged.""" + if pd.completeness != "complete" and not (ack and ack.covers_history(pd)): detail = (f"{pd.missing_sessions} session(s) without a log" if pd.history == "ok" else f"session inventory {pd.history}") self.unknown.append(f"today's project total is {pd.completeness} ({detail})") - others = pd.unverified(exclude=exclude) + others = {k: r for k, r in pd.unverified(exclude=exclude).items() + if not (ack and ack.covers(k, r))} if others: reasons = sorted({r for rs in others.values() for r in rs}) who = "other session(s)" if exclude else "session(s)" @@ -161,11 +195,11 @@ def finish(self) -> Verdict: def evaluate(policy: BudgetPolicy, prices: PriceTable, snap: Snapshot, *, - acknowledged: bool = False) -> Verdict: - """The decision for one session at one hook. ``acknowledged``: a signed - ``aegis budget reset --day`` covers today, so an incomplete project - history no longer counts as unverifiable (the session's own log problems - still do).""" + ack: Acknowledgement | None = None) -> Verdict: + """The decision for one session at one hook. ``ack``: what a signed + ``aegis budget reset --day`` acknowledged for today; only that part of + an incomplete project history stops counting as unverifiable (the + session's own log problems always count).""" agent = snap.record.agent if not policy.measures(agent): return Verdict(ALLOW) @@ -188,12 +222,12 @@ def evaluate(policy: BudgetPolicy, prices: PriceTable, snap: Snapshot, *, if rec.buckets().get(UNKNOWN_DAY): b.unknown.append("usage with no recorded time (counted for the session only)") if snap.project is not None: - b.project_history(snap.project, rec.key, acknowledged) + b.project_history(snap.project, rec.key, ack) return b.finish() def evaluate_project(policy: BudgetPolicy, prices: PriceTable, pd: ProjectDay, *, - acknowledged: bool = False) -> Verdict: + ack: Acknowledgement | None = None) -> Verdict: """The project/day limit alone, for ``aegis budget check`` without a session: the day's total against ``project_day``, with every contributing session's verification status.""" @@ -201,5 +235,5 @@ def evaluate_project(policy: BudgetPolicy, prices: PriceTable, pd: ProjectDay, * if policy.project_day: token_usage = {m: u for m, u in pd.usage().items() if "requests" not in u} b.measure("project_day", policy.unit, b.amount(token_usage), policy.project_day) - b.project_history(pd, None, acknowledged) + b.project_history(pd, None, ack) return b.finish() diff --git a/src/aegis_core/budget/gate.py b/src/aegis_core/budget/gate.py index 56029eb..18a8c31 100644 --- a/src/aegis_core/budget/gate.py +++ b/src/aegis_core/budget/gate.py @@ -98,9 +98,9 @@ def check(agent: str, payload: Mapping[str, Any], cwd: str | None, policy: Budge snap = refresh(store, logs, loc, policy.zone, now=time.time() if now is None else now, fallback_cwd=cwd) prices = PriceTable(load_builtin_table(), policy.pricing) - acknowledged = bool(snap.record.project_root) and reset_acknowledged( - store, snap.record.project_root, snap.day, key, insecure) - verdict = evaluate(policy, prices, snap, acknowledged=acknowledged) + ack = reset_acknowledged(store, snap.record.project_root, snap.day, key, insecure) \ + if snap.record.project_root else None + verdict = evaluate(policy, prices, snap, ack=ack) if verdict.decision == DENY: over = [m.name for m in verdict.measures if m.level == DENY] ways = sorted({_HOW_TO_CONTINUE[n] for n in over}) diff --git a/src/aegis_core/budget/prices.json b/src/aegis_core/budget/prices.json new file mode 100644 index 0000000..8e81bef --- /dev/null +++ b/src/aegis_core/budget/prices.json @@ -0,0 +1,282 @@ +{ + "as_of": "2026-10-07", + "vendors": { + "anthropic": { + "models": { + "claude-fable-5": { + "cache_read": 1.0, + "cache_write": 20.0, + "input": 10.0, + "output": 50.0 + }, + "claude-fable-5-1": { + "cache_read": 0.25, + "cache_write": 20.0, + "input": 10.0, + "output": 50.0 + }, + "claude-haiku-3-5": { + "cache_read": 0.08, + "cache_write": 1.6, + "input": 0.8, + "output": 4.0 + }, + "claude-haiku-4-5": { + "cache_read": 0.1, + "cache_write": 2.0, + "input": 1.0, + "output": 5.0 + }, + "claude-haiku-5-5": { + "cache_read": 0.05, + "cache_write": 1.0, + "input": 0.5, + "output": 2.5 + }, + "claude-mythos-5": { + "cache_read": 1.0, + "cache_write": 20.0, + "input": 10.0, + "output": 50.0 + }, + "claude-mythos-5-1": { + "cache_read": 0.25, + "cache_write": 20.0, + "input": 10.0, + "output": 50.0 + }, + "claude-opus-4": { + "cache_read": 1.5, + "cache_write": 30.0, + "input": 15.0, + "output": 75.0 + }, + "claude-opus-4-1": { + "cache_read": 1.5, + "cache_write": 30.0, + "input": 15.0, + "output": 75.0 + }, + "claude-opus-4-5": { + "cache_read": 0.5, + "cache_write": 10.0, + "input": 5.0, + "output": 25.0 + }, + "claude-opus-4-6": { + "cache_read": 0.5, + "cache_write": 10.0, + "input": 5.0, + "output": 25.0 + }, + "claude-opus-4-7": { + "cache_read": 0.5, + "cache_write": 10.0, + "input": 5.0, + "output": 25.0 + }, + "claude-opus-4-8": { + "cache_read": 0.5, + "cache_write": 10.0, + "input": 5.0, + "output": 25.0 + }, + "claude-opus-5": { + "cache_read": 0.5, + "cache_write": 10.0, + "input": 5.0, + "output": 25.0 + }, + "claude-opus-5-5": { + "cache_read": 0.2, + "cache_write": 8.0, + "input": 4.0, + "output": 20.0 + }, + "claude-sonnet-4": { + "cache_read": 0.3, + "cache_write": 6.0, + "input": 3.0, + "output": 15.0 + }, + "claude-sonnet-4-5": { + "cache_read": 0.3, + "cache_write": 6.0, + "input": 3.0, + "output": 15.0 + }, + "claude-sonnet-4-6": { + "cache_read": 0.3, + "cache_write": 6.0, + "input": 3.0, + "output": 15.0 + }, + "claude-sonnet-5": { + "cache_read": 0.2, + "cache_write": 4.0, + "input": 2.0, + "output": 10.0 + }, + "claude-sonnet-5-5": { + "cache_read": 0.1, + "cache_write": 4.0, + "input": 2.0, + "output": 10.0 + } + }, + "prefixes": [ + "claude-" + ] + }, + "google": { + "models": { + "gemini-2.5-flash": { + "cache_read": 0.03, + "cache_write": 0.3, + "input": 0.3, + "output": 2.5 + }, + "gemini-2.5-flash-lite": { + "cache_read": 0.01, + "cache_write": 0.1, + "input": 0.1, + "output": 0.4 + }, + "gemini-2.5-pro": { + "cache_read": 0.25, + "cache_write": 2.5, + "input": 2.5, + "output": 15.0 + }, + "gemini-3.1-flash-lite": { + "cache_read": 0.025, + "cache_write": 0.25, + "input": 0.25, + "output": 1.5 + }, + "gemini-3.1-pro-preview": { + "cache_read": 0.4, + "cache_write": 4.0, + "input": 4.0, + "output": 18.0 + }, + "gemini-3.5-flash": { + "cache_read": 0.15, + "cache_write": 1.5, + "input": 1.5, + "output": 9.0 + }, + "gemini-3.5-flash-lite": { + "cache_read": 0.03, + "cache_write": 0.3, + "input": 0.3, + "output": 2.5 + }, + "gemini-3.6-flash": { + "cache_read": 0.15, + "cache_write": 1.5, + "input": 1.5, + "output": 7.5 + }, + "gemini-3.7-flash": { + "cache_read": 0.15, + "cache_write": 1.5, + "input": 1.5, + "output": 7.5 + }, + "gemini-3.8-flash": { + "cache_read": 0.15, + "cache_write": 1.5, + "input": 1.5, + "output": 7.5 + } + }, + "prefixes": [ + "gemini-" + ] + }, + "openai": { + "models": { + "gpt-4.1": { + "cache_read": 0.5, + "cache_write": 2.0, + "input": 2.0, + "output": 8.0 + }, + "gpt-5": { + "cache_read": 0.125, + "cache_write": 1.25, + "input": 1.25, + "output": 10.0 + }, + "gpt-5-mini": { + "cache_read": 0.025, + "cache_write": 0.25, + "input": 0.25, + "output": 2.0 + }, + "gpt-5-nano": { + "cache_read": 0.005, + "cache_write": 0.05, + "input": 0.05, + "output": 0.4 + }, + "gpt-5.3-codex": { + "cache_read": 0.175, + "cache_write": 1.75, + "input": 1.75, + "output": 14.0 + }, + "gpt-5.4": { + "cache_read": 0.25, + "cache_write": 2.5, + "input": 2.5, + "output": 15.0 + }, + "gpt-5.4-mini": { + "cache_read": 0.075, + "cache_write": 0.75, + "input": 0.75, + "output": 4.5 + }, + "gpt-5.4-nano": { + "cache_read": 0.02, + "cache_write": 0.2, + "input": 0.2, + "output": 1.25 + }, + "gpt-5.5": { + "cache_read": 0.5, + "cache_write": 5.0, + "input": 5.0, + "output": 30.0 + }, + "gpt-5.5-pro": { + "cache_write": 30.0, + "input": 30.0, + "output": 180.0 + }, + "o3": { + "cache_read": 0.5, + "cache_write": 2.0, + "input": 2.0, + "output": 8.0 + }, + "o4-mini": { + "cache_read": 0.275, + "cache_write": 1.1, + "input": 1.1, + "output": 4.4 + } + }, + "prefixes": [ + "gpt-", + "o1", + "o3", + "o4", + "codex-" + ] + } + }, + "version": 1 +} diff --git a/src/aegis_core/budget/pricing.py b/src/aegis_core/budget/pricing.py index 213b2c9..3aaa2f6 100644 --- a/src/aegis_core/budget/pricing.py +++ b/src/aegis_core/budget/pricing.py @@ -14,16 +14,12 @@ from __future__ import annotations -import contextlib -import hashlib import json import math -import os import re from collections.abc import Mapping from dataclasses import dataclass, field from importlib import resources -from pathlib import Path from typing import Any import yaml @@ -166,34 +162,25 @@ def cost(self, usage_by_model: Mapping[str, Mapping[str, int]], return total -def _table_cache_dir() -> Path: - home = Path(os.environ.get("HOME") or Path.home()) - return Path(os.environ.get("XDG_CACHE_HOME") or home / ".cache") / "aegis" / "budget" - - def load_builtin_table() -> dict[str, Any]: - """The built-in price table. Parsing it with PyYAML's pure-Python loader - costs ~15 ms, paid on every hook call, so the parsed table is cached as - JSON under the cache directory, keyed by a hash of the YAML text: an - edited table is re-parsed, and a damaged or missing cache only costs a - re-parse.""" - text = resources.files("aegis_core.budget").joinpath("prices.yaml").read_text() - digest = hashlib.sha256(text.encode()).hexdigest()[:16] - cache = _table_cache_dir() / f"prices-{digest}.json" - try: - table = json.loads(cache.read_text()) - except (OSError, ValueError): - table = yaml.safe_load(text) - with contextlib.suppress(OSError): - cache.parent.mkdir(parents=True, exist_ok=True) - tmp = cache.with_suffix(f".{os.getpid()}.tmp") - tmp.write_text(json.dumps(table)) - os.replace(tmp, cache) + """The built-in price table, from ``prices.json``: the same table as + ``prices.yaml`` (the file people edit), generated by + ``scripts/sync_prices.py`` and kept equal by a test. It is read from the + installed package, never from a writable cache, so nothing but a signed + ``pricing:`` override can change a price; and JSON loads in well under a + millisecond where parsing the YAML on every hook call cost ~15 ms.""" + table = json.loads(resources.files("aegis_core.budget").joinpath("prices.json").read_text()) if not isinstance(table, dict) or table.get("version") != 1: - raise ValueError("aegis_core/budget/prices.yaml: unsupported shape or version") + raise ValueError("aegis_core/budget/prices.json: unsupported shape or version") return table +def load_source_table() -> dict[str, Any]: + """``prices.yaml`` itself (for the sync script and its test).""" + return yaml.safe_load( + resources.files("aegis_core.budget").joinpath("prices.yaml").read_text()) + + def check_override(where: str, entry: Any) -> dict[str, float]: """A ``budget.yaml`` price override: all four categories, numbers >= 0.""" return _check_rates(where, entry, require_all=True) diff --git a/src/aegis_core/budget/report.py b/src/aegis_core/budget/report.py index 28d0856..8d26a1c 100644 --- a/src/aegis_core/budget/report.py +++ b/src/aegis_core/budget/report.py @@ -20,7 +20,14 @@ from typing import Any from aegis_core.budget.accounting import BudgetStore, day_of, rebuild -from aegis_core.budget.evaluate import ALLOW, DENY, Measure, evaluate_project, tokens_of +from aegis_core.budget.evaluate import ( + ALLOW, + DENY, + Acknowledgement, + Measure, + evaluate_project, + tokens_of, +) from aegis_core.budget.logs import AgentLog from aegis_core.budget.policy import BudgetPolicy from aegis_core.budget.pricing import PriceTable @@ -43,13 +50,15 @@ def reset_path(store: BudgetStore, root: str, day: str) -> Path: return store.project_dir(root) / f"reset-{day}.json" -def write_reset(store: BudgetStore, root: str, day: str, key: bytes, +def write_reset(store: BudgetStore, root: str, day: str, key: bytes, baseline: Acknowledgement, now: float | None = None) -> Path: - """Acknowledges ``day``'s project history as it stands, signed with - ``key`` (the policy signing key).""" + """Acknowledges ``day``'s project history as it stands -- ``baseline``, + the day's verification state now -- signed with ``key`` (the policy + signing key). Later problems are not covered.""" fields = {"header": RESET_HEADER, "root": root, "day": day, "at": datetime.fromtimestamp(time.time() if now is None else now, - tz=UTC).isoformat()} + tz=UTC).isoformat(), + "baseline": baseline.to_dict()} path = reset_path(store, root, day) path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps({**fields, "mac": _mac(_canonical(fields), key)}, indent=2)) @@ -57,23 +66,32 @@ def write_reset(store: BudgetStore, root: str, day: str, key: bytes, def reset_acknowledged(store: BudgetStore, root: str, day: str, key: bytes | None, - insecure: bool = False) -> bool: - """Whether a valid reset marker covers ``day`` in project ``root``. - Without a key (``--insecure``) nothing is verified, so a marker's - presence is taken as it is, like every other policy file then.""" + insecure: bool = False) -> Acknowledgement | None: + """What a valid reset marker for ``day`` in project ``root`` + acknowledged, or ``None``. Without a key (``--insecure``) nothing is + verified, so a marker is taken as it is, like every other policy file + then.""" try: doc = json.loads(reset_path(store, root, day).read_text()) except (OSError, ValueError): - return False + return None if not isinstance(doc, dict): - return False - fields = {k: doc.get(k) for k in ("header", "root", "day", "at")} + return None + fields = {k: doc.get(k) for k in ("header", "root", "day", "at", "baseline")} if fields["header"] != RESET_HEADER or fields["root"] != root or fields["day"] != day: - return False + return None if key is None: - return insecure - mac = doc.get("mac") - return isinstance(mac, str) and hmac.compare_digest(mac, _mac(_canonical(fields), key)) + if not insecure: + return None + else: + mac = doc.get("mac") + if not (isinstance(mac, str) + and hmac.compare_digest(mac, _mac(_canonical(fields), key))): + return None + try: + return Acknowledgement.from_dict(fields["baseline"]) + except (KeyError, TypeError, ValueError): + return None # --- status / check -------------------------------------------------------------------- @@ -98,8 +116,8 @@ def project_status(policy: BudgetPolicy, prices: PriceTable, store: BudgetStore, day = day_of(datetime.fromtimestamp(now, tz=UTC), zone) measured = {a: lg for a, lg in logs.items() if policy.measures(a)} pd = rebuild(store, measured, root, day, zone) - acknowledged = reset_acknowledged(store, root, day, key, insecure) - verdict = evaluate_project(policy, prices, pd, acknowledged=acknowledged) + ack = reset_acknowledged(store, root, day, key, insecure) + verdict = evaluate_project(policy, prices, pd, ack=ack) sessions = [] for entry_key, entry in sorted(pd.entries.items()): @@ -131,7 +149,7 @@ def project_status(policy: BudgetPolicy, prices: PriceTable, store: BudgetStore, return { "project": root, "day": day, "tz": policy.tz, "unit": policy.unit, "project_day": project, "completeness": pd.completeness, "history": pd.history, - "missing_sessions": pd.missing_sessions, "acknowledged": acknowledged, + "missing_sessions": pd.missing_sessions, "acknowledged": ack is not None, "decision": verdict.decision, "messages": verdict.messages, "estimated": verdict.estimated, "sessions": sessions, "warnings": list(policy.warnings), diff --git a/src/aegis_core/cli.py b/src/aegis_core/cli.py index 12bc12a..8118d47 100644 --- a/src/aegis_core/cli.py +++ b/src/aegis_core/cli.py @@ -1411,7 +1411,8 @@ def load_budget_with_key(config_dir: Path, extra_args: list[str] | None = None): def _run_budget(args: argparse.Namespace) -> int: """``aegis budget status|check|reset`` (docs/budget.md).""" from aegis_core.budget import BudgetStore, default_logs, evaluate, price_table, refresh - from aegis_core.budget.accounting import day_of, find_project_root + from aegis_core.budget.accounting import Snapshot, day_of, find_project_root, rebuild + from aegis_core.budget.evaluate import Acknowledgement from aegis_core.budget.report import ( check_exit, project_status, @@ -1432,6 +1433,7 @@ def _run_budget(args: argparse.Namespace) -> int: raise ConfigNotFoundError(f"no budget: {config_dir / 'budget.yaml'} does not exist") root = find_project_root(str(project)) store, logs, prices = BudgetStore(), default_logs(), price_table(policy) + measured = {a: lg for a, lg in logs.items() if policy.measures(a)} now = time.time() today = day_of(datetime.fromtimestamp(now, tz=UTC), policy.zone) @@ -1447,7 +1449,9 @@ def _run_budget(args: argparse.Namespace) -> int: datetime.strptime(day, "%Y-%m-%d") except ValueError: raise UsageError(f"--day {day!r} is not YYYY-MM-DD") from None - path = write_reset(store, root, day, key, now) + # the baseline is the day as the logs show it now: only that is excused + pd = rebuild(store, measured, root, day, policy.zone) + path = write_reset(store, root, day, key, Acknowledgement.of(pd), now) print(f"aegis budget: history for {day} acknowledged for {root} ({path})") return 0 @@ -1472,11 +1476,16 @@ def _run_budget(args: argparse.Namespace) -> int: loc = logs[args.agent].locate(args.session) if loc is None: raise DataError(f"cannot find the log of {args.agent} session {args.session}") - measured = {a: lg for a, lg in logs.items() if policy.measures(a)} snap = refresh(store, measured, loc, policy.zone, now=now, fallback_cwd=str(project)) - acknowledged = bool(snap.record.project_root) and reset_acknowledged( - store, snap.record.project_root, snap.day, key, insecure) - verdict = evaluate(policy, prices, snap, acknowledged=acknowledged) + ack = None + if snap.record.project_root: + # a refresh only updates this session; rebuild the project day from + # every measured agent's logs, as status and an unscoped check do + pd = rebuild(store, measured, snap.record.project_root, snap.day, policy.zone) + snap = Snapshot(snap.record, snap.day, pd, snap.session_usage, pd.usage(), + {**pd.providers(), **snap.record.providers}) + ack = reset_acknowledged(store, snap.record.project_root, snap.day, key, insecure) + verdict = evaluate(policy, prices, snap, ack=ack) decision, messages = verdict.decision, verdict.messages elif root is None: print(f"aegis budget: {project} is not inside a project with .aegis/; pass --agent " diff --git a/tests/test_budget_cli.py b/tests/test_budget_cli.py index 553b8e4..0cd5066 100644 --- a/tests/test_budget_cli.py +++ b/tests/test_budget_cli.py @@ -150,15 +150,18 @@ def test_reset_markers_are_verified(env, capsys): # noqa: F811 key = bytes.fromhex("".join( line.strip() for line in (env["project"] / ".aegis" / "example-signing.key") .read_text().splitlines() if not line.lstrip().startswith("#"))) - assert reset_acknowledged(store, root, TODAY, key) + assert reset_acknowledged(store, root, TODAY, key) is not None doc = json.loads(marker.read_text()) marker.write_text(json.dumps({**doc, "day": "2001-01-01"})) # moved to another day - assert not reset_acknowledged(store, root, TODAY, key) + assert reset_acknowledged(store, root, TODAY, key) is None marker.write_text(json.dumps({**doc, "mac": "0" * 64})) # forged without the key - assert not reset_acknowledged(store, root, TODAY, key) + assert reset_acknowledged(store, root, TODAY, key) is None + widened = {**doc["baseline"], "missing_sessions": 99} + marker.write_text(json.dumps({**doc, "baseline": widened})) # baseline widened + assert reset_acknowledged(store, root, TODAY, key) is None assert hook_call(env, "now")[0] == 2 - assert not reset_acknowledged(store, root, TODAY, None) # no key: not verified - assert reset_acknowledged(store, root, TODAY, None, insecure=True) is True + assert reset_acknowledged(store, root, TODAY, None) is None # no key: not verified + assert reset_acknowledged(store, root, TODAY, None, insecure=True) is not None def test_reset_needs_the_key_and_a_project(env, capsys, tmp_path): # noqa: F811 @@ -169,3 +172,40 @@ def test_reset_needs_the_key_and_a_project(env, capsys, tmp_path): # noqa: F811 "--project", str(env["project"])]) == 64 assert main(["budget", "reset", "--day", "2026-10-01", "--project", str(env["project"])]) == 0 + + +# --- review of #40 --------------------------------------------------------------------------- + + +def test_a_session_check_sees_other_sessions_unhooked_spend(env, capsys): # noqa: F811 + """check --agent/--session rebuilds the project day too: session a's + hook recorded $0.50, then session b spent $2.00 with no hook firing.""" + write_budget(env, DAY_BUDGET) + transcript(env, "a", 125_000) + assert hook_call(env, "a")[0] == 0 # the project cache now exists + transcript(env, "b", 500_000) + assert _check(env, "--agent", "claude", "--session", "a") == 3 + assert "project day limit reached" in capsys.readouterr().out + + +def test_a_reset_does_not_excuse_problems_that_appear_after_it(env, capsys): # noqa: F811 + write_budget(env, STRICT) + transcript(env, "a", 1_000) + assert hook_call(env, "a")[0] == 0 + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + transcript(env, "c", 1_000, malformed=True) # a new, unreadable session + code, _, err = hook_call(env, "c") + assert code == 2 and "unreadable log record" in err # its own hook denies + code, _, err = hook_call(env, "a") + assert code == 2 and "lower bound" in err # and it is not excused for a + + +def test_a_reset_still_excuses_what_it_acknowledged(env, capsys): # noqa: F811 + write_budget(env, STRICT) + transcript(env, "a", 1_000) + transcript(env, "c", 1_000, malformed=True) + hook_call(env, "a") + hook_call(env, "c") + assert hook_call(env, "a")[0] == 2 # c is a lower bound: strict denies + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + assert hook_call(env, "a")[0] == 0 # acknowledged as it stood diff --git a/tests/test_budget_policy.py b/tests/test_budget_policy.py index e2d7a32..5efbcf3 100644 --- a/tests/test_budget_policy.py +++ b/tests/test_budget_policy.py @@ -267,29 +267,38 @@ def test_an_override_reprices_earlier_estimated_usage(table): assert after.estimated == {} and after.dollars == pytest.approx(4.0) -def test_the_parsed_price_table_is_cached_by_content(tmp_path, monkeypatch): - """Parsing prices.yaml costs ~15 ms on every hook call; the parsed table - is cached as JSON keyed by a hash of the YAML text.""" - import aegis_core.budget.pricing as pricing +def test_prices_json_is_prices_yaml(): + """The package loads prices.json (fast); people edit prices.yaml. They + must be the same table: run scripts/sync_prices.py after an edit.""" + import subprocess + import sys + from aegis_core.budget.pricing import load_source_table + + assert load_builtin_table() == load_source_table() + assert subprocess.run([sys.executable, "scripts/sync_prices.py", "--check"], + capture_output=True).returncode == 0 + + +def test_no_writable_file_can_change_a_price(tmp_path, monkeypatch): + """Review of #40: a cached copy of the table under ~/.cache acted as an + unsigned price override. The table is read only from the package now; a + planted 'cache' with zero rates, or with a bogus version, changes + nothing.""" monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) - first = load_builtin_table() - [cache] = (tmp_path / "aegis" / "budget").glob("prices-*.json") - assert json.loads(cache.read_text()) == first - cache.write_text("{damaged") # a damaged cache re-parses - assert load_builtin_table() == first - assert json.loads(cache.read_text()) == first - real = pricing.resources.files - - class Edited: - def __init__(self, pkg): - self.pkg = pkg - - def joinpath(self, name): - text = real(self.pkg).joinpath(name).read_text().replace("as_of: \"2026-10-07\"", - "as_of: \"2027-01-01\"") - return type("T", (), {"read_text": lambda self: text})() - - monkeypatch.setattr(pricing.resources, "files", Edited) - assert load_builtin_table()["as_of"] == "2027-01-01" # an edited table is re-read - assert len(list((tmp_path / "aegis" / "budget").glob("prices-*.json"))) == 2 + before = load_builtin_table() + cache = tmp_path / "aegis" / "budget" + cache.mkdir(parents=True, exist_ok=True) + zero = {**before, "vendors": {v: {**spec, "models": {m: {c: 0 for c in r} + for m, r in spec["models"].items()}} + for v, spec in before["vendors"].items()}} + import hashlib + from importlib import resources + + yaml_text = resources.files("aegis_core.budget").joinpath("prices.yaml").read_text() + old_name = f"prices-{hashlib.sha256(yaml_text.encode()).hexdigest()[:16]}.json" + for name in ("prices.json", old_name): # old_name: the cache the review exploited + (cache / name).write_text(json.dumps(zero)) + (cache / "prices-bogus.json").write_text(json.dumps({"version": 999})) + assert load_builtin_table() == before + assert PriceTable(load_builtin_table()).price("claude-opus-5-5").rates["input"] == 4.0 From caec6cffe61d5e6925e7c52e0e01209708059c93 Mon Sep 17 00:00:00 2001 From: moneytools Date: Fri, 9 Oct 2026 21:16:16 -0500 Subject: [PATCH 4/5] chore: re-measure budget latency with prices.json (+9.5 ms, 74.2 vs 64.7 ms) --- results/budget-latency.json | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/results/budget-latency.json b/results/budget-latency.json index 870947c..f612708 100644 --- a/results/budget-latency.json +++ b/results/budget-latency.json @@ -4,13 +4,13 @@ "python": "3.12.12", "platform": "Darwin arm64", "hook_without_budget": { - "median_ms": 65.3, - "p95_ms": 68.0 + "median_ms": 64.7, + "p95_ms": 67.3 }, "hook_with_budget": { - "median_ms": 74.1, - "p95_ms": 77.6 + "median_ms": 74.2, + "p95_ms": 76.7 }, - "added_median_ms": 8.8, + "added_median_ms": 9.5, "target_ms": 20 } From fd583220d21eba143dd2e5186e278d142238bb7f Mon Sep 17 00:00:00 2001 From: moneytools Date: Fri, 9 Oct 2026 21:20:30 -0500 Subject: [PATCH 5/5] fix(budget): reset baselines record quantities, so a gap that grows counts again Project entries carry each session's verification quantities (log missing, unreadable records, usage with no recorded time); an acknowledgement covers a session only while none of them has grown. Review of #40. --- docs/budget.md | 8 ++++--- src/aegis_core/budget/accounting.py | 22 ++++++++++++++++-- src/aegis_core/budget/evaluate.py | 34 +++++++++++++++++++++------ tests/test_budget_cli.py | 36 +++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 12 deletions(-) diff --git a/docs/budget.md b/docs/budget.md index 77e3ded..a3a394b 100644 --- a/docs/budget.md +++ b/docs/budget.md @@ -126,10 +126,12 @@ were deleted, or the first day of a budget when other sessions had already run t - `deny`: covered calls are denied. A session's own unreadable log stays denied (start a new session). For an incomplete **project day**, `aegis budget reset --day` acknowledges that day's history **as it stands**: it records the day's verification state at that moment (how - complete the history is, which sessions are lower bounds and why) and is signed with the + complete the history is, and for each session that is a lower bound: whether its log is + missing, how many records could not be read, how much usage has no recorded time) and is + signed with the policy key, so an agent cannot write one. Exactly that is excused; anything that becomes - unverifiable afterwards — a new session with an unreadable log, more sessions without one — - counts again. Otherwise it clears at the next day. + unverifiable afterwards — a new session with an unreadable log, another unreadable record in a + session already acknowledged, more sessions without a log — counts again. Otherwise it clears at the next day. ## Commands diff --git a/src/aegis_core/budget/accounting.py b/src/aegis_core/budget/accounting.py index d084142..c8aa56c 100644 --- a/src/aegis_core/budget/accounting.py +++ b/src/aegis_core/budget/accounting.py @@ -195,6 +195,15 @@ def unverified(self) -> list[str]: reasons.append("unknown-time") return reasons + def verification(self) -> dict[str, Any]: + """How much of this session's usage is unverifiable, as quantities + (not only reasons), so an acknowledgement can tell an unchanged gap + from one that grew: whether the log is missing, how many records + could not be read, and how much usage carries no recorded time.""" + unknown = self.buckets().get(UNKNOWN_DAY, {}) + return {"log_missing": self.log_missing, "unreadable_records": self.problem_total, + "unknown_time_usage": sum(v for u in unknown.values() for v in u.values())} + def to_json(self) -> str: return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) @@ -231,6 +240,12 @@ def unverified(self, exclude: str | None = None) -> dict[str, list[str]]: return {k: list(e["unverified"]) for k, e in self.entries.items() if e.get("unverified") and k != exclude} + def verification(self, key: str) -> dict[str, Any] | None: + """The quantities behind an entry's ``unverified`` reasons (see + :meth:`SessionRecord.verification`); ``None`` if not recorded.""" + v = self.entries.get(key, {}).get("verification") + return v if isinstance(v, dict) else None + def to_json(self) -> str: return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) @@ -519,7 +534,7 @@ def _entry(rec: SessionRecord, day: str) -> dict[str, Any]: usage = rec.usage_on(day) return {"revision": rec.revision, "usage": usage, "providers": {m: p for m, p in rec.providers.items() if m in usage}, - "unverified": rec.unverified()} + "unverified": rec.unverified(), "verification": rec.verification()} # --- rebuild ------------------------------------------------------------------------- @@ -574,7 +589,10 @@ def rebuild(store: BudgetStore, logs: Mapping[str, AgentLog], root: str, day: st if cp: incoming[key] = {"revision": int(cp.get("revision", 0)), "usage": cp.get("usage", {}), "providers": {}, - "unverified": ["log-missing", "from-checkpoint"]} + "unverified": ["log-missing", "from-checkpoint"], + "verification": {"log_missing": True, + "unreadable_records": 0, + "unknown_time_usage": 0}} continue # A missing log with a cached record keeps its usage as a lower bound; # the entry says so (its "unverified" reasons). diff --git a/src/aegis_core/budget/evaluate.py b/src/aegis_core/budget/evaluate.py index b4ac4c1..14f1f0a 100644 --- a/src/aegis_core/budget/evaluate.py +++ b/src/aegis_core/budget/evaluate.py @@ -87,28 +87,48 @@ class Acknowledgement: completeness: str history: str missing_sessions: int - unverified: Mapping[str, list[str]] + # session -> its verification quantities at reset time + # (SessionRecord.verification: log_missing, unreadable_records, unknown_time_usage) + unverified: Mapping[str, Mapping[str, Any]] @classmethod def of(cls, pd: ProjectDay) -> Acknowledgement: - return cls(pd.completeness, pd.history, pd.missing_sessions, pd.unverified()) + return cls(pd.completeness, pd.history, pd.missing_sessions, + {k: pd.verification(k) or {} for k in pd.unverified()}) def to_dict(self) -> dict[str, Any]: return {"completeness": self.completeness, "history": self.history, "missing_sessions": self.missing_sessions, - "unverified": {k: sorted(v) for k, v in sorted(self.unverified.items())}} + "unverified": {k: dict(v) for k, v in sorted(self.unverified.items())}} @classmethod def from_dict(cls, d: Mapping[str, Any]) -> Acknowledgement: + unverified = dict(d["unverified"]) + if not all(isinstance(v, dict) for v in unverified.values()): + raise ValueError("reset baseline: unverified sessions need their quantities") return cls(str(d["completeness"]), str(d["history"]), int(d["missing_sessions"]), - {str(k): list(v) for k, v in dict(d["unverified"]).items()}) + {str(k): dict(v) for k, v in unverified.items()}) def covers_history(self, pd: ProjectDay) -> bool: return (pd.completeness == self.completeness and pd.history == self.history and pd.missing_sessions <= self.missing_sessions) - def covers(self, key: str, reasons: list[str]) -> bool: - return key in self.unverified and set(reasons) <= set(self.unverified[key]) + def covers(self, key: str, now: Mapping[str, Any] | None) -> bool: + """Whether session ``key``'s unverifiable usage is no worse than it + was when acknowledged: the same session, a log no more missing, no + more unreadable records and no more usage without a recorded time. + Unknown quantities are never covered.""" + then = self.unverified.get(key) + if then is None or now is None: + return False + try: + return (bool(now.get("log_missing")) <= bool(then.get("log_missing")) + and int(now.get("unreadable_records", 0)) + <= int(then.get("unreadable_records", 0)) + and float(now.get("unknown_time_usage", 0)) + <= float(then.get("unknown_time_usage", 0))) + except (TypeError, ValueError): + return False class _Builder: @@ -142,7 +162,7 @@ def project_history(self, pd: ProjectDay, exclude: str | None, else f"session inventory {pd.history}") self.unknown.append(f"today's project total is {pd.completeness} ({detail})") others = {k: r for k, r in pd.unverified(exclude=exclude).items() - if not (ack and ack.covers(k, r))} + if not (ack and ack.covers(k, pd.verification(k)))} if others: reasons = sorted({r for rs in others.values() for r in rs}) who = "other session(s)" if exclude else "session(s)" diff --git a/tests/test_budget_cli.py b/tests/test_budget_cli.py index 0cd5066..0b81a07 100644 --- a/tests/test_budget_cli.py +++ b/tests/test_budget_cli.py @@ -209,3 +209,39 @@ def test_a_reset_still_excuses_what_it_acknowledged(env, capsys): # noqa: F811 assert hook_call(env, "a")[0] == 2 # c is a lower bound: strict denies assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 assert hook_call(env, "a")[0] == 0 # acknowledged as it stood + + + +def test_more_unreadable_records_in_an_acknowledged_session_count_again(env): # noqa: F811 + """Review of #40: the baseline must hold quantities, not only reasons. A + second malformed record appended to an already acknowledged session is new + unverifiable usage, so a healthy session is denied again.""" + write_budget(env, STRICT) + transcript(env, "a", 1_000) + broken = transcript(env, "broken", 1_000, malformed=True) + hook_call(env, "a") + hook_call(env, "broken") + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + assert hook_call(env, "a")[0] == 0 # the one bad record: excused + with open(broken, "a") as f: + f.write("{also broken\n") + assert hook_call(env, "broken")[0] == 2 # its own hook records the change + code, _, err = hook_call(env, "a") + assert code == 2 and "lower bound" in err + + +def test_the_acknowledged_quantities_are_signed(env): # noqa: F811 + write_budget(env, STRICT) + transcript(env, "a", 1_000) + transcript(env, "broken", 1_000, malformed=True) + hook_call(env, "a") + hook_call(env, "broken") + assert main(["budget", "reset", "--day", "--project", str(env["project"])]) == 0 + store = BudgetStore() + marker = reset_path(store, str(env["project"].resolve()), TODAY) + doc = json.loads(marker.read_text()) + [state] = doc["baseline"]["unverified"].values() + assert state == {"log_missing": False, "unreadable_records": 1, "unknown_time_usage": 0} + state["unreadable_records"] = 50 # widened after signing + marker.write_text(json.dumps(doc)) + assert hook_call(env, "a")[0] == 2