From 2c20bae4831a8f1107129725b68ef2f81a0e1c67 Mon Sep 17 00:00:00 2001 From: moneytools Date: Wed, 7 Oct 2026 22:47:08 -0500 Subject: [PATCH 1/2] =?UTF-8?q?feat(budget):=20v1.0=20budget=20cap=20core?= =?UTF-8?q?=20=E2=80=94=20policy,=20prices,=20log=20readers,=20accounting,?= =?UTF-8?q?=20evaluation;=20version=201.0.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit aegis_core.budget implements design v1.0 §3-§6 as a library; the hook integration and the 'aegis budget' commands follow in separate PRs. - budget.yaml loader: signed, principal needs the new 'budget' class, unit required, only listed agents measured, every shape problem a load error - dated price table (Anthropic/OpenAI/Google, 2026-10-07); per-category fallback estimates; on_unknown_price; usage kept as tokens so overrides re-price - log readers for Claude Code, Codex, Gemini CLI, OpenCode, Copilot CLI with linked sub-agents, tested on redacted real-log samples - accounting per §4.1: revisioned merge with a durable counter, session inventory, rebuild with completeness, recorded-time day buckets - evaluation: allow/warn/deny, overshoot, unknown-log handling - version 1.0.0 (unreleased); example authority gives admin 'budget' --- CHANGELOG.md | 29 + data/authority.example.yaml | 1 + data/authority.example.yaml.sig | 2 +- gemini-extension.json | 2 +- plugins/copilot/plugin.json | 2 +- pyproject.toml | 2 +- src/aegis_core/__init__.py | 2 +- .../_examples/authority.example.yaml | 1 + .../_examples/authority.example.yaml.sig | 2 +- src/aegis_core/budget/__init__.py | 27 + src/aegis_core/budget/accounting.py | 608 ++++++++++++++++ src/aegis_core/budget/evaluate.py | 159 ++++ src/aegis_core/budget/logs.py | 678 ++++++++++++++++++ src/aegis_core/budget/policy.py | 232 ++++++ src/aegis_core/budget/prices.yaml | 77 ++ src/aegis_core/budget/pricing.py | 171 +++++ tests/fixtures/budget/claude/session.jsonl | 18 + tests/fixtures/budget/claude/subagent.jsonl | 7 + tests/fixtures/budget/codex/child.jsonl | 10 + tests/fixtures/budget/codex/parent.jsonl | 20 + tests/fixtures/budget/copilot/events.jsonl | 17 + tests/fixtures/budget/gemini/session.jsonl | 26 + tests/fixtures/budget/opencode/db.json | 385 ++++++++++ tests/test_authority.py | 4 +- tests/test_budget_accounting.py | 406 +++++++++++ tests/test_budget_evaluate.py | 130 ++++ tests/test_budget_logs.py | 365 ++++++++++ tests/test_budget_policy.py | 241 +++++++ 28 files changed, 3616 insertions(+), 8 deletions(-) create mode 100644 src/aegis_core/budget/__init__.py create mode 100644 src/aegis_core/budget/accounting.py create mode 100644 src/aegis_core/budget/evaluate.py create mode 100644 src/aegis_core/budget/logs.py create mode 100644 src/aegis_core/budget/policy.py create mode 100644 src/aegis_core/budget/prices.yaml create mode 100644 src/aegis_core/budget/pricing.py create mode 100644 tests/fixtures/budget/claude/session.jsonl create mode 100644 tests/fixtures/budget/claude/subagent.jsonl create mode 100644 tests/fixtures/budget/codex/child.jsonl create mode 100644 tests/fixtures/budget/codex/parent.jsonl create mode 100644 tests/fixtures/budget/copilot/events.jsonl create mode 100644 tests/fixtures/budget/gemini/session.jsonl create mode 100644 tests/fixtures/budget/opencode/db.json create mode 100644 tests/test_budget_accounting.py create mode 100644 tests/test_budget_evaluate.py create mode 100644 tests/test_budget_logs.py create mode 100644 tests/test_budget_policy.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 875c69f..8b5bb0f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,35 @@ All notable changes to this project are documented here. The format follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html) — while the major version is 0, any release may change behaviour. +## [1.0.0] — Unreleased + +The cross-agent session budget cap, and the 1.0 stability promise +([design](docs/dev/DESIGN-v1.0-budget-cap.md)). In progress: this entry grows as the parts land; +the release is cut when the hook integration, the `aegis budget` commands, the docs and the live +test on each agent are in. + +### Added +- `aegis_core.budget`, the budget cap's core (design §3–§6): + - `budget.yaml`, signed, whose `principal` must hold the new **`budget`** authority class: + a per-session and a per-project-per-day limit in estimated USD or tokens (`unit` is + required), a warning threshold, the day's time zone, the agents measured (only listed + agents are), `on_unknown_log`, `on_unknown_price`, signed price overrides and a + premium-request limit for Copilot CLI. Every shape problem is a load error. + - A dated price table (`aegis_core/budget/prices.yaml`: Anthropic, OpenAI and Google list + prices read 2026-10-07). An unpriced model is estimated per billing category from its + vendor's (or the whole table's) highest rate — an estimate, not an upper bound — or, with + `on_unknown_price: deny`, refused until priced. Usage is stored as tokens per model, so a + new override re-prices earlier usage. + - Usage readers for the session logs of Claude Code, Codex CLI, Gemini CLI, OpenCode and + Copilot CLI (premium requests), with linked sub-agents counted inside their parent session, + tested against redacted samples of real logs. + - Accounting that the logs can always rebuild: per-session records, project/day totals merged + by revision from a durable per-session counter, a session inventory outside the cache, + rebuilds that report whether the day is complete, partial or of unknown completeness, and + usage bucketed by the time the log recorded, never the time it was read. + - Evaluation to allow / warn / deny with the measured usage, limits and observed overshoot. +- The example `authority.yaml` gives `admin` the `budget` class. + ## [0.3.2] — 2026-09-29 Package metadata only; no change to behaviour. diff --git a/data/authority.example.yaml b/data/authority.example.yaml index 3f6a079..7c8c9d9 100644 --- a/data/authority.example.yaml +++ b/data/authority.example.yaml @@ -6,6 +6,7 @@ principals: - deletion - configuration - identity # may change agents.yaml, the identity model + - budget # may change budget.yaml, the budget cap sre_lead: - scaling - configuration diff --git a/data/authority.example.yaml.sig b/data/authority.example.yaml.sig index f395450..fe83479 100644 --- a/data/authority.example.yaml.sig +++ b/data/authority.example.yaml.sig @@ -1,2 +1,2 @@ aegis-sig-v1 -b0d414e6dcaee5b9ca5d5ed5aa8405f35ed1ea2c7a623124faed7bfe670a104d +b2ceb3e6cb27ac15787c4d6b9015065e03b804a6ac837c8fcc4f9895961bd130 diff --git a/gemini-extension.json b/gemini-extension.json index c1f8830..e298b77 100644 --- a/gemini-extension.json +++ b/gemini-extension.json @@ -1,5 +1,5 @@ { "name": "aegis-devops", - "version": "0.3.2", + "version": "1.0.0", "description": "Checks the infrastructure commands Gemini CLI wants to run (kubectl, terraform/tofu, aws/az/gcloud, helm, argocd, git, SQL) against your signed, authority-checked policy before they run. Only active in projects with a .aegis/ policy. Needs 'pip install aegis-devops'." } diff --git a/plugins/copilot/plugin.json b/plugins/copilot/plugin.json index b384d81..a5acfd0 100644 --- a/plugins/copilot/plugin.json +++ b/plugins/copilot/plugin.json @@ -1,7 +1,7 @@ { "$schema": "https://agent-plugins.org/schemas/1.0.0/plugin.schema.json", "name": "aegis-devops", - "version": "0.3.2", + "version": "1.0.0", "description": "Checks the infrastructure commands Copilot wants to run (kubectl, terraform/tofu, aws/az/gcloud, helm, argocd, git, SQL) against your signed, authority-checked policy before they run. Only active in projects with a .aegis/ policy directory. Needs 'pip install aegis-devops'.", "author": { "name": "Vamshidher Reddy Jannapu Reddy" }, "homepage": "https://github.com/moneytool/aegis-devops", diff --git a/pyproject.toml b/pyproject.toml index 8e32edb..3545b50 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/*"] +aegis_core = ["_examples/**/*", "integrations/*", "compile/actions/*", "budget/*.yaml"] [tool.ruff] line-length = 100 diff --git a/src/aegis_core/__init__.py b/src/aegis_core/__init__.py index 4440148..2e0aa02 100644 --- a/src/aegis_core/__init__.py +++ b/src/aegis_core/__init__.py @@ -4,4 +4,4 @@ See the README for what this does and, importantly, what it does not yet do. """ -__version__ = "0.3.2" +__version__ = "1.0.0" diff --git a/src/aegis_core/_examples/authority.example.yaml b/src/aegis_core/_examples/authority.example.yaml index 3f6a079..7c8c9d9 100644 --- a/src/aegis_core/_examples/authority.example.yaml +++ b/src/aegis_core/_examples/authority.example.yaml @@ -6,6 +6,7 @@ principals: - deletion - configuration - identity # may change agents.yaml, the identity model + - budget # may change budget.yaml, the budget cap sre_lead: - scaling - configuration diff --git a/src/aegis_core/_examples/authority.example.yaml.sig b/src/aegis_core/_examples/authority.example.yaml.sig index f395450..fe83479 100644 --- a/src/aegis_core/_examples/authority.example.yaml.sig +++ b/src/aegis_core/_examples/authority.example.yaml.sig @@ -1,2 +1,2 @@ aegis-sig-v1 -b0d414e6dcaee5b9ca5d5ed5aa8405f35ed1ea2c7a623124faed7bfe670a104d +b2ceb3e6cb27ac15787c4d6b9015065e03b804a6ac837c8fcc4f9895961bd130 diff --git a/src/aegis_core/budget/__init__.py b/src/aegis_core/budget/__init__.py new file mode 100644 index 0000000..79b6f63 --- /dev/null +++ b/src/aegis_core/budget/__init__.py @@ -0,0 +1,27 @@ +"""Cross-agent session budget cap (design: docs/dev/DESIGN-v1.0-budget-cap.md). + +* :mod:`~aegis_core.budget.policy` -- the signed ``budget.yaml``; +* :mod:`~aegis_core.budget.pricing` -- tokens to estimated dollars; +* :mod:`~aegis_core.budget.logs` -- reading usage from each agent's logs; +* :mod:`~aegis_core.budget.accounting` -- session records, project/day + totals, revisions, locks, the inventory and rebuilds; +* :mod:`~aegis_core.budget.evaluate` -- allow / warn / deny. +""" + +from aegis_core.budget.accounting import BudgetStore, Snapshot, rebuild, refresh +from aegis_core.budget.evaluate import ALLOW, DENY, WARN, Verdict, evaluate +from aegis_core.budget.logs import Locator, default_logs +from aegis_core.budget.policy import BUDGET_CLASS, BUDGET_FILE, BudgetPolicy, load_budget_policy +from aegis_core.budget.pricing import PriceTable, load_builtin_table + + +def price_table(policy: BudgetPolicy) -> PriceTable: + """The built-in prices with the policy's signed overrides.""" + return PriceTable(load_builtin_table(), policy.pricing) + + +__all__ = [ + "ALLOW", "BUDGET_CLASS", "BUDGET_FILE", "DENY", "WARN", "BudgetPolicy", "BudgetStore", + "Locator", "PriceTable", "Snapshot", "Verdict", "default_logs", "evaluate", + "load_budget_policy", "load_builtin_table", "price_table", "rebuild", "refresh", +] diff --git a/src/aegis_core/budget/accounting.py b/src/aegis_core/budget/accounting.py new file mode 100644 index 0000000..343a948 --- /dev/null +++ b/src/aegis_core/budget/accounting.py @@ -0,0 +1,608 @@ +"""Session and project/day usage accounting (design v1.0, §4.1). + +The agents' logs are the source of truth; everything under the cache +directory is derived from them and can be rebuilt. Three rules make +updates idempotent: + +* a session's usage is only ever **replaced** in the project map, never + added to a running total; +* a replacement only wins if its **revision** is newer, and revisions come + from a durable per-session counter outside the cache; +* a project/day total is always a **sum computed** from per-session entries. + +Locks are always taken in the order session -> project, and nothing takes a +session lock while holding a project lock, so hooks and rebuilds cannot +deadlock. + +Layout:: + + $XDG_CACHE_HOME/aegis/budget/ (default ~/.cache; disposable) + sessions//.json session record + projects//.json project contributions for a day + $XDG_STATE_HOME/aegis/budget/ (default ~/.local/state; durable) + sessions//.rev, .lock revision counter and lock + /.jsonl session inventory for a day + /project.json, .lock days with an inventory; project lock +""" + +from __future__ import annotations + +import contextlib +import fcntl +import hashlib +import json +import os +import re +import tempfile +import time +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from zoneinfo import ZoneInfo + +from aegis_core.budget.logs import AgentLog, Event, Locator, LogMissing, SourceRead, parse_ts + +UNKNOWN_DAY = "unknown" +CHECKPOINT_EVERY = 60.0 # seconds +_MAX_PROBLEMS = 20 +_SAFE_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}") + +Usage = dict[str, float] # category -> amount +ByModel = dict[str, Usage] # model -> usage +Buckets = dict[str, ByModel] # day -> model -> usage + + +# --- small helpers ------------------------------------------------------------------- + + +def add_usage(into: ByModel, model: str, usage: Mapping[str, float], sign: int = 1) -> None: + slot = into.setdefault(model, {}) + for cat, n in usage.items(): + slot[cat] = slot.get(cat, 0) + sign * n + if abs(slot[cat]) < 1e-9: + slot[cat] = 0 + if not any(slot.values()): + del into[model] + + +def sum_by_model(*parts: Mapping[str, Mapping[str, float]]) -> ByModel: + out: ByModel = {} + for part in parts: + for model, usage in part.items(): + add_usage(out, model, usage) + return out + + +def find_project_root(cwd: str | None) -> str | None: + """The nearest directory at or above ``cwd`` holding ``.aegis/`` (the + lookup the hook uses for policy); ``None`` if there is none.""" + if not cwd: + return None + start = Path(cwd) + for directory in (start, *start.parents): + if (directory / ".aegis").is_dir(): + return str(directory.resolve()) + return None + + +def project_key(root: str) -> str: + return hashlib.sha256(root.encode()).hexdigest() + + +def day_of(ts: datetime | None, zone: ZoneInfo) -> str: + return ts.astimezone(zone).date().isoformat() if ts else UNKNOWN_DAY + + +def day_start(day: str, zone: ZoneInfo) -> float: + d = datetime.fromisoformat(day).replace(tzinfo=zone) + return d.timestamp() + + +def _safe_name(session_id: str) -> str: + if _SAFE_ID.fullmatch(session_id): + return session_id + return "h-" + hashlib.sha256(session_id.encode()).hexdigest() + + +def _atomic_write(path: Path, text: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "w") as f: + f.write(text) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + except BaseException: + with contextlib.suppress(OSError): + os.unlink(tmp) + raise + + +def _read_json(path: Path) -> Any: + try: + return json.loads(path.read_text()) + except (OSError, json.JSONDecodeError, UnicodeDecodeError): + return None + + +@contextlib.contextmanager +def _flock(path: Path) -> Iterator[None]: + path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + yield + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +# --- records --------------------------------------------------------------------------- + + +@dataclass +class SessionRecord: + """The cached, derived state of one session (the record), plus the + revision it was written at.""" + + agent: str + session_id: str + locator: dict[str, str] + tz: str + revision: int = 0 + project_root: str | None = None + start: str | None = None + parent: str | None = None + sources: dict[str, dict[str, Any]] = field(default_factory=dict) + providers: dict[str, str] = field(default_factory=dict) + problems: list[str] = field(default_factory=list) + problem_total: int = 0 + log_missing: bool = False + inventory_days: list[str] = field(default_factory=list) + checkpoint_at: float = 0.0 + + @property + def key(self) -> str: + return f"{self.agent}/{self.session_id}" + + def buckets(self) -> Buckets: + """day -> model -> usage, summed over the session's sources.""" + out: Buckets = {} + for src in self.sources.values(): + for day, by_model in src.get("buckets", {}).items(): + target = out.setdefault(day, {}) + for model, usage in by_model.items(): + add_usage(target, model, usage) + return out + + def usage_on(self, day: str) -> ByModel: + return self.buckets().get(day, {}) + + def usage_total(self) -> ByModel: + return sum_by_model(*self.buckets().values()) + + def to_json(self) -> str: + return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> SessionRecord: + d = {k: v for k, v in d.items() if k in cls.__dataclass_fields__} + return cls(**d) + + +@dataclass +class ProjectDay: + """One project's contributions for one day.""" + + root: str + day: str + entries: dict[str, dict[str, Any]] = field(default_factory=dict) + completeness: str = "complete" # complete | partial | unknown + missing_sessions: int = 0 + known_sessions: int = 0 + + def usage(self) -> ByModel: + return sum_by_model(*(e.get("usage", {}) for e in self.entries.values())) + + def providers(self) -> dict[str, str]: + out: dict[str, str] = {} + for e in self.entries.values(): + out.update(e.get("providers", {})) + return out + + def to_json(self) -> str: + return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> ProjectDay: + d = {k: v for k, v in d.items() if k in cls.__dataclass_fields__} + return cls(**d) + + +@dataclass +class Snapshot: + """What a hook evaluates: the session's usage over its lifetime and the + project's usage today, with how much of either is unverifiable.""" + + record: SessionRecord + day: str + project: ProjectDay | None + session_usage: ByModel + project_usage: ByModel + providers: dict[str, str] + + +# --- the store ------------------------------------------------------------------------- + + +class BudgetStore: + """The cache and state directories, and the locks and files in them.""" + + def __init__(self, cache_dir: str | Path | None = None, state_dir: str | Path | None = None, + env: Mapping[str, str] | None = None): + env = os.environ if env is None else env + home = Path(env.get("HOME") or Path.home()) + self.cache = Path(cache_dir) if cache_dir else Path( + env.get("XDG_CACHE_HOME") or home / ".cache") / "aegis" / "budget" + self.state = Path(state_dir) if state_dir else Path( + env.get("XDG_STATE_HOME") or home / ".local" / "state") / "aegis" / "budget" + + # paths + def record_path(self, agent: str, sid: str) -> Path: + return self.cache / "sessions" / agent / f"{_safe_name(sid)}.json" + + def rev_path(self, agent: str, sid: str) -> Path: + return self.state / "sessions" / agent / f"{_safe_name(sid)}.rev" + + def session_lock(self, agent: str, sid: str): + return _flock(self.state / "sessions" / agent / f"{_safe_name(sid)}.lock") + + def project_dir(self, root: str) -> Path: + return self.state / project_key(root) + + def project_lock(self, root: str): + return _flock(self.project_dir(root) / "project.lock") + + def day_path(self, root: str, day: str) -> Path: + return self.cache / "projects" / project_key(root) / f"{day}.json" + + def inventory_path(self, root: str, day: str) -> Path: + return self.project_dir(root) / f"{day}.jsonl" + + def index_path(self, root: str) -> Path: + return self.project_dir(root) / "project.json" + + # session records + def load_record(self, agent: str, sid: str) -> SessionRecord | None: + d = _read_json(self.record_path(agent, sid)) + if not isinstance(d, dict): + return None + try: + return SessionRecord.from_dict(d) + except TypeError: + return None + + def save_record(self, rec: SessionRecord) -> None: + _atomic_write(self.record_path(rec.agent, rec.session_id), rec.to_json()) + + # the durable revision counter (caller holds the session lock) + def read_revision(self, agent: str, sid: str) -> int | None: + try: + return int(self.rev_path(agent, sid).read_text().strip()) + except (OSError, ValueError): + return None + + def write_revision(self, agent: str, sid: str, rev: int) -> None: + _atomic_write(self.rev_path(agent, sid), f"{rev}\n") + + def recover_revision(self, agent: str, sid: str, root: str | None, + rec: SessionRecord | None) -> int: + """The counter is missing (the state directory was deleted): one + more than the highest revision for the session anywhere a merge + could compare it with -- the cached record, every cached project + map entry and every inventory checkpoint for its project.""" + key = f"{agent}/{sid}" + best = rec.revision if rec else 0 + if root: + for p in (self.cache / "projects" / project_key(root)).glob("*.json"): + d = _read_json(p) + if isinstance(d, dict): + best = max(best, int((d.get("entries", {}).get(key) or {}).get("revision", + 0))) + for p in self.project_dir(root).glob("*.jsonl"): + for line in self._inventory_lines(p): + if line.get("session") == key: + best = max(best, int(line.get("revision", 0))) + return best + 1 + + # project maps (caller holds the project lock to write) + def load_day(self, root: str, day: str) -> ProjectDay | None: + d = _read_json(self.day_path(root, day)) + if not isinstance(d, dict): + return None + try: + return ProjectDay.from_dict(d) + except TypeError: + return None + + def save_day(self, pd: ProjectDay) -> None: + _atomic_write(self.day_path(pd.root, pd.day), pd.to_json()) + + @staticmethod + def merge_entries(pd: ProjectDay, incoming: Mapping[str, dict[str, Any]]) -> bool: + """The merge rule: set an entry only if its revision is newer. + Entries are merged one by one; a map is never replaced wholesale.""" + changed = False + for key, entry in incoming.items(): + stored = pd.entries.get(key) + if stored is None or int(entry["revision"]) > int(stored.get("revision", 0)): + pd.entries[key] = entry + changed = True + return changed + + # the inventory (caller holds the project lock to write) + @staticmethod + def _inventory_lines(path: Path) -> list[dict[str, Any]]: + out = [] + try: + with open(path) as f: + for line in f: + with contextlib.suppress(json.JSONDecodeError): + d = json.loads(line) + if isinstance(d, dict): + out.append(d) + except OSError: + pass + return out + + def inventory(self, root: str, day: str) -> list[dict[str, Any]] | None: + path = self.inventory_path(root, day) + return self._inventory_lines(path) if path.exists() else None + + def append_inventory(self, root: str, day: str, line: dict[str, Any]) -> None: + path = self.inventory_path(root, day) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a") as f: + f.write(json.dumps(line, sort_keys=True) + "\n") + f.flush() + os.fsync(f.fileno()) + index = _read_json(self.index_path(root)) or {} + days = set(index.get("days", [])) + if day not in days: + days.add(day) + _atomic_write(self.index_path(root), + json.dumps({"root": root, "days": sorted(days)}, sort_keys=True)) + + +# --- reading a session (steps 1-3 of the update protocol) -------------------------------- + + +def _apply(rec: SessionRecord, reads: list[SourceRead], zone: ZoneInfo) -> bool: + """Folds new events into the record's per-source buckets. Returns + whether anything changed.""" + start = parse_ts(rec.start) + changed = False + for sr in reads: + src = rec.sources.get(sr.source) + if src is None or sr.reset: + src = {"buckets": {}, "keyed": {}} + rec.sources[sr.source] = src + changed = True + if src.get("cursor") != sr.cursor: + changed = True + src["cursor"] = sr.cursor + buckets: Buckets = src.setdefault("buckets", {}) + keyed: dict[str, list[Any]] = src.setdefault("keyed", {}) + for ev in sr.events: + changed = True + day = day_of(ev.ts or start, zone) + if ev.provider: + rec.providers[ev.model] = ev.provider + if ev.key is not None: + old = keyed.get(ev.key) + if old is not None: + add_usage(buckets.setdefault(old[0], {}), old[1], old[2], sign=-1) + if not buckets[old[0]]: + del buckets[old[0]] + keyed[ev.key] = [day, ev.model, ev.usage] + add_usage(buckets.setdefault(day, {}), ev.model, ev.usage) + if sr.problems: + changed = True + rec.problem_total += len(sr.problems) + rec.problems = (rec.problems + sr.problems)[-_MAX_PROBLEMS:] + return changed + + +def update_session(store: BudgetStore, log: AgentLog, loc: Locator, zone: ZoneInfo, *, + fallback_cwd: str | None = None) -> SessionRecord: + """Steps 1-3 for one session. **The caller holds the session's lock.** + + Reads the log from where the record left off, folds the new usage in, + advances the durable revision counter and writes the record at the new + revision. A record that would not change is not rewritten (and the + revision is not advanced).""" + tz = str(zone.key) + rec = store.load_record(loc.agent, loc.session_id) + if rec is None or rec.tz != tz or rec.locator.get("path") != loc.path: + old_rev = rec.revision if rec else 0 + rec = SessionRecord(loc.agent, loc.session_id, loc.to_dict(), tz, revision=old_rev) + if rec.start is None and rec.project_root is None: + info = log.info(loc) + cwd = (info.cwd if info else None) or fallback_cwd + rec.project_root = find_project_root(cwd) + rec.start = info.start.isoformat() if info and info.start else None + rec.parent = info.parent if info else None + + rev = store.read_revision(loc.agent, loc.session_id) + if rev is None: + rev = store.recover_revision(loc.agent, loc.session_id, rec.project_root, rec) + store.write_revision(loc.agent, loc.session_id, rev) + rev = max(rev, rec.revision) + + try: + reads = log.read(loc, {name: s.get("cursor", {}) for name, s in rec.sources.items()}) + missing = False + except LogMissing as e: + reads, missing = [], True + rec.problems = (rec.problems + [str(e)])[-_MAX_PROBLEMS:] + changed = _apply(rec, reads, zone) or missing != rec.log_missing + rec.log_missing = missing + if changed or not store.record_path(loc.agent, loc.session_id).exists(): + rec.revision = rev + 1 + store.write_revision(loc.agent, loc.session_id, rec.revision) # durable first + store.save_record(rec) + return rec + + +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}} + + +# --- rebuild ------------------------------------------------------------------------- + + +def rebuild(store: BudgetStore, logs: Mapping[str, AgentLog], root: str, day: str, + zone: ZoneInfo, *, exclude: str | None = None) -> ProjectDay: + """Rebuilds a project's day from the logs and the inventory. + + 1. Without the project lock: the sessions to read are the ones + ``discover`` finds for the project today plus those in the inventory. + 2. Each session is brought up to date under **its own lock alone**. + 3. Under the project lock, every entry is merged by the merge rule, so a + newer entry written meanwhile by a hook is kept. + """ + since = day_start(day, zone) + candidates: dict[str, Locator] = {} + for agent, log in logs.items(): + for loc in log.discover(since): + key = f"{agent}/{loc.session_id}" + info = log.info(loc) + if info and info.parent is None and find_project_root(info.cwd) == root: + candidates[key] = loc + inventory = store.inventory(root, day) + checkpoints: dict[str, dict[str, Any]] = {} + in_inventory: set[str] = set() + for line in inventory or []: + key = line.get("session") + if not isinstance(key, str): + continue + in_inventory.add(key) + if line.get("event") == "seen" and key not in candidates: + with contextlib.suppress(KeyError, TypeError): + candidates[key] = Locator.from_dict(line["locator"]) + elif line.get("event") == "checkpoint": + if int(line.get("revision", 0)) >= int(checkpoints.get(key, {}).get("revision", 0)): + checkpoints[key] = line + + incoming: dict[str, dict[str, Any]] = {} + missing = 0 + for key, loc in candidates.items(): + if key == exclude or loc.agent not in logs: + continue + with store.session_lock(loc.agent, loc.session_id): + rec = update_session(store, logs[loc.agent], loc, zone) + if rec.log_missing and not rec.sources: + # Neither the log nor a cached record holds this session's usage: + # the inventory checkpoint is all that is left. + missing += 1 + cp = checkpoints.get(key) + if cp: + incoming[key] = {"revision": int(cp.get("revision", 0)), + "usage": cp.get("usage", {}), "providers": {}} + continue + incoming[key] = _entry(rec, day) + + if inventory is not None: + completeness = "partial" if missing else "complete" + else: + index = _read_json(store.index_path(root)) + if isinstance(index, dict): + completeness = "unknown" if day in index.get("days", []) else "complete" + else: + others = [k for k in candidates if k != exclude] + completeness = "unknown" if others else "complete" + + with store.project_lock(root): + pd = store.load_day(root, day) or ProjectDay(root, day) + store.merge_entries(pd, incoming) + pd.completeness = completeness + pd.missing_sessions = missing + pd.known_sessions = len(in_inventory) + store.save_day(pd) + return pd + + +# --- the per-hook refresh ------------------------------------------------------------ + + +def refresh(store: BudgetStore, logs: Mapping[str, AgentLog], loc: Locator, zone: ZoneInfo, *, + now: float | None = None, fallback_cwd: str | None = None) -> Snapshot: + """The update protocol for one hook call on session ``loc``. + + A linked sub-agent is accounted inside its parent: if the log names a + parent session that can be located, the parent is refreshed instead. + If the project's day file is missing, the day is rebuilt first, holding + no lock. Then, under the session lock: read the log (steps 1-3); under + the project lock as well: merge the session's entry and keep the + inventory (step 4).""" + now = time.time() if now is None else now + log = logs[loc.agent] + info = log.info(loc) + if info and info.parent: + parent = log.locate(info.parent) + if parent is not None: + loc, info = parent, log.info(parent) + day = day_of(datetime.fromtimestamp(now, tz=UTC), zone) + + peek = store.load_record(loc.agent, loc.session_id) + root = peek.project_root if peek else find_project_root( + (info.cwd if info else None) or fallback_cwd) + if root and store.load_day(root, day) is None: + rebuild(store, logs, root, day, zone, exclude=f"{loc.agent}/{loc.session_id}") + + pd: ProjectDay | None = None + with store.session_lock(loc.agent, loc.session_id): + rec = update_session(store, log, loc, zone, fallback_cwd=fallback_cwd) + if rec.project_root: + with store.project_lock(rec.project_root): + pd = store.load_day(rec.project_root, day) or ProjectDay(rec.project_root, day) + if store.merge_entries(pd, {rec.key: _entry(rec, day)}): + store.save_day(pd) + if _keep_inventory(store, rec, day, now): + store.save_record(rec) + + session_usage = rec.usage_total() + project_usage = pd.usage() if pd else {} + providers = {**(pd.providers() if pd else {}), **rec.providers} + return Snapshot(rec, day, pd, session_usage, project_usage, providers) + + +def _keep_inventory(store: BudgetStore, rec: SessionRecord, day: str, now: float) -> bool: + """Appends the session to the day's inventory the first time it is seen + that day, and a usage checkpoint at most once a minute. **The caller + holds the project lock.** Returns whether the record changed.""" + root = rec.project_root + assert root + changed = False + if day not in rec.inventory_days: + store.append_inventory(root, day, {"event": "seen", "session": rec.key, + "locator": rec.locator, "start": rec.start}) + rec.inventory_days = sorted({*rec.inventory_days, day})[-7:] + changed = True + if now - rec.checkpoint_at >= CHECKPOINT_EVERY: + store.append_inventory(root, day, {"event": "checkpoint", "session": rec.key, + "revision": rec.revision, + "usage": rec.usage_on(day), "at": now}) + rec.checkpoint_at = now + changed = True + return changed + + +def event_day(ev: Event, rec: SessionRecord, zone: ZoneInfo) -> str: + """The day an event belongs to (exposed for tests and status output).""" + return day_of(ev.ts or parse_ts(rec.start), zone) diff --git a/src/aegis_core/budget/evaluate.py b/src/aegis_core/budget/evaluate.py new file mode 100644 index 0000000..a222ec1 --- /dev/null +++ b/src/aegis_core/budget/evaluate.py @@ -0,0 +1,159 @@ +"""Deciding allow / warn / deny from a usage snapshot (design v1.0, §6). + +The guarantee is exactly the design's (§6.1): covered tool calls and +prompts are denied when the usage *available at that hook* reaches a limit. +Spend Aegis cannot see in time is not bounded; usage above a limit is +reported as observed overshoot. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from aegis_core.budget.accounting import UNKNOWN_DAY, Snapshot +from aegis_core.budget.policy import BudgetPolicy, Limit +from aegis_core.budget.pricing import CATEGORIES, PriceTable + +ALLOW, WARN, DENY = "allow", "warn", "deny" + + +@dataclass +class Measure: + """One limit, measured: ``used`` against ``limit`` in ``unit``.""" + + name: str # session | project_day | copilot_premium_requests + unit: str # usd | tokens | requests + used: float + limit: float + warn_at: float + + @property + def fraction(self) -> float: + return self.used / self.limit if self.limit else 0.0 + + @property + def level(self) -> str: + if self.used >= self.limit: + return DENY + if self.used >= self.warn_at * self.limit: + return WARN + return ALLOW + + @property + def overshoot(self) -> float: + return max(0.0, self.used - self.limit) + + def to_dict(self) -> dict[str, Any]: + return {"name": self.name, "unit": self.unit, "used": self.used, "limit": self.limit, + "warn_at": self.warn_at, "fraction": self.fraction, "level": self.level, + "overshoot": self.overshoot} + + +@dataclass +class Verdict: + decision: str # allow | warn | deny + measures: list[Measure] = field(default_factory=list) + reasons: list[str] = field(default_factory=list) # codes, for deny or warn + messages: list[str] = field(default_factory=list) # one line each, for people + estimated: dict[str, list[str]] = field(default_factory=dict) # model -> categories + + def to_dict(self) -> dict[str, Any]: + return {"decision": self.decision, "measures": [m.to_dict() for m in self.measures], + "reasons": self.reasons, "messages": self.messages, + "estimated": self.estimated} + + +def tokens_of(usage_by_model: Mapping[str, Mapping[str, float]]) -> float: + return float(sum(u.get(c, 0) for u in usage_by_model.values() for c in CATEGORIES)) + + +def _fmt(unit: str, amount: float) -> str: + if unit == "usd": + return f"${amount:,.2f}" + if unit == "tokens": + return f"{amount:,.0f} tokens" + 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) + + def amount(usage: Mapping[str, Mapping[str, float]]) -> float: + if policy.unit == "tokens": + return tokens_of(usage) + cost = prices.cost(usage, snap.providers) + for model, cats in cost.estimated.items(): + verdict.estimated[model] = sorted(set(verdict.estimated.get(model, [])) | set(cats)) + 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)) + + 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) + else: + if policy.session: + measure("session", policy.unit, 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] = [] + rec = snap.record + if rec.log_missing: + unknown.append("the session's log is missing") + if rec.problem_total: + 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)") + if snap.project is not None and snap.project.completeness != "complete": + unknown.append(f"today's project total is {snap.project.completeness} " + f"({snap.project.missing_sessions} session(s) without a log)") + 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 diff --git a/src/aegis_core/budget/logs.py b/src/aegis_core/budget/logs.py new file mode 100644 index 0000000..4c72272 --- /dev/null +++ b/src/aegis_core/budget/logs.py @@ -0,0 +1,678 @@ +"""Reading model usage from each agent's own session logs (design v1.0, §4). + +No hook payload carries token usage, so usage comes from the logs the agents +already write. One reader per agent, each recognising its format by +structure; anything it does not recognise is reported as a *problem* (an +"unknown log", handled by ``on_unknown_log``), never guessed at. + +Reads are incremental. A session's log is one or more **sources** (the main +log, and a linked sub-agent's logs), each with a cursor the caller keeps: +for a JSONL file the byte offset of the last complete line plus the file's +identity (device, inode, and a hash of its first bytes). A file that was +replaced or truncated comes back with ``reset=True`` and is read again from +the start. + +Usage categories (see :mod:`aegis_core.budget.pricing`): ``input`` is +uncached input, ``output`` includes reasoning, ``cache_read`` and +``cache_write``; Copilot CLI reports premium ``requests`` instead. + +Field names below were read from the logs of Claude Code 2.1.281, Codex CLI +0.160.1, Gemini CLI (session JSONL, 2026-09), OpenCode (SQLite) and Copilot +CLI 1.0.89 on 2026-10-07; ``tests/fixtures/budget/`` holds redacted samples. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import sqlite3 +from collections.abc import Iterator +from dataclasses import dataclass, field +from datetime import UTC, date, datetime, timedelta +from pathlib import Path +from typing import Any + +HEAD_BYTES = 4096 +_CLAUDE_SEEN_MAX = 512 + + +class LogMissing(Exception): + """The session's log cannot be found or read.""" + + +@dataclass(frozen=True) +class Locator: + """Where one session's log is: an agent, its session id and a path (the + log file, or OpenCode's database).""" + + agent: str + session_id: str + path: str + + def to_dict(self) -> dict[str, str]: + return {"agent": self.agent, "session_id": self.session_id, "path": self.path} + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> Locator: + return cls(str(d["agent"]), str(d["session_id"]), str(d["path"])) + + +@dataclass(frozen=True) +class SessionInfo: + """What the log records about the session itself: the working directory + it started in, its start time, and its parent session if it is a linked + sub-agent.""" + + cwd: str | None + start: datetime | None + parent: str | None = None + + +@dataclass +class Event: + """Usage of one model request (or, with ``key``, the current usage of a + record the agent may still rewrite: a later event with the same key + replaces the earlier one).""" + + ts: datetime | None + model: str + usage: dict[str, float] + provider: str | None = None + key: str | None = None + + +@dataclass +class SourceRead: + source: str + cursor: dict[str, Any] + events: list[Event] = field(default_factory=list) + reset: bool = False + problems: list[str] = field(default_factory=list) + + +def parse_ts(value: Any) -> datetime | None: + """ISO-8601 strings and epoch milliseconds -> aware UTC datetimes.""" + if isinstance(value, bool) or value is None: + return None + if isinstance(value, (int, float)): + try: + return datetime.fromtimestamp(value / 1000, tz=UTC) + except (OverflowError, OSError, ValueError): + return None + if isinstance(value, str): + try: + dt = datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + return dt if dt.tzinfo else dt.replace(tzinfo=UTC) + return None + + +def _ts_str(dt: datetime | None) -> str | None: + return dt.astimezone(UTC).isoformat() if dt else None + + +def _int(value: Any) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else 0 + + +def _head_hash(path: Path, length: int) -> str: + with open(path, "rb") as f: + return hashlib.sha256(f.read(length)).hexdigest() + + +def tail_jsonl(path: Path, cursor: dict[str, Any] | None + ) -> tuple[list[Any], list[str], dict[str, Any], bool]: + """New complete lines of a JSONL file since ``cursor``. + + Returns ``(records, problems, new_cursor, reset)``. A half-written last + line is left for the next read. ``reset`` means the file is not the one + the cursor described (another inode, a different head, or shorter than + the offset), so it was read from the start; the caller discards what it + counted from it before. Raises :class:`LogMissing` if it cannot be read. + """ + try: + st = os.stat(path) + except OSError as e: + raise LogMissing(f"{path}: {e.strerror or e}") from None + reset = False + offset = 0 + if cursor: + same = (cursor.get("dev"), cursor.get("ino")) == (st.st_dev, st.st_ino) + if (not same or st.st_size < cursor.get("offset", 0) + or _head_hash(path, cursor.get("head_len", 0)) != cursor.get("head")): + reset = True + else: + offset = cursor.get("offset", 0) + try: + with open(path, "rb") as f: + f.seek(offset) + data = f.read() + except OSError as e: + raise LogMissing(f"{path}: {e.strerror or e}") from None + end = data.rfind(b"\n") + complete = data[: end + 1] if end >= 0 else b"" + new_offset = offset + len(complete) + records: list[Any] = [] + problems: list[str] = [] + for n, line in enumerate(complete.splitlines()): + if not line.strip(): + continue + try: + records.append(json.loads(line)) + except (json.JSONDecodeError, UnicodeDecodeError): + problems.append(f"{path}: unreadable line at byte {offset}+{n}") + head_len = min(HEAD_BYTES, new_offset) + base = {} if reset or not cursor else dict(cursor) + base.update({"dev": st.st_dev, "ino": st.st_ino, "offset": new_offset, + "head_len": head_len, "head": _head_hash(path, head_len)}) + return records, problems, base, reset + + +def _first_json_line(path: Path) -> Any: + try: + with open(path, "rb") as f: + line = f.readline() + return json.loads(line) if line.strip() else None + except (OSError, json.JSONDecodeError, UnicodeDecodeError): + return None + + +def _modified_since(paths: Iterator[Path], since: float) -> Iterator[Path]: + for p in paths: + try: + if p.is_file() and p.stat().st_mtime >= since: + yield p + except OSError: + continue + + +class AgentLog: + """One agent's log reader. ``base`` is the agent's data directory + (``~/.claude``, ``~/.codex``, ...), injectable for tests.""" + + agent = "" + + def __init__(self, base: str | Path): + self.base = Path(base) + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + raise NotImplementedError + + def info(self, loc: Locator) -> SessionInfo | None: + raise NotImplementedError + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + """One :class:`SourceRead` per source of the session (its main log + and linked sub-agents). Raises :class:`LogMissing` if the main log + is gone.""" + raise NotImplementedError + + def discover(self, since: float) -> list[Locator]: + """Top-level sessions (not linked sub-agents) whose logs changed at + or after ``since`` (epoch seconds).""" + raise NotImplementedError + + # -- shared JSONL source reading ---------------------------------------- + + def _read_jsonl_source(self, source: str, path: Path, cursor: dict[str, Any] | None, + handle) -> SourceRead: + records, problems, cur, reset = tail_jsonl(path, cursor) + if reset: + cur = {k: cur[k] for k in ("dev", "ino", "offset", "head_len", "head")} + out = SourceRead(source, cur, reset=reset, problems=problems) + last_ts = parse_ts(cur.get("last_ts")) + for rec in records: + if not isinstance(rec, dict): + out.problems.append(f"{path}: a record is not a JSON object") + continue + first_new = len(out.events) + ts = handle(rec, cur, out) + if ts is not None: + last_ts = ts + for ev in out.events[first_new:]: + if ev.ts is None: + ev.ts = last_ts # the nearest earlier timestamp in this source + cur["last_ts"] = _ts_str(last_ts) + return out + + +# --- Claude Code ------------------------------------------------------------------ + + +class ClaudeLog(AgentLog): + """``~/.claude/projects//.jsonl``; sub-agents (linked) in + ``//subagents/agent-*.jsonl``. Each response's usage is + repeated on every content block of the message, so a request is counted + once per ``(message.id, requestId)``.""" + + agent = "claude" + + def _projects(self) -> Path: + return self.base / "projects" + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + if transcript_path and Path(transcript_path).is_file(): + return Locator(self.agent, session_id, str(transcript_path)) + for p in self._projects().glob(f"*/{session_id}.jsonl"): + return Locator(self.agent, session_id, str(p)) + return None + + def info(self, loc: Locator) -> SessionInfo | None: + start = None + try: + with open(loc.path, "rb") as f: + for n, line in enumerate(f): + if n > 200: + break + try: + rec = json.loads(line) + except (json.JSONDecodeError, UnicodeDecodeError): + continue + if not isinstance(rec, dict): + continue + start = start or parse_ts(rec.get("timestamp")) + if isinstance(rec.get("cwd"), str): + return SessionInfo(rec["cwd"], parse_ts(rec.get("timestamp")) or start) + except OSError: + return None + return SessionInfo(None, start) + + def _handle(self, rec: dict[str, Any], cur: dict[str, Any], out: SourceRead): + ts = parse_ts(rec.get("timestamp")) + if rec.get("type") != "assistant": + return ts + msg = rec.get("message") + usage = msg.get("usage") if isinstance(msg, dict) else None + if not isinstance(usage, dict): + out.problems.append("claude: an assistant record without message.usage") + return ts + model = msg.get("model") + if not isinstance(model, str) or model == "": + return ts + key = f"{msg.get('id')}:{rec.get('requestId')}" + seen = cur.setdefault("seen", []) + if key in seen: + return ts + seen.append(key) + del seen[:-_CLAUDE_SEEN_MAX] + counts = {"input": _int(usage.get("input_tokens")), + "output": _int(usage.get("output_tokens")), + "cache_read": _int(usage.get("cache_read_input_tokens")), + "cache_write": _int(usage.get("cache_creation_input_tokens"))} + if any(counts.values()): + out.events.append(Event(ts, model, counts, "anthropic")) + return ts + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + main = Path(loc.path) + reads = [self._read_jsonl_source("main", main, cursors.get("main"), self._handle)] + subdir = main.with_suffix("") / "subagents" + if subdir.is_dir(): + for p in sorted(subdir.glob("agent-*.jsonl")): + name = f"sub:{p.name}" + try: + reads.append(self._read_jsonl_source(name, p, cursors.get(name), + self._handle)) + except LogMissing: + continue + return reads + + def discover(self, since: float) -> list[Locator]: + return [Locator(self.agent, p.stem, str(p)) + for p in _modified_since(self._projects().glob("*/*.jsonl"), since)] + + +# --- Codex CLI -------------------------------------------------------------------- + + +def _codex_meta(path: Path) -> dict[str, Any] | None: + first = _first_json_line(path) + if isinstance(first, dict) and first.get("type") == "session_meta": + payload = first.get("payload") + return payload if isinstance(payload, dict) else None + return None + + +class CodexLog(AgentLog): + """``~/.codex/sessions/YYYY/MM/DD/rollout-*.jsonl``. ``token_count`` + events carry running totals; usage is the difference between successive + totals. The model is the latest ``turn_context`` model. A sub-agent has + its own rollout file whose ``session_meta`` names the parent + (``thread_source: subagent``, ``parent_thread_id``).""" + + agent = "codex" + + def _sessions(self) -> Path: + return self.base / "sessions" + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + if transcript_path and Path(transcript_path).is_file(): + return Locator(self.agent, session_id, str(transcript_path)) + for p in self._sessions().glob(f"*/*/*/rollout-*-{session_id}.jsonl"): + return Locator(self.agent, session_id, str(p)) + return None + + def info(self, loc: Locator) -> SessionInfo | None: + meta = _codex_meta(Path(loc.path)) + if meta is None: + return None + parent = meta.get("parent_thread_id") if meta.get("thread_source") == "subagent" else None + cwd = meta.get("cwd") if isinstance(meta.get("cwd"), str) else None + return SessionInfo(cwd, parse_ts(meta.get("timestamp")), parent) + + def _handle(self, rec: dict[str, Any], cur: dict[str, Any], out: SourceRead): + ts = parse_ts(rec.get("timestamp")) + kind = rec.get("type") + payload = rec.get("payload") + if not isinstance(payload, dict): + return ts + if kind == "session_meta" and isinstance(payload.get("model_provider"), str): + cur["provider"] = payload["model_provider"] + elif kind == "turn_context" and isinstance(payload.get("model"), str): + cur["model"] = payload["model"] + elif kind == "event_msg" and payload.get("type") == "token_count": + info = payload.get("info") + if info is None: + return ts # rate-limit only update + total = info.get("total_token_usage") if isinstance(info, dict) else None + if not isinstance(total, dict): + out.problems.append("codex: a token_count event without total_token_usage") + return ts + cached = _int(total.get("cached_input_tokens")) + now = {"input": max(0, _int(total.get("input_tokens")) - cached), + "output": _int(total.get("output_tokens")), + "cache_read": cached, + "cache_write": _int(total.get("cache_write_input_tokens"))} + before = cur.get("totals") or {} + delta = {c: max(0, now[c] - before.get(c, 0)) for c in now} + cur["totals"] = now + if any(delta.values()): + out.events.append(Event(ts, cur.get("model") or "unknown", + delta, cur.get("provider") or "openai")) + return ts + + def _children(self, loc: Locator, main_cursor: dict[str, Any]) -> dict[str, str]: + """Rollout files whose session_meta links them to this session, as + ``{file name: path}`` with ``""`` for files that are not children. + Classifications are kept in the main cursor, so each file's first + line is read once.""" + main = Path(loc.path) + known: dict[str, str] = dict(main_cursor.get("children") or {}) + start = parse_ts((_codex_meta(main) or {}).get("timestamp")) + floor = (start - timedelta(days=1)).date() if start else date.min + for day_dir in self._sessions().glob("*/*/*"): + try: + d = date(int(day_dir.parts[-3]), int(day_dir.parts[-2]), int(day_dir.parts[-1])) + except (ValueError, IndexError): + continue + if d < floor: + continue + for p in day_dir.glob("rollout-*.jsonl"): + if p == main or p.name in known: + continue + meta = _codex_meta(p) or {} + linked = (meta.get("thread_source") == "subagent" + and loc.session_id in (meta.get("parent_thread_id"), + meta.get("session_id"))) + known[p.name] = str(p) if linked else "" + return known + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + main = self._read_jsonl_source("main", Path(loc.path), cursors.get("main"), self._handle) + children = self._children(loc, cursors.get("main") or {}) + main.cursor["children"] = children + reads = [main] + for name, path in sorted(children.items()): + if not path: + continue + source = f"child:{name}" + try: + reads.append(self._read_jsonl_source(source, Path(path), cursors.get(source), + self._handle)) + except LogMissing: + continue + return reads + + def discover(self, since: float) -> list[Locator]: + out = [] + for p in _modified_since(self._sessions().glob("*/*/*/rollout-*.jsonl"), since): + meta = _codex_meta(p) + if not meta or meta.get("thread_source") == "subagent": + continue + sid = meta.get("id") or meta.get("session_id") + if isinstance(sid, str): + out.append(Locator(self.agent, sid, str(p))) + return out + + +# --- Gemini CLI ------------------------------------------------------------------- + + +class GeminiLog(AgentLog): + """``~/.gemini/tmp//chats/session-*.jsonl``; the project's + directory is in ``/.project_root``. Messages are written more + than once as they update, so each is keyed by its id.""" + + agent = "gemini" + + def _chats(self) -> Iterator[Path]: + return (self.base / "tmp").glob("*/chats/session-*.jsonl") + + @staticmethod + def _header(path: Path) -> dict[str, Any] | None: + first = _first_json_line(path) + return first if isinstance(first, dict) and "sessionId" in first else None + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + if transcript_path and Path(transcript_path).is_file(): + return Locator(self.agent, session_id, str(transcript_path)) + for p in self._chats(): + if session_id[:8] in p.name and (self._header(p) or {}).get("sessionId") == session_id: + return Locator(self.agent, session_id, str(p)) + return None + + def info(self, loc: Locator) -> SessionInfo | None: + path = Path(loc.path) + header = self._header(path) or {} + try: + cwd = (path.parent.parent / ".project_root").read_text().strip() or None + except OSError: + cwd = None + return SessionInfo(cwd, parse_ts(header.get("startTime"))) + + def _message(self, m: Any, out: SourceRead) -> None: + if not isinstance(m, dict) or m.get("type") != "gemini" or "tokens" not in m: + return + t = m["tokens"] + if not isinstance(t, dict): + out.problems.append("gemini: a message whose tokens is not a mapping") + return + cached = _int(t.get("cached")) + counts = {"input": max(0, _int(t.get("input")) - cached) + _int(t.get("tool")), + "output": _int(t.get("output")) + _int(t.get("thoughts")), + "cache_read": cached, "cache_write": 0} + model = m.get("model") if isinstance(m.get("model"), str) else "unknown" + out.events.append(Event(parse_ts(m.get("timestamp")), model, counts, "google", + key=str(m.get("id")))) + + def _handle(self, rec: dict[str, Any], cur: dict[str, Any], out: SourceRead): + if rec.get("type") == "gemini": + self._message(rec, out) + elif isinstance(rec.get("$set"), dict): + for m in rec["$set"].get("messages") or []: + self._message(m, out) + return parse_ts(rec.get("timestamp") or rec.get("startTime")) + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + return [self._read_jsonl_source("main", Path(loc.path), cursors.get("main"), + self._handle)] + + def discover(self, since: float) -> list[Locator]: + out = [] + for p in _modified_since(self._chats(), since): + sid = (self._header(p) or {}).get("sessionId") + if isinstance(sid, str): + out.append(Locator(self.agent, sid, str(p))) + return out + + +# --- OpenCode --------------------------------------------------------------------- + + +class OpenCodeLog(AgentLog): + """``opencode.db`` (SQLite, opened read-only; only the ``session`` and + ``message`` tables are read). Message rows are updated in place as a + response completes, so each is keyed by its id; sub-agents are child + sessions (``session.parent_id``).""" + + agent = "opencode" + + def _db(self) -> Path: + return self.base / "opencode.db" + + def _connect(self, path: str) -> sqlite3.Connection: + if not Path(path).is_file(): + raise LogMissing(f"{path}: no such file") + try: + return sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=1.0) + except sqlite3.Error as e: + raise LogMissing(f"{path}: {e}") from None + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + try: + with self._connect(str(self._db())) as db: + row = db.execute("SELECT id FROM session WHERE id = ?", (session_id,)).fetchone() + except (LogMissing, sqlite3.Error): + return None + return Locator(self.agent, session_id, str(self._db())) if row else None + + def info(self, loc: Locator) -> SessionInfo | None: + try: + with self._connect(loc.path) as db: + row = db.execute("SELECT directory, time_created, parent_id FROM session " + "WHERE id = ?", (loc.session_id,)).fetchone() + except (LogMissing, sqlite3.Error): + return None + if not row: + return None + return SessionInfo(row[0], parse_ts(row[1]), row[2]) + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + cur = dict(cursors.get("db") or {}) + out = SourceRead("db", cur) + try: + with self._connect(loc.path) as db: + if not db.execute("SELECT 1 FROM session WHERE id = ?", + (loc.session_id,)).fetchone(): + raise LogMissing(f"{loc.path}: session {loc.session_id} not found") + rows = db.execute( + "WITH RECURSIVE tree(id) AS (SELECT ? UNION " + "SELECT s.id FROM session s JOIN tree ON s.parent_id = tree.id) " + "SELECT m.id, m.time_created, m.time_updated, m.data FROM message m " + "WHERE m.session_id IN (SELECT id FROM tree) AND m.time_updated >= ? " + "ORDER BY m.time_updated", (loc.session_id, cur.get("since", 0))).fetchall() + except sqlite3.Error as e: + raise LogMissing(f"{loc.path}: {e}") from None + for mid, created, updated, data in rows: + cur["since"] = max(cur.get("since", 0), updated) + try: + m = json.loads(data) + except (json.JSONDecodeError, TypeError): + out.problems.append(f"opencode: message {mid} is not JSON") + continue + if not isinstance(m, dict) or m.get("role") != "assistant": + continue + t = m.get("tokens") + if t is None: + continue + if not isinstance(t, dict): + out.problems.append(f"opencode: message {mid} has malformed tokens") + continue + cache = t.get("cache") if isinstance(t.get("cache"), dict) else {} + counts = {"input": _int(t.get("input")), + "output": _int(t.get("output")) + _int(t.get("reasoning")), + "cache_read": _int(cache.get("read")), + "cache_write": _int(cache.get("write"))} + times = m.get("time") if isinstance(m.get("time"), dict) else {} + ts = parse_ts(times.get("completed") or times.get("created") or created) + model = m.get("modelID") if isinstance(m.get("modelID"), str) else "unknown" + provider = m.get("providerID") if isinstance(m.get("providerID"), str) else None + out.events.append(Event(ts, model, counts, provider, key=str(mid))) + return [out] + + def discover(self, since: float) -> list[Locator]: + try: + with self._connect(str(self._db())) as db: + rows = db.execute("SELECT id FROM session WHERE parent_id IS NULL AND " + "time_updated >= ?", (int(since * 1000),)).fetchall() + except (LogMissing, sqlite3.Error): + return [] + return [Locator(self.agent, r[0], str(self._db())) for r in rows] + + +# --- Copilot CLI ------------------------------------------------------------------ + + +class CopilotLog(AgentLog): + """``~/.copilot/session-state//events.jsonl``. Copilot CLI + writes token counts only at shutdown, but running totals of premium + requests in ``session.usage_checkpoint`` (and ``session.shutdown``), so + it is measured in premium requests.""" + + agent = "copilot" + MODEL = "copilot-premium-requests" + + def _state(self) -> Path: + return self.base / "session-state" + + def locate(self, session_id: str, transcript_path: str | None = None) -> Locator | None: + p = self._state() / session_id / "events.jsonl" + return Locator(self.agent, session_id, str(p)) if p.is_file() else None + + def info(self, loc: Locator) -> SessionInfo | None: + first = _first_json_line(Path(loc.path)) + if not isinstance(first, dict) or first.get("type") != "session.start": + return None + data = first.get("data") if isinstance(first.get("data"), dict) else {} + ctx = data.get("context") if isinstance(data.get("context"), dict) else {} + cwd = ctx.get("cwd") if isinstance(ctx.get("cwd"), str) else None + return SessionInfo(cwd, parse_ts(data.get("startTime") or first.get("timestamp"))) + + def _handle(self, rec: dict[str, Any], cur: dict[str, Any], out: SourceRead): + ts = parse_ts(rec.get("timestamp")) + if rec.get("type") in ("session.usage_checkpoint", "session.shutdown"): + data = rec.get("data") if isinstance(rec.get("data"), dict) else {} + total = data.get("totalPremiumRequests") + if isinstance(total, bool) or not isinstance(total, (int, float)): + out.problems.append(f"copilot: {rec.get('type')} without totalPremiumRequests") + return ts + delta = max(0.0, float(total) - float(cur.get("premium", 0.0))) + cur["premium"] = max(float(total), float(cur.get("premium", 0.0))) + if delta: + out.events.append(Event(ts, self.MODEL, {"requests": delta}, "github")) + return ts + + def read(self, loc: Locator, cursors: dict[str, dict[str, Any]]) -> list[SourceRead]: + return [self._read_jsonl_source("main", Path(loc.path), cursors.get("main"), + self._handle)] + + def discover(self, since: float) -> list[Locator]: + return [Locator(self.agent, p.parent.name, str(p)) + for p in _modified_since(self._state().glob("*/events.jsonl"), since)] + + +def default_logs(env: dict[str, str] | None = None) -> dict[str, AgentLog]: + """The readers for this machine, honouring each agent's own location + overrides (``CLAUDE_CONFIG_DIR``, ``CODEX_HOME``, ``XDG_DATA_HOME``).""" + env = dict(os.environ if env is None else env) + home = Path(env.get("HOME") or Path.home()) + data_home = Path(env.get("XDG_DATA_HOME") or home / ".local" / "share") + return { + "claude": ClaudeLog(env.get("CLAUDE_CONFIG_DIR") or home / ".claude"), + "codex": CodexLog(env.get("CODEX_HOME") or home / ".codex"), + "gemini": GeminiLog(home / ".gemini"), + "opencode": OpenCodeLog(data_home / "opencode"), + "copilot": CopilotLog(home / ".copilot"), + } diff --git a/src/aegis_core/budget/policy.py b/src/aegis_core/budget/policy.py new file mode 100644 index 0000000..55db13f --- /dev/null +++ b/src/aegis_core/budget/policy.py @@ -0,0 +1,232 @@ +"""The budget policy (``budget.yaml``, design v1.0 §3). + +Signed like every policy file; its ``principal`` must hold the ``budget`` +class in ``authority.yaml``, so raising a limit is a policy change an agent +cannot make by editing a file. Any shape problem is a load error, never a +guess. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +import yaml + +from aegis_core.budget.pricing import check_override +from aegis_core.signing import check_signature + +BUDGET_FILE = "budget.yaml" +BUDGET_CLASS = "budget" +AGENTS = ("claude", "codex", "gemini", "opencode", "copilot") +UNITS = ("usd", "tokens") +ON_UNKNOWN_LOG = ("allow", "deny") +ON_UNKNOWN_PRICE = ("estimate", "deny") +DEFAULT_WARN_AT = 0.8 + +_TOP_LEVEL = {"version", "principal", "unit", "session", "project_day", "agents", + "on_unknown_log", "on_unknown_price", "pricing", "copilot"} + + +@dataclass(frozen=True) +class Limit: + limit: float + warn_at: float = DEFAULT_WARN_AT + + +@dataclass(frozen=True) +class BudgetPolicy: + path: str + principal: str + unit: str + session: Limit | None + project_day: Limit | None + tz: str + agents: tuple[str, ...] + on_unknown_log: str + on_unknown_price: str + pricing: dict[str, dict[str, float]] = field(default_factory=dict) + copilot_premium_requests: Limit | None = None + warnings: tuple[str, ...] = () + + @property + def zone(self) -> ZoneInfo: + return ZoneInfo(self.tz) + + def measures(self, agent: str) -> bool: + """Only listed agents are measured and enforced (§11, decision 4). + Copilot CLI has no token data, so it is measured only when a + premium-request limit is set.""" + if agent not in self.agents: + return False + if agent == "copilot": + return self.copilot_premium_requests is not None + return self.session is not None or self.project_day is not None + + def to_dict(self) -> dict[str, Any]: + def lim(x: Limit | None) -> dict[str, float] | None: + return None if x is None else {"limit": x.limit, "warn_at": x.warn_at} + + return {"path": self.path, "principal": self.principal, "unit": self.unit, + "session": lim(self.session), "project_day": lim(self.project_day), + "tz": self.tz, "agents": list(self.agents), + "on_unknown_log": self.on_unknown_log, + "on_unknown_price": self.on_unknown_price, "pricing": self.pricing, + "copilot_premium_requests": lim(self.copilot_premium_requests), + "warnings": list(self.warnings)} + + +def _number(where: str, value: Any, *, integer: bool) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{where}: must be a number") + if integer and value != int(value): + raise ValueError(f"{where}: must be a whole number of tokens") + if value <= 0: + raise ValueError(f"{where}: must be greater than 0") + return float(value) + + +def _limit(path: Path, name: str, raw: Any, *, integer: bool, + extra: frozenset[str] = frozenset()) -> tuple[Limit, dict[str, Any]]: + where = f"{path}: {name}" + if not isinstance(raw, dict): + raise ValueError(f"{where}: must be a mapping with 'limit' (and optional 'warn_at')") + allowed = {"limit", "warn_at"} | extra + unknown = set(raw) - allowed + if unknown: + raise ValueError(f"{where}: unknown field(s) {sorted(unknown)}") + if "limit" not in raw: + raise ValueError(f"{where}: 'limit' is required") + limit = _number(f"{where}.limit", raw["limit"], integer=integer) + warn_at = raw.get("warn_at", DEFAULT_WARN_AT) + if isinstance(warn_at, bool) or not isinstance(warn_at, (int, float)) or not ( + 0 < warn_at < 1 + ): + raise ValueError(f"{where}.warn_at: must be a fraction between 0 and 1 (e.g. 0.8)") + return Limit(limit, float(warn_at)), {k: raw[k] for k in extra if k in raw} + + +def load_budget_policy( + path: str | Path, + *, + authority_map: dict[str, set[str]] | None, + key: bytes | None = None, + insecure: bool = False, +) -> BudgetPolicy: + """``budget.yaml`` -> :class:`BudgetPolicy`. Shape:: + + version: 1 + principal: admin # must hold the 'budget' class + unit: usd # usd | tokens; required + session: {limit: 20, warn_at: 0.8} + project_day: {limit: 100, warn_at: 0.8, tz: America/Chicago} + agents: [claude, codex, gemini, opencode, copilot] + on_unknown_log: allow # allow | deny + on_unknown_price: estimate # estimate | deny (usd only) + pricing: # 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} + """ + path = Path(path) + warnings: list[str] = [] + check_signature(path, key, insecure, warnings) + raw = yaml.safe_load(path.read_text()) + if not isinstance(raw, dict): + raise ValueError(f"{path}: must be a mapping") + unknown = set(raw) - _TOP_LEVEL + if unknown: + raise ValueError(f"{path}: unknown field(s) {sorted(unknown)}") + if raw.get("version", 1) != 1: + raise ValueError(f"{path}: unsupported version {raw.get('version')!r} (expected 1)") + + principal = raw.get("principal") + if not isinstance(principal, str) or not principal: + raise ValueError(f"{path}: 'principal' is required (who asserts this budget)") + if authority_map is None: + raise ValueError(f"{path}: an authority map is required to check '{principal}'") + if BUDGET_CLASS not in authority_map.get(principal, set()): + raise ValueError(f"{path}: principal {principal!r} does not hold the " + f"'{BUDGET_CLASS}' class in the authority map") + + if "unit" not in raw: + raise ValueError(f"{path}: 'unit' is required: usd (estimated dollars) or tokens") + unit = raw["unit"] + if unit not in UNITS: + raise ValueError(f"{path}: unit must be one of {list(UNITS)}, got {unit!r}") + integer = unit == "tokens" + + session = None + if raw.get("session") is not None: + session, _ = _limit(path, "session", raw["session"], integer=integer) + project_day, tz = None, "UTC" + if raw.get("project_day") is not None: + project_day, extra = _limit(path, "project_day", raw["project_day"], integer=integer, + extra=frozenset({"tz"})) + tz = extra.get("tz", "UTC") + if not isinstance(tz, str): + raise ValueError(f"{path}: project_day.tz must be a time zone name") + try: + ZoneInfo(tz) + except (ZoneInfoNotFoundError, ValueError): + raise ValueError(f"{path}: project_day.tz: unknown time zone {tz!r}") from None + + copilot = None + if raw.get("copilot") is not None: + c = raw["copilot"] + if not isinstance(c, dict) or "premium_requests" not in c: + raise ValueError(f"{path}: copilot must be a mapping with 'premium_requests'") + unknown_c = set(c) - {"premium_requests", "warn_at"} + if unknown_c: + raise ValueError(f"{path}: copilot: unknown field(s) {sorted(unknown_c)}") + copilot, _ = _limit(path, "copilot", {"limit": c["premium_requests"], + **({"warn_at": c["warn_at"]} if "warn_at" in c + else {})}, integer=True) + + if session is None and project_day is None and copilot is None: + raise ValueError(f"{path}: set at least one limit: session, project_day or " + "copilot.premium_requests") + + agents_raw = raw.get("agents", list(AGENTS)) + if not isinstance(agents_raw, list) or not agents_raw or not all( + isinstance(a, str) for a in agents_raw + ): + raise ValueError(f"{path}: agents must be a non-empty list of {list(AGENTS)}") + bad = [a for a in agents_raw if a not in AGENTS] + if bad: + raise ValueError(f"{path}: unknown agent(s) {bad}; supported: {list(AGENTS)}") + if len(set(agents_raw)) != len(agents_raw): + raise ValueError(f"{path}: an agent is listed twice in 'agents'") + agents = tuple(a for a in AGENTS if a in agents_raw) + + on_unknown_log = raw.get("on_unknown_log", "allow") + if on_unknown_log not in ON_UNKNOWN_LOG: + raise ValueError(f"{path}: on_unknown_log must be one of {list(ON_UNKNOWN_LOG)}") + on_unknown_price = raw.get("on_unknown_price", "estimate") + if on_unknown_price not in ON_UNKNOWN_PRICE: + raise ValueError(f"{path}: on_unknown_price must be one of {list(ON_UNKNOWN_PRICE)}") + + pricing_raw = raw.get("pricing") or {} + if not isinstance(pricing_raw, dict): + raise ValueError(f"{path}: pricing must be a mapping of model id to prices") + pricing = {str(model): check_override(f"{path}: pricing.{model}", entry) + for model, entry in pricing_raw.items()} + + if unit == "tokens": + if "on_unknown_price" in raw: + warnings.append("ignored: on_unknown_price has no effect with unit: tokens") + if pricing: + warnings.append("ignored: pricing has no effect with unit: tokens") + if "copilot" in agents and copilot is None: + warnings.append("copilot-unmeasured: copilot is listed but copilot.premium_requests " + "is not set; Copilot CLI reports no tokens, so its sessions are not " + "measured") + if copilot is not None and "copilot" not in agents: + warnings.append("ignored: copilot.premium_requests is set but copilot is not in agents") + if session is None and project_day is None and any(a != "copilot" for a in agents): + warnings.append("token-agents-unmeasured: no session or project_day limit, so only " + "Copilot CLI is measured") + + return BudgetPolicy(str(path), principal, unit, session, project_day, tz, agents, + on_unknown_log, on_unknown_price, pricing, copilot, tuple(warnings)) diff --git a/src/aegis_core/budget/prices.yaml b/src/aegis_core/budget/prices.yaml new file mode 100644 index 0000000..67a491e --- /dev/null +++ b/src/aegis_core/budget/prices.yaml @@ -0,0 +1,77 @@ +# Built-in model prices for the budget cap (docs/dev/DESIGN-v1.0-budget-cap.md §5). +# +# USD per million tokens, standard (non-batch, global) tier, per billing category: +# input uncached input tokens +# output output tokens, reasoning/thinking included +# cache_read cache hits +# cache_write cache writes +# A category left out is unknown for that model and falls back like an unknown model +# (on_unknown_price), for that category only. +# +# Choices that keep an estimate from running low: +# - Anthropic cache_write is the 1-hour rate (2x input); 5-minute writes cost less. +# - Google: the regular price, not the promotional one (Flash 3.6-3.8 are half price until +# 2026-12-31), and the >200k-token tier where prices depend on prompt length. Google has no +# per-token cache write price (storage is billed per hour), so cache_write is the input rate. +# - OpenAI cache_write is the input rate. +# +# Sources, read 2026-10-07: +# https://platform.claude.com/docs/en/about-claude/pricing +# https://developers.openai.com/api/docs/pricing +# https://ai.google.dev/gemini-api/docs/pricing +version: 1 +as_of: "2026-10-07" +vendors: + anthropic: + prefixes: [claude-] + models: + claude-fable-5-1: {input: 10.0, output: 50.0, cache_read: 0.25, cache_write: 20.0} + claude-mythos-5-1: {input: 10.0, output: 50.0, cache_read: 0.25, cache_write: 20.0} + claude-fable-5: {input: 10.0, output: 50.0, cache_read: 1.0, cache_write: 20.0} + claude-mythos-5: {input: 10.0, output: 50.0, cache_read: 1.0, cache_write: 20.0} + claude-opus-5-5: {input: 4.0, output: 20.0, cache_read: 0.20, cache_write: 8.0} + claude-opus-5: {input: 5.0, output: 25.0, cache_read: 0.50, cache_write: 10.0} + claude-opus-4-8: {input: 5.0, output: 25.0, cache_read: 0.50, cache_write: 10.0} + claude-opus-4-7: {input: 5.0, output: 25.0, cache_read: 0.50, cache_write: 10.0} + claude-opus-4-6: {input: 5.0, output: 25.0, cache_read: 0.50, cache_write: 10.0} + claude-opus-4-5: {input: 5.0, output: 25.0, cache_read: 0.50, cache_write: 10.0} + claude-opus-4-1: {input: 15.0, output: 75.0, cache_read: 1.50, cache_write: 30.0} + claude-opus-4: {input: 15.0, output: 75.0, cache_read: 1.50, cache_write: 30.0} + claude-sonnet-5-5: {input: 2.0, output: 10.0, cache_read: 0.10, cache_write: 4.0} + claude-sonnet-5: {input: 2.0, output: 10.0, cache_read: 0.20, cache_write: 4.0} + claude-sonnet-4-6: {input: 3.0, output: 15.0, cache_read: 0.30, cache_write: 6.0} + claude-sonnet-4-5: {input: 3.0, output: 15.0, cache_read: 0.30, cache_write: 6.0} + claude-sonnet-4: {input: 3.0, output: 15.0, cache_read: 0.30, cache_write: 6.0} + # Haiku 5.5: the over-100k-token prompt prices. + claude-haiku-5-5: {input: 0.50, output: 2.50, cache_read: 0.05, cache_write: 1.0} + claude-haiku-4-5: {input: 1.0, output: 5.0, cache_read: 0.10, cache_write: 2.0} + claude-haiku-3-5: {input: 0.80, output: 4.0, cache_read: 0.08, cache_write: 1.60} + openai: + prefixes: [gpt-, o1, o3, o4, codex-] + models: + gpt-5.5: {input: 5.0, output: 30.0, cache_read: 0.50, cache_write: 5.0} + # No cached-input price is published for gpt-5.5-pro. + gpt-5.5-pro: {input: 30.0, output: 180.0, cache_write: 30.0} + gpt-5.4: {input: 2.50, output: 15.0, cache_read: 0.25, cache_write: 2.50} + gpt-5.4-mini: {input: 0.75, output: 4.50, cache_read: 0.075, cache_write: 0.75} + gpt-5.4-nano: {input: 0.20, output: 1.25, cache_read: 0.02, cache_write: 0.20} + gpt-5.3-codex: {input: 1.75, output: 14.0, cache_read: 0.175, cache_write: 1.75} + gpt-5: {input: 1.25, output: 10.0, cache_read: 0.125, cache_write: 1.25} + gpt-5-mini: {input: 0.25, output: 2.0, cache_read: 0.025, cache_write: 0.25} + gpt-5-nano: {input: 0.05, output: 0.40, cache_read: 0.005, cache_write: 0.05} + o3: {input: 2.0, output: 8.0, cache_read: 0.50, cache_write: 2.0} + o4-mini: {input: 1.10, output: 4.40, cache_read: 0.275, cache_write: 1.10} + gpt-4.1: {input: 2.0, output: 8.0, cache_read: 0.50, cache_write: 2.0} + google: + prefixes: [gemini-] + models: + gemini-3.8-flash: {input: 1.50, output: 7.50, cache_read: 0.15, cache_write: 1.50} + gemini-3.7-flash: {input: 1.50, output: 7.50, cache_read: 0.15, cache_write: 1.50} + gemini-3.6-flash: {input: 1.50, output: 7.50, cache_read: 0.15, cache_write: 1.50} + gemini-3.5-flash: {input: 1.50, output: 9.0, cache_read: 0.15, cache_write: 1.50} + gemini-3.5-flash-lite: {input: 0.30, output: 2.50, cache_read: 0.03, cache_write: 0.30} + gemini-3.1-flash-lite: {input: 0.25, output: 1.50, cache_read: 0.025, cache_write: 0.25} + gemini-3.1-pro-preview: {input: 4.0, output: 18.0, cache_read: 0.40, cache_write: 4.0} + gemini-2.5-pro: {input: 2.50, output: 15.0, cache_read: 0.25, cache_write: 2.50} + gemini-2.5-flash: {input: 0.30, output: 2.50, cache_read: 0.03, cache_write: 0.30} + gemini-2.5-flash-lite: {input: 0.10, output: 0.40, cache_read: 0.01, cache_write: 0.10} diff --git a/src/aegis_core/budget/pricing.py b/src/aegis_core/budget/pricing.py new file mode 100644 index 0000000..8f64343 --- /dev/null +++ b/src/aegis_core/budget/pricing.py @@ -0,0 +1,171 @@ +"""Turning token counts into estimated dollars (design v1.0, §5). + +Usage is stored as tokens per billing category under the model id the log +recorded; dollars are computed here, when read, from the built-in table +(``prices.yaml``) and the signed overrides in ``budget.yaml``. Adding an +override therefore re-prices earlier usage without re-reading any log. + +A price is **explicit** when an override or a table entry gives it. Any +other price is a fallback and only an estimate: per category, the highest +rate among the model's vendor's entries, or across the whole table when the +vendor cannot be identified. A new model can cost more than every model the +table knows, so the fallback is not an upper bound. +""" + +from __future__ import annotations + +import re +from collections.abc import Mapping +from dataclasses import dataclass, field +from importlib import resources +from typing import Any + +import yaml + +CATEGORIES = ("input", "output", "cache_read", "cache_write") +PER_TOKENS = 1_000_000 + +# Log provider fields that name one vendor. Others (github-copilot, openrouter, +# ...) route to many vendors, so the model id's prefix decides instead. +_PROVIDER_VENDOR = { + "anthropic": "anthropic", + "openai": "openai", + "google": "google", + "gemini": "google", + "google-vertex": "google", +} +_DATE_SUFFIX = re.compile(r"-(?:\d{8}|\d{4}-\d{2}-\d{2})$") +_BRACKET_SUFFIX = re.compile(r"\[[^\]]*\]$") + + +def normalize_model(model: str) -> str: + """The key a model id is matched on: lower case, without a provider + prefix (``anthropic/``), a context-size tag (``[1m]``) or a date suffix, + with dots read as dashes (``claude-sonnet-4.5`` is ``claude-sonnet-4-5``).""" + m = model.strip().lower() + m = m.rsplit("/", 1)[-1] + m = _BRACKET_SUFFIX.sub("", m) + m = _DATE_SUFFIX.sub("", m) + return m.replace(".", "-") + + +@dataclass(frozen=True) +class Price: + """What one model costs, per category (USD per million tokens). A + category in ``estimated`` came from a fallback, not a published price.""" + + model: str + vendor: str | None + rates: Mapping[str, float] + estimated: frozenset[str] = frozenset() + + def cost(self, usage: Mapping[str, int]) -> float: + return sum(usage.get(c, 0) * self.rates[c] for c in CATEGORIES) / PER_TOKENS + + +@dataclass +class Cost: + """Dollars for a set of usage, and which models were (partly) estimated.""" + + dollars: float = 0.0 + estimated: dict[str, list[str]] = field(default_factory=dict) # model -> categories + + def add(self, other: Cost) -> None: + self.dollars += other.dollars + for model, cats in other.estimated.items(): + merged = sorted(set(self.estimated.get(model, [])) | set(cats)) + self.estimated[model] = merged + + +def _check_rates(where: str, entry: Any, *, require_all: bool) -> dict[str, float]: + if not isinstance(entry, dict): + raise ValueError(f"{where}: must be a mapping of {', '.join(CATEGORIES)}") + unknown = set(entry) - set(CATEGORIES) + if unknown: + raise ValueError(f"{where}: unknown price field(s) {sorted(unknown)}") + missing = [c for c in CATEGORIES if c not in entry] + if require_all and missing: + raise ValueError(f"{where}: missing price field(s) {missing}") + rates: dict[str, float] = {} + for cat, value in entry.items(): + if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: + raise ValueError(f"{where}.{cat}: must be a number >= 0 (USD per million tokens)") + rates[cat] = float(value) + return rates + + +class PriceTable: + """The built-in table plus ``budget.yaml`` overrides.""" + + def __init__(self, table: Mapping[str, Any], overrides: Mapping[str, Mapping[str, float]] + | None = None): + self.as_of = str(table.get("as_of", "")) + self._vendor_prefixes: dict[str, tuple[str, ...]] = {} + self._models: dict[str, tuple[str, dict[str, float]]] = {} # key -> (vendor, rates) + for vendor, spec in (table.get("vendors") or {}).items(): + self._vendor_prefixes[vendor] = tuple(spec.get("prefixes") or ()) + for model, entry in (spec.get("models") or {}).items(): + rates = _check_rates(f"prices.yaml: {model}", entry, require_all=False) + self._models[normalize_model(model)] = (vendor, rates) + self._overrides = {normalize_model(m): dict(r) for m, r in (overrides or {}).items()} + self._vendor_max: dict[str, dict[str, float]] = {} + self._global_max: dict[str, float] = {} + for vendor, rates in self._models.values(): + vmax = self._vendor_max.setdefault(vendor, {}) + for cat, rate in rates.items(): + vmax[cat] = max(vmax.get(cat, 0.0), rate) + self._global_max[cat] = max(self._global_max.get(cat, 0.0), rate) + + def vendor_of(self, model: str, provider: str | None = None) -> str | None: + """The vendor, from the log's provider field when it names one, else + from the table entry or the model id's prefix; ``None`` if unknown.""" + if provider and provider.lower() in _PROVIDER_VENDOR: + return _PROVIDER_VENDOR[provider.lower()] + key = normalize_model(model) + if key in self._models: + return self._models[key][0] + for vendor, prefixes in self._vendor_prefixes.items(): + if any(key.startswith(p.replace(".", "-")) for p in prefixes): + return vendor + return None + + def price(self, model: str, provider: str | None = None) -> Price: + key = normalize_model(model) + vendor = self.vendor_of(model, provider) + if key in self._overrides: + return Price(model, vendor, self._overrides[key]) + known = self._models[key][1] if key in self._models else {} + fallback = self._vendor_max.get(vendor or "", {}) or self._global_max + rates: dict[str, float] = {} + estimated = set() + for cat in CATEGORIES: + if cat in known: + rates[cat] = known[cat] + else: + rates[cat] = fallback.get(cat, self._global_max.get(cat, 0.0)) + estimated.add(cat) + return Price(model, vendor, rates, frozenset(estimated)) + + def cost(self, usage_by_model: Mapping[str, Mapping[str, int]], + providers: Mapping[str, str] | None = None) -> Cost: + """Dollars for ``{model: {category: tokens}}``. A model is reported as + estimated only for categories it actually used.""" + total = Cost() + for model, usage in usage_by_model.items(): + price = self.price(model, (providers or {}).get(model)) + used_estimates = sorted(c for c in price.estimated if usage.get(c, 0)) + total.add(Cost(price.cost(usage), {model: used_estimates} if used_estimates else {})) + return total + + +def load_builtin_table() -> dict[str, Any]: + text = resources.files("aegis_core.budget").joinpath("prices.yaml").read_text() + table = yaml.safe_load(text) + if not isinstance(table, dict) or table.get("version") != 1: + raise ValueError("aegis_core/budget/prices.yaml: unsupported shape or version") + return table + + +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/tests/fixtures/budget/claude/session.jsonl b/tests/fixtures/budget/claude/session.jsonl new file mode 100644 index 0000000..d102641 --- /dev/null +++ b/tests/fixtures/budget/claude/session.jsonl @@ -0,0 +1,18 @@ +{"type": "user", "timestamp": "2026-09-21T02:28:37.709Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0001", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:39.753Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0003", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0002", "message": {"model": "claude-opus-5", "id": "msg-0004", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 30871, "cache_read_input_tokens": 36463, "output_tokens": 318, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:40.181Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0003", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0005", "message": {"model": "claude-opus-5", "id": "msg-0004", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 30871, "cache_read_input_tokens": 36463, "output_tokens": 318, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:42.565Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0003", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0006", "message": {"model": "claude-opus-5", "id": "msg-0004", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 30871, "cache_read_input_tokens": 36463, "output_tokens": 318, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-21T02:28:44.192Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0007", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:46.860Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0009", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0008", "message": {"model": "claude-opus-5", "id": "msg-0010", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 1818, "cache_read_input_tokens": 67334, "output_tokens": 148, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-21T02:28:46.916Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0011", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:51.023Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0013", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0012", "message": {"model": "claude-opus-5", "id": "msg-0014", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 9510, "cache_read_input_tokens": 69152, "output_tokens": 541, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:51.029Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0013", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0015", "message": {"model": "claude-opus-5", "id": "msg-0014", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 9510, "cache_read_input_tokens": 69152, "output_tokens": 541, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:54.015Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0013", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0016", "message": {"model": "claude-opus-5", "id": "msg-0014", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 9510, "cache_read_input_tokens": 69152, "output_tokens": 541, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-21T02:28:55.169Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0017", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:28:57.851Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0019", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0018", "message": {"model": "claude-opus-5", "id": "msg-0020", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 891, "cache_read_input_tokens": 78662, "output_tokens": 199, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-21T02:28:57.899Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0021", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:29:07.634Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0023", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0022", "message": {"model": "claude-opus-5", "id": "msg-0024", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 3266, "cache_read_input_tokens": 79553, "output_tokens": 855, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:29:07.636Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0023", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0025", "message": {"model": "claude-opus-5", "id": "msg-0024", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 3266, "cache_read_input_tokens": 79553, "output_tokens": 855, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-21T02:29:08.750Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0023", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0026", "message": {"model": "claude-opus-5", "id": "msg-0024", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 3266, "cache_read_input_tokens": 79553, "output_tokens": 855, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-21T02:29:14.488Z", "isSidechain": false, "version": "2.1.275", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0027", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-21T02:29:39.620Z", "isSidechain": false, "version": "2.1.275", "requestId": "req-0029", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}", "uuid": "u-0028", "message": {"model": "claude-opus-5", "id": "msg-0030", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 894, "cache_read_input_tokens": 82819, "output_tokens": 2290, "service_tier": "standard"}}} diff --git a/tests/fixtures/budget/claude/subagent.jsonl b/tests/fixtures/budget/claude/subagent.jsonl new file mode 100644 index 0000000..a9c20e6 --- /dev/null +++ b/tests/fixtures/budget/claude/subagent.jsonl @@ -0,0 +1,7 @@ +{"type": "user", "timestamp": "2026-09-27T14:39:55.110Z", "isSidechain": true, "version": "2.1.281", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0031", "agentId": "a0000000000000001", "message": {"role": "user", "content": ""}} +{"type": "user", "timestamp": "2026-09-27T14:39:55.311Z", "isSidechain": true, "version": "2.1.281", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0032", "agentId": "a0000000000000001", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-27T14:39:57.722Z", "isSidechain": true, "version": "2.1.281", "requestId": "req-0034", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0033", "agentId": "a0000000000000001", "message": {"model": "claude-sonnet-5", "id": "msg-0035", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 60039, "cache_read_input_tokens": 0, "output_tokens": 4, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-27T14:39:59.099Z", "isSidechain": true, "version": "2.1.281", "requestId": "req-0034", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0036", "agentId": "a0000000000000001", "message": {"model": "claude-sonnet-5", "id": "msg-0035", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 60039, "cache_read_input_tokens": 0, "output_tokens": 4, "service_tier": "standard"}}} +{"type": "user", "timestamp": "2026-09-27T14:40:00.197Z", "isSidechain": true, "version": "2.1.281", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0037", "agentId": "a0000000000000001", "message": {"role": "user", "content": ""}} +{"type": "assistant", "timestamp": "2026-09-27T14:40:03.331Z", "isSidechain": true, "version": "2.1.281", "requestId": "req-0039", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0038", "agentId": "a0000000000000001", "message": {"model": "claude-sonnet-5", "id": "msg-0040", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 5574, "cache_read_input_tokens": 60039, "output_tokens": 2, "service_tier": "standard"}}} +{"type": "assistant", "timestamp": "2026-09-27T14:40:05.450Z", "isSidechain": true, "version": "2.1.281", "requestId": "req-0039", "sessionId": "11111111-2222-3333-4444-555555555555", "cwd": "{PROJECT}/sub", "uuid": "u-0041", "agentId": "a0000000000000001", "message": {"model": "claude-sonnet-5", "id": "msg-0040", "type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}], "usage": {"input_tokens": 2, "cache_creation_input_tokens": 5574, "cache_read_input_tokens": 60039, "output_tokens": 2, "service_tier": "standard"}}} diff --git a/tests/fixtures/budget/codex/child.jsonl b/tests/fixtures/budget/codex/child.jsonl new file mode 100644 index 0000000..9c92958 --- /dev/null +++ b/tests/fixtures/budget/codex/child.jsonl @@ -0,0 +1,10 @@ +{"timestamp": "2026-09-19T13:56:16.065Z", "type": "session_meta", "payload": {"timestamp": "2026-09-19T13:56:16.004Z", "cli_version": "0.154.0-alpha.6.2", "source": {"subagent": {"thread_spawn": {"parent_thread_id": "01a0b657-8357-7262-ad42-8e7b9f05772f", "depth": 1, "agent_path": "/root/update_documentation", "agent_nickname": "Ohm", "agent_role": null}}}, "model_provider": "openai", "thread_source": "subagent", "id": "01a00000-0000-7000-8000-000000000002", "session_id": "01a00000-0000-7000-8000-000000000001", "cwd": "{PROJECT}", "parent_thread_id": "01a00000-0000-7000-8000-000000000001", "agent_path": "/root/worker"}} +{"timestamp": "2026-09-19T13:56:17.305Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-09-19T13:56:17.305Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-09-19T13:56:17.305Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-09-19T13:56:17.305Z", "type": "response_item", "payload": {"type": "message", "role": "user", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-09-19T13:56:17.307Z", "type": "turn_context", "payload": {"model": "gpt-5.6-luna", "effort": "medium", "cwd": "{PROJECT}"}} +{"timestamp": "2026-09-19T13:56:20.269Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-09-19T13:56:23.069Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 32450, "cached_input_tokens": 14080, "cache_write_input_tokens": 0, "output_tokens": 247, "reasoning_output_tokens": 36, "total_tokens": 32697}, "last_token_usage": {"input_tokens": 32450, "cached_input_tokens": 14080, "cache_write_input_tokens": 0, "output_tokens": 247, "reasoning_output_tokens": 36, "total_tokens": 32697}, "model_context_window": 258400}}} +{"timestamp": "2026-09-19T13:56:27.099Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 65223, "cached_input_tokens": 45568, "cache_write_input_tokens": 0, "output_tokens": 413, "reasoning_output_tokens": 36, "total_tokens": 65636}, "last_token_usage": {"input_tokens": 32773, "cached_input_tokens": 31488, "cache_write_input_tokens": 0, "output_tokens": 166, "reasoning_output_tokens": 0, "total_tokens": 32939}, "model_context_window": 258400}}} +{"timestamp": "2026-09-19T13:56:28.849Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 65223, "cached_input_tokens": 45568, "cache_write_input_tokens": 0, "output_tokens": 413, "reasoning_output_tokens": 36, "total_tokens": 65636}, "last_token_usage": {"input_tokens": 32773, "cached_input_tokens": 31488, "cache_write_input_tokens": 0, "output_tokens": 166, "reasoning_output_tokens": 0, "total_tokens": 32939}, "model_context_window": 258400}}} diff --git a/tests/fixtures/budget/codex/parent.jsonl b/tests/fixtures/budget/codex/parent.jsonl new file mode 100644 index 0000000..1a58352 --- /dev/null +++ b/tests/fixtures/budget/codex/parent.jsonl @@ -0,0 +1,20 @@ +{"timestamp": "2026-10-01T00:29:53.107Z", "type": "session_meta", "payload": {"timestamp": "2026-10-01T00:29:49.834Z", "cli_version": "0.159.2", "source": "vscode", "model_provider": "openai", "thread_source": "user", "id": "01a00000-0000-7000-8000-000000000001", "session_id": "01a00000-0000-7000-8000-000000000001", "cwd": "{PROJECT}"}} +{"timestamp": "2026-10-01T00:29:53.409Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:29:53.409Z", "type": "response_item", "payload": {"type": "message", "role": "user", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:29:53.411Z", "type": "turn_context", "payload": {"model": "gpt-5.5", "effort": "medium", "cwd": "{PROJECT}"}} +{"timestamp": "2026-10-01T00:29:53.475Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:29:53.476Z", "type": "response_item", "payload": {"type": "message", "role": "user", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:29:53.477Z", "type": "response_item", "payload": {"type": "message", "role": "developer", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:29:53.479Z", "type": "response_item", "payload": {"type": "message", "role": "user", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:01.153Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:01.302Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 28856, "cached_input_tokens": 1408, "cache_write_input_tokens": 0, "output_tokens": 240, "reasoning_output_tokens": 62, "total_tokens": 29096}, "last_token_usage": {"input_tokens": 28856, "cached_input_tokens": 1408, "cache_write_input_tokens": 0, "output_tokens": 240, "reasoning_output_tokens": 62, "total_tokens": 29096}, "model_context_window": 258400}}} +{"timestamp": "2026-10-01T00:30:16.487Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:17.711Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 58425, "cached_input_tokens": 29440, "cache_write_input_tokens": 0, "output_tokens": 620, "reasoning_output_tokens": 203, "total_tokens": 59045}, "last_token_usage": {"input_tokens": 29569, "cached_input_tokens": 28032, "cache_write_input_tokens": 0, "output_tokens": 380, "reasoning_output_tokens": 141, "total_tokens": 29949}, "model_context_window": 258400}}} +{"timestamp": "2026-10-01T00:30:27.306Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:27.469Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 89923, "cached_input_tokens": 58496, "cache_write_input_tokens": 0, "output_tokens": 950, "reasoning_output_tokens": 216, "total_tokens": 90873}, "last_token_usage": {"input_tokens": 31498, "cached_input_tokens": 29056, "cache_write_input_tokens": 0, "output_tokens": 330, "reasoning_output_tokens": 13, "total_tokens": 31828}, "model_context_window": 258400}}} +{"timestamp": "2026-10-01T00:30:38.179Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:38.724Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 129097, "cached_input_tokens": 89600, "cache_write_input_tokens": 0, "output_tokens": 1292, "reasoning_output_tokens": 228, "total_tokens": 130389}, "last_token_usage": {"input_tokens": 39174, "cached_input_tokens": 31104, "cache_write_input_tokens": 0, "output_tokens": 342, "reasoning_output_tokens": 12, "total_tokens": 39516}, "model_context_window": 258400}}} +{"timestamp": "2026-10-01T00:30:48.038Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:30:48.275Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 176648, "cached_input_tokens": 127872, "cache_write_input_tokens": 0, "output_tokens": 1674, "reasoning_output_tokens": 228, "total_tokens": 178322}, "last_token_usage": {"input_tokens": 47551, "cached_input_tokens": 38272, "cache_write_input_tokens": 0, "output_tokens": 382, "reasoning_output_tokens": 0, "total_tokens": 47933}, "model_context_window": 258400}}} +{"timestamp": "2026-10-01T00:31:22.354Z", "type": "response_item", "payload": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": ""}]}} +{"timestamp": "2026-10-01T00:31:22.398Z", "type": "event_msg", "payload": {"type": "token_count", "info": {"total_token_usage": {"input_tokens": 230221, "cached_input_tokens": 175360, "cache_write_input_tokens": 0, "output_tokens": 3352, "reasoning_output_tokens": 744, "total_tokens": 233573}, "last_token_usage": {"input_tokens": 53573, "cached_input_tokens": 47488, "cache_write_input_tokens": 0, "output_tokens": 1678, "reasoning_output_tokens": 516, "total_tokens": 55251}, "model_context_window": 258400}}} diff --git a/tests/fixtures/budget/copilot/events.jsonl b/tests/fixtures/budget/copilot/events.jsonl new file mode 100644 index 0000000..3f408f3 --- /dev/null +++ b/tests/fixtures/budget/copilot/events.jsonl @@ -0,0 +1,17 @@ +{"type": "session.start", "timestamp": "2026-09-07T02:56:14.445Z", "data": {"version": 1, "producer": "copilot-agent", "copilotVersion": "1.0.80", "startTime": "2026-09-07T02:56:14.434Z", "selectedModel": "gpt-5.3-codex", "sessionId": "cccccccc-0000-0000-0000-000000000001", "context": {"cwd": "{PROJECT}"}}} +{"type": "assistant.message", "timestamp": "2026-09-07T02:56:26.136Z", "data": {"model": "gpt-5.3-codex", "content": ""}} +{"type": "assistant.message", "timestamp": "2026-09-07T02:56:29.084Z", "data": {"model": "gpt-5.3-codex", "content": ""}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T02:58:00.123Z", "data": {"totalPremiumRequests": 0, "totalNanoAiu": 35156800000}} +{"type": "session.shutdown", "timestamp": "2026-09-07T03:08:00.186Z", "data": {"totalPremiumRequests": 0, "totalNanoAiu": 35156800000, "tokenDetails": {"input": {"tokenCount": 99776}, "cache_read": {"tokenCount": 747520}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 3296}}, "shutdownType": "routine"}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T15:01:06.752Z", "data": {"totalPremiumRequests": 2, "totalNanoAiu": 97203085000}} +{"type": "session.shutdown", "timestamp": "2026-09-07T15:11:06.820Z", "data": {"totalPremiumRequests": 2, "totalNanoAiu": 97203085000, "tokenDetails": {"input": {"tokenCount": 218595}, "cache_read": {"tokenCount": 2094592}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 15924}}, "shutdownType": "routine"}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T15:17:58.038Z", "data": {"totalPremiumRequests": 3, "totalNanoAiu": 115026800000}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T15:20:27.821Z", "data": {"totalPremiumRequests": 4, "totalNanoAiu": 147106085000}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T15:23:21.977Z", "data": {"totalPremiumRequests": 5, "totalNanoAiu": 180271875000}} +{"type": "session.usage_checkpoint", "timestamp": "2026-09-07T15:23:57.884Z", "data": {"totalPremiumRequests": 6, "totalNanoAiu": 190150765000}} +{"type": "session.shutdown", "timestamp": "2026-09-07T15:50:06.176Z", "data": {"totalPremiumRequests": 17, "totalNanoAiu": 348510750000, "tokenDetails": {"input": {"tokenCount": 418314}, "cache_read": {"tokenCount": 12241920}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 43623}}, "shutdownType": "routine"}} +{"type": "session.shutdown", "timestamp": "2026-09-07T16:07:15.645Z", "data": {"totalPremiumRequests": 19, "totalNanoAiu": 395409000000, "tokenDetails": {"input": {"tokenCount": 658488}, "cache_read": {"tokenCount": 12357120}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 45660}}, "shutdownType": "routine"}} +{"type": "session.shutdown", "timestamp": "2026-09-08T01:40:48.266Z", "data": {"totalPremiumRequests": 28, "totalNanoAiu": 564707535000, "tokenDetails": {"input": {"tokenCount": 1235953}, "cache_read": {"tokenCount": 15636992}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 53406}}, "shutdownType": "routine"}} +{"type": "session.shutdown", "timestamp": "2026-09-08T20:38:52.312Z", "data": {"totalPremiumRequests": 36, "totalNanoAiu": 901701855000, "tokenDetails": {"input": {"tokenCount": 1458314}, "cache_read": {"tokenCount": 23972176}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 66893}}, "shutdownType": "routine"}} +{"type": "session.shutdown", "timestamp": "2026-09-08T21:56:02.364Z", "data": {"totalPremiumRequests": 40, "totalNanoAiu": 1032399045000, "tokenDetails": {"input": {"tokenCount": 1683291}, "cache_read": {"tokenCount": 25979579}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 68881}}, "shutdownType": "routine"}} +{"type": "session.shutdown", "timestamp": "2026-09-08T23:23:41.271Z", "data": {"totalPremiumRequests": 44, "totalNanoAiu": 1295207325000, "tokenDetails": {"input": {"tokenCount": 1922163}, "cache_read": {"tokenCount": 32009885}, "cache_write": {"tokenCount": 0}, "output": {"tokenCount": 75706}}, "shutdownType": "routine"}} diff --git a/tests/fixtures/budget/gemini/session.jsonl b/tests/fixtures/budget/gemini/session.jsonl new file mode 100644 index 0000000..6a10f41 --- /dev/null +++ b/tests/fixtures/budget/gemini/session.jsonl @@ -0,0 +1,26 @@ +{"startTime": "2026-09-27T20:58:34.394Z", "lastUpdated": "2026-09-27T20:58:34.394Z", "kind": "main", "sessionId": "77777777-aaaa-bbbb-cccc-000000000001", "projectHash": "0000000000000000000000000000000000000000000000000000000000000000"} +{"$set": {"lastUpdated": "2026-09-27T20:58:34.394Z", "messages": [{"id": "d04923d38bb0f6017037e74183378ef4", "timestamp": "2026-09-27T20:58:34.394Z", "type": "user", "content": [{"text": ""}]}]}} +{"timestamp": "2026-09-27T20:58:41.042Z", "type": "user", "id": "gu-0042", "content": [{"text": ""}]} +{"$set": {"lastUpdated": "2026-09-27T20:58:41.043Z"}} +{"timestamp": "2026-09-27T20:58:51.974Z", "type": "gemini", "tokens": {"input": 11391, "output": 115, "cached": 0, "thoughts": 887, "tool": 0, "total": 12393}, "model": "gemini-3.8-flash", "id": "g-0043", "content": ""} +{"$set": {"lastUpdated": "2026-09-27T20:58:51.974Z"}} +{"timestamp": "2026-09-27T20:58:51.974Z", "type": "gemini", "tokens": {"input": 11391, "output": 115, "cached": 0, "thoughts": 887, "tool": 0, "total": 12393}, "model": "gemini-3.8-flash", "id": "g-0043", "content": ""} +{"timestamp": "2026-09-27T20:58:52.640Z", "type": "user", "id": "gu-0044", "content": [{"text": ""}]} +{"$set": {"lastUpdated": "2026-09-27T20:58:52.640Z"}} +{"timestamp": "2026-09-27T20:58:53.737Z", "type": "gemini", "tokens": {"input": 12698, "output": 44, "cached": 0, "thoughts": 85, "tool": 0, "total": 12827}, "model": "gemini-3.8-flash", "id": "g-0045", "content": ""} +{"$set": {"lastUpdated": "2026-09-27T20:58:53.737Z"}} +{"timestamp": "2026-09-27T20:58:53.737Z", "type": "gemini", "tokens": {"input": 12698, "output": 44, "cached": 0, "thoughts": 85, "tool": 0, "total": 12827}, "model": "gemini-3.8-flash", "id": "g-0045", "content": ""} +{"timestamp": "2026-09-27T20:58:53.943Z", "type": "user", "id": "gu-0046", "content": [{"text": ""}]} +{"$set": {"lastUpdated": "2026-09-27T20:58:53.943Z"}} +{"timestamp": "2026-09-27T20:58:55.529Z", "type": "gemini", "tokens": {"input": 12888, "output": 82, "cached": 0, "thoughts": 172, "tool": 0, "total": 13142}, "model": "gemini-3.8-flash", "id": "g-0047", "content": ""} +{"$set": {"lastUpdated": "2026-09-27T20:58:55.529Z"}} +{"timestamp": "2026-09-27T20:58:55.529Z", "type": "gemini", "tokens": {"input": 12888, "output": 82, "cached": 0, "thoughts": 172, "tool": 0, "total": 13142}, "model": "gemini-3.8-flash", "id": "g-0047", "content": ""} +{"timestamp": "2026-09-27T20:58:55.631Z", "type": "user", "id": "gu-0048", "content": [{"text": ""}]} +{"$set": {"lastUpdated": "2026-09-27T20:58:55.631Z"}} +{"timestamp": "2026-09-27T20:58:57.427Z", "type": "gemini", "tokens": {"input": 13181, "output": 93, "cached": 0, "thoughts": 180, "tool": 0, "total": 13454}, "model": "gemini-3.8-flash", "id": "g-0049", "content": ""} +{"$set": {"lastUpdated": "2026-09-27T20:58:57.428Z"}} +{"timestamp": "2026-09-27T20:58:57.427Z", "type": "gemini", "tokens": {"input": 13181, "output": 93, "cached": 0, "thoughts": 180, "tool": 0, "total": 13454}, "model": "gemini-3.8-flash", "id": "g-0049", "content": ""} +{"timestamp": "2026-09-27T20:58:57.433Z", "type": "user", "id": "gu-0050", "content": [{"text": ""}]} +{"$set": {"lastUpdated": "2026-09-27T20:58:57.433Z"}} +{"timestamp": "2026-09-27T20:58:58.830Z", "type": "gemini", "tokens": {"input": 13547, "output": 114, "cached": 0, "thoughts": 129, "tool": 0, "total": 13790}, "model": "gemini-3.8-flash", "id": "g-0051", "content": ""} +{"$set": {"lastUpdated": "2026-09-27T20:58:58.830Z"}} diff --git a/tests/fixtures/budget/opencode/db.json b/tests/fixtures/budget/opencode/db.json new file mode 100644 index 0000000..b5e0af8 --- /dev/null +++ b/tests/fixtures/budget/opencode/db.json @@ -0,0 +1,385 @@ +{ + "schema": { + "message": "CREATE TABLE `message` (\n\t`id` text PRIMARY KEY,\n\t`session_id` text NOT NULL,\n\t`time_created` integer NOT NULL,\n\t`time_updated` integer NOT NULL,\n\t`data` text NOT NULL,\n\tCONSTRAINT `fk_message_session_id_session_id_fk` FOREIGN KEY (`session_id`) REFERENCES `session`(`id`) ON DELETE CASCADE\n)", + "session": "CREATE TABLE `session` (\n\t`id` text PRIMARY KEY,\n\t`project_id` text NOT NULL,\n\t`parent_id` text,\n\t`slug` text NOT NULL,\n\t`directory` text NOT NULL,\n\t`title` text NOT NULL,\n\t`version` text NOT NULL,\n\t`share_url` text,\n\t`summary_additions` integer,\n\t`summary_deletions` integer,\n\t`summary_files` integer,\n\t`summary_diffs` text,\n\t`revert` text,\n\t`permission` text,\n\t`time_created` integer NOT NULL,\n\t`time_updated` integer NOT NULL,\n\t`time_compacting` integer,\n\t`time_archived` integer, `workspace_id` text, `path` text, `agent` text, `model` text, `cost` real DEFAULT 0 NOT NULL, `tokens_input` integer DEFAULT 0 NOT NULL, `tokens_output` integer DEFAULT 0 NOT NULL, `tokens_reasoning` integer DEFAULT 0 NOT NULL, `tokens_cache_read` integer DEFAULT 0 NOT NULL, `tokens_cache_write` integer DEFAULT 0 NOT NULL, `metadata` text,\n\tCONSTRAINT `fk_session_project_id_project_id_fk` FOREIGN KEY (`project_id`) REFERENCES `project`(`id`) ON DELETE CASCADE\n)" + }, + "session": [ + { + "id": "ses-0052", + "project_id": "proj-1", + "parent_id": "ses-0053", + "slug": "s", + "directory": "{PROJECT}", + "title": "", + "version": "1.18.26", + "time_created": 1789305823015, + "time_updated": 1789305860777 + }, + { + "id": "ses-0054", + "project_id": "proj-1", + "parent_id": "ses-0053", + "slug": "s", + "directory": "{PROJECT}", + "title": "", + "version": "1.18.26", + "time_created": 1789305733593, + "time_updated": 1789305817370 + }, + { + "id": "ses-0055", + "project_id": "proj-1", + "parent_id": "ses-0053", + "slug": "s", + "directory": "{PROJECT}", + "title": "", + "version": "1.18.26", + "time_created": 1788796502399, + "time_updated": 1789305726468 + }, + { + "id": "ses-0056", + "project_id": "proj-1", + "parent_id": "ses-0053", + "slug": "s", + "directory": "{PROJECT}", + "title": "", + "version": "1.18.26", + "time_created": 1788796455180, + "time_updated": 1788796500341 + }, + { + "id": "ses-0053", + "project_id": "proj-1", + "parent_id": null, + "slug": "s", + "directory": "{PROJECT}", + "title": "", + "version": "1.18.26", + "time_created": 1788796401974, + "time_updated": 1789305900310 + } + ], + "message": [ + { + "id": "msg-0057", + "session_id": "ses-0053", + "time_created": 1788796402096, + "time_updated": 1788796437388, + "data": { + "role": "user", + "time": { + "created": 1788796402096 + }, + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0058", + "session_id": "ses-0053", + "time_created": 1788796402101, + "time_updated": 1788796417004, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 8202, + "input": 5900, + "output": 510, + "reasoning": 0, + "cache": { + "write": 0, + "read": 1792 + } + }, + "cost": 0, + "time": { + "created": 1788796402101, + "completed": 1788796417003 + }, + "finish": "tool-calls", + "mode": "build", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0059", + "session_id": "ses-0053", + "time_created": 1788796417005, + "time_updated": 1788796427783, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 8770, + "input": 66, + "output": 512, + "reasoning": 0, + "cache": { + "write": 0, + "read": 8192 + } + }, + "cost": 0, + "time": { + "created": 1788796417005, + "completed": 1788796427783 + }, + "finish": "tool-calls", + "mode": "build", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0060", + "session_id": "ses-0053", + "time_created": 1788796427784, + "time_updated": 1788796437383, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 9258, + "input": 109, + "output": 445, + "reasoning": 0, + "cache": { + "write": 0, + "read": 8704 + } + }, + "cost": 0, + "time": { + "created": 1788796427784, + "completed": 1788796437383 + }, + "finish": "tool-calls", + "mode": "build", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0061", + "session_id": "ses-0053", + "time_created": 1788796437384, + "time_updated": 1788796447807, + "data": { + "role": "user", + "time": { + "created": 1788796437384 + }, + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0062", + "session_id": "ses-0053", + "time_created": 1788796437385, + "time_updated": 1788796447803, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 2506, + "input": 1953, + "output": 553, + "reasoning": 0, + "cache": { + "write": 0, + "read": 0 + } + }, + "cost": 0, + "time": { + "created": 1788796437385, + "completed": 1788796447803 + }, + "finish": "stop", + "mode": "compaction", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0063", + "session_id": "ses-0053", + "time_created": 1788796447803, + "time_updated": 1789305860781, + "data": { + "role": "user", + "time": { + "created": 1788796447803 + }, + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0064", + "session_id": "ses-0053", + "time_created": 1788796447805, + "time_updated": 1788796500331, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 8287, + "input": 534, + "output": 329, + "reasoning": 0, + "cache": { + "write": 0, + "read": 7424 + } + }, + "cost": 0, + "time": { + "created": 1788796447805, + "completed": 1788796500331 + }, + "finish": "tool-calls", + "mode": "build", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0065", + "session_id": "ses-0056", + "time_created": 1788796455188, + "time_updated": 1788796500343, + "data": { + "role": "user", + "time": { + "created": 1788796455188 + }, + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0066", + "session_id": "ses-0056", + "time_created": 1788796455203, + "time_updated": 1788796467083, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 3696, + "input": 3114, + "output": 582, + "reasoning": 0, + "cache": { + "write": 0, + "read": 0 + } + }, + "cost": 0, + "time": { + "created": 1788796455203, + "completed": 1788796467083 + }, + "finish": "tool-calls", + "mode": "explore", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0067", + "session_id": "ses-0056", + "time_created": 1788796467085, + "time_updated": 1788796474925, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 4073, + "input": 2435, + "output": 358, + "reasoning": 0, + "cache": { + "write": 0, + "read": 1280 + } + }, + "cost": 0, + "time": { + "created": 1788796467085, + "completed": 1788796474925 + }, + "finish": "tool-calls", + "mode": "explore", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + }, + { + "id": "msg-0068", + "session_id": "ses-0056", + "time_created": 1788796474926, + "time_updated": 1788796480254, + "data": { + "role": "assistant", + "modelID": "Qwen3-1.7B-MLX-bf16", + "providerID": "omlx", + "tokens": { + "total": 4347, + "input": 252, + "output": 255, + "reasoning": 0, + "cache": { + "write": 0, + "read": 3840 + } + }, + "cost": 0, + "time": { + "created": 1788796474926, + "completed": 1788796480254 + }, + "finish": "tool-calls", + "mode": "explore", + "path": { + "cwd": "{PROJECT}", + "root": "{PROJECT}" + } + } + } + ] +} \ No newline at end of file diff --git a/tests/test_authority.py b/tests/test_authority.py index 4f0c0ab..b49ccbf 100644 --- a/tests/test_authority.py +++ b/tests/test_authority.py @@ -20,7 +20,7 @@ def test_load_authority_map_from_yaml(tmp_path): def test_example_authority_file_loads(): authority_map = load_authority_map("data/authority.example.yaml") - assert authority_map["admin"] == {"scaling", "deletion", "configuration", "identity"} + assert authority_map["admin"] == {"scaling", "deletion", "configuration", "identity", "budget"} assert authority_map["sre_lead"] == {"scaling", "configuration"} assert authority_map["developer"] == {"configuration"} @@ -72,7 +72,7 @@ def test_example_authority_file_verifies_under_example_key(): key = load_key("file:data/example-signing.key") authority_map = load_authority_map("data/authority.example.yaml", key=key) - assert authority_map["admin"] == {"scaling", "deletion", "configuration", "identity"} + assert authority_map["admin"] == {"scaling", "deletion", "configuration", "identity", "budget"} assert authority_map.warnings == [] diff --git a/tests/test_budget_accounting.py b/tests/test_budget_accounting.py new file mode 100644 index 0000000..d14b00a --- /dev/null +++ b/tests/test_budget_accounting.py @@ -0,0 +1,406 @@ +"""Session and project/day accounting (design v1.0 §4.1) and the validation +cases of §10 step 5: crash recovery, cache deletion with a high-revision +project entry, log replacement, a lost revision counter, rebuilds racing +hook updates, linked sub-agents, missing history and midnight rebuilds.""" + +import json +import os +import shutil +import threading +from datetime import UTC, datetime +from pathlib import Path +from zoneinfo import ZoneInfo + +import pytest +from test_budget_logs import CODEX_CHILD, CODEX_PARENT, home # noqa: F401 (fixture) + +from aegis_core.budget.accounting import ( + UNKNOWN_DAY, + BudgetStore, + ProjectDay, + find_project_root, + rebuild, + refresh, + update_session, +) +from aegis_core.budget.logs import ClaudeLog, CodexLog + +UTC_ZONE = ZoneInfo("UTC") +DAY = "2026-10-07" +NOW = datetime(2026, 10, 7, 18, 0, tzinfo=UTC).timestamp() +MODEL = "claude-opus-5-5" + + +class World: + """A project with ``.aegis/``, a fake ``~/.claude`` and a budget store.""" + + def __init__(self, tmp: Path): + self.tmp = tmp + self.project = tmp / "project" + (self.project / ".aegis").mkdir(parents=True) + self.root = str(self.project.resolve()) + self.claude = ClaudeLog(tmp / "home" / ".claude") + self.logs = {"claude": self.claude} + self.store = BudgetStore(cache_dir=tmp / "cache", state_dir=tmp / "state") + self._n = 0 + + def path(self, sid: str) -> Path: + return self.tmp / "home" / ".claude" / "projects" / "-project" / f"{sid}.jsonl" + + def start(self, sid: str, ts: str = f"{DAY}T09:00:00Z", cwd: Path | None = None, + path: Path | None = None) -> Path: + p = path or self.path(sid) + p.parent.mkdir(parents=True, exist_ok=True) + rec = {"type": "user", "sessionId": sid, "cwd": str(cwd or self.project)} + if ts: + rec["timestamp"] = ts + p.write_text(json.dumps(rec) + "\n") + return p + + def add(self, sid: str, tokens: int = 1000, ts: str | None = f"{DAY}T12:00:00Z", + path: Path | None = None) -> None: + self._n += 1 + rec = {"type": "assistant", "sessionId": sid, "requestId": f"r{self._n}", + "message": {"id": f"m{self._n}", "model": MODEL, + "usage": {"input_tokens": tokens, "output_tokens": 0}}} + if ts: + rec["timestamp"] = ts + with open(path or self.path(sid), "a") as f: + f.write(json.dumps(rec) + "\n") + + def loc(self, sid: str): + return self.claude.locate(sid) + + def refresh(self, sid: str, now: float = NOW, zone=UTC_ZONE): + return refresh(self.store, self.logs, self.loc(sid), zone, now=now) + + def project_input(self, day: str = DAY) -> float: + pd = self.store.load_day(self.root, day) + return sum(u.get("input", 0) for u in pd.usage().values()) if pd else 0 + + +def _input(by_model) -> float: + return sum(u.get("input", 0) for u in by_model.values()) + + +@pytest.fixture +def world(tmp_path): + return World(tmp_path) + + +# --- the basic protocol ------------------------------------------------------------------- + + +def test_a_session_is_counted_once_in_its_project(world): + world.start("s1") + world.add("s1", 100) + world.add("s1", 200) + snap = world.refresh("s1") + assert snap.record.project_root == world.root + assert _input(snap.session_usage) == 300 and _input(snap.project_usage) == 300 + rev = snap.record.revision + again = world.refresh("s1") + assert _input(again.project_usage) == 300 + assert again.record.revision == rev # nothing new: no new revision, no rewrite + + +def test_sessions_of_a_project_sum_and_other_projects_stay_out(world): + world.start("s1") + world.add("s1", 100) + world.start("s2") + world.add("s2", 50) + other = world.tmp / "other" + (other / ".aegis").mkdir(parents=True) + world.start("s3", cwd=other) + world.add("s3", 7) + world.refresh("s1") + world.refresh("s3") + snap = world.refresh("s2") + assert _input(snap.project_usage) == 150 + assert set(snap.project.entries) == {"claude/s1", "claude/s2"} + + +def test_a_session_belongs_to_the_directory_it_started_in(world): + nested = world.project / "sub" / "deeper" + nested.mkdir(parents=True) + assert find_project_root(str(nested)) == world.root + inner = world.project / "inner" + (inner / ".aegis").mkdir(parents=True) + assert find_project_root(str(inner / "x")) == str(inner.resolve()) # innermost wins + assert find_project_root(str(world.tmp)) is None + + +# --- crash recovery ------------------------------------------------------------------------- + + +def test_a_crash_between_the_session_write_and_the_project_merge_heals(world): + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + world.add("s1", 40) + with world.store.session_lock("claude", "s1"): # steps 1-3 only, then "crash" + update_session(world.store, world.claude, world.loc("s1"), UTC_ZONE) + assert world.project_input() == 100 # the map is one revision behind + snap = world.refresh("s1") + assert _input(snap.project_usage) == 140 # merged again, not double counted + assert world.refresh("s1").project.entries["claude/s1"]["revision"] == snap.record.revision + + +# --- the durable revision counter (review round 3) ------------------------------------------- + + +def test_deleting_only_the_session_cache_keeps_updates_flowing(world): + """The reviewer's case: the project map holds a high-revision entry, the + session cache is deleted, usage is appended -- the next hook updates the + project total immediately, and a stale rebuild cannot overwrite it.""" + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + # Simulate a long-lived session: its counter and project entry are at revision 100. + world.store.write_revision("claude", "s1", 100) + with world.store.project_lock(world.root): + pd = world.store.load_day(world.root, DAY) + pd.entries["claude/s1"]["revision"] = 100 + world.store.save_day(pd) + stale = dict(pd.entries["claude/s1"]) + + world.store.record_path("claude", "s1").unlink() # only the session cache + world.add("s1", 25) + snap = world.refresh("s1") + assert _input(snap.project_usage) == 125 + assert snap.project.entries["claude/s1"]["revision"] == 101 + + # A rebuild holding the older snapshot of s1 merges after the hook: it is ignored. + with world.store.project_lock(world.root): + pd = world.store.load_day(world.root, DAY) + assert not world.store.merge_entries(pd, {"claude/s1": stale}) + assert world.project_input() == 125 + + +def test_a_replaced_log_restarts_the_count_but_never_the_revision(world): + p = world.start("s1") + world.add("s1", 100) + world.add("s1", 100) + first = world.refresh("s1") + replacement = p.with_suffix(".new") + world.start("s1", path=replacement) + world.add("s1", 30, path=replacement) + os.replace(replacement, p) # another inode, other content + snap = world.refresh("s1") + assert _input(snap.session_usage) == 30 and _input(snap.project_usage) == 30 + assert snap.record.revision > first.record.revision + + +def test_a_lost_revision_counter_is_recovered_above_every_stored_revision(world): + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + with world.store.project_lock(world.root): + pd = world.store.load_day(world.root, DAY) + pd.entries["claude/s1"]["revision"] = 57 + world.store.save_day(pd) + world.store.rev_path("claude", "s1").unlink() # the state counter is gone + world.store.record_path("claude", "s1").unlink() # and so is the cache record + world.add("s1", 1) + snap = world.refresh("s1") + assert snap.record.revision == 59 # 57 + 1 recovered, + 1 written + assert _input(snap.project_usage) == 101 + assert world.store.read_revision("claude", "s1") == 59 + + +# --- concurrency ------------------------------------------------------------------------------- + + +def test_rebuilds_racing_hook_updates_neither_deadlock_nor_lose_usage(world): + sids = [f"s{i}" for i in range(4)] + for sid in sids: + world.start(sid) + errors: list[BaseException] = [] + stop = threading.Event() + + def hooks(sid): + try: + for _ in range(15): + world.add(sid, 10) + world.refresh(sid) + except BaseException as e: # pragma: no cover - reported below + errors.append(e) + + def rebuilds(): + try: + while not stop.is_set(): + world.store.day_path(world.root, DAY).unlink(missing_ok=True) + rebuild(world.store, world.logs, world.root, DAY, UTC_ZONE) + except BaseException as e: # pragma: no cover + errors.append(e) + + threads = [threading.Thread(target=hooks, args=(sid,)) for sid in sids] + rb = threading.Thread(target=rebuilds) + rb.start() + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + stop.set() + rb.join(timeout=60) + assert not any(t.is_alive() for t in [*threads, rb]), "deadlock" + assert not errors, errors + for sid in sids: + world.refresh(sid) + assert world.project_input() == 4 * 15 * 10 + + +# --- linked sub-agents --------------------------------------------------------------------------- + + +def test_a_linked_claude_subagent_is_counted_once_inside_its_parent(world): + p = world.start("s1") + world.add("s1", 100) + sub = p.with_suffix("") / "subagents" / "agent-a1.jsonl" + world.start("s1", path=sub, cwd=world.project / ".claude" / "worktrees" / "a1") + world.add("s1", 40, path=sub) + snap = world.refresh("s1") + assert _input(snap.session_usage) == 140 + assert set(snap.project.entries) == {"claude/s1"} + world.store.day_path(world.root, DAY).unlink() + pd = rebuild(world.store, world.logs, world.root, DAY, UTC_ZONE) + assert set(pd.entries) == {"claude/s1"} and _input(pd.usage()) == 140 + + +def test_a_codex_subagent_hook_is_accounted_in_its_parent(home, tmp_path): # noqa: F811 + logs = {"codex": CodexLog(home["home"] / ".codex")} + store = BudgetStore(cache_dir=tmp_path / "c", state_dir=tmp_path / "s") + child = logs["codex"].locate(CODEX_CHILD) + now = datetime(2026, 10, 1, 12, tzinfo=UTC).timestamp() + snap = refresh(store, logs, child, UTC_ZONE, now=now) + assert snap.record.session_id == CODEX_PARENT + assert set(snap.project.entries) == {f"codex/{CODEX_PARENT}"} + assert any(s.startswith("child:") for s in snap.record.sources) + + +# --- missing history --------------------------------------------------------------------------- + + +def _wipe_cache(world): + shutil.rmtree(world.store.cache, ignore_errors=True) + + +def test_logs_gone_with_the_inventory_present_is_partial_from_checkpoints(world): + world.start("s1") + world.add("s1", 100) + world.start("s2") + world.add("s2", 30) + world.refresh("s1") + world.refresh("s2") + world.path("s2").unlink() # s2's log is gone ... + _wipe_cache(world) # ... and so is the cache + snap = world.refresh("s1") + assert snap.project.completeness == "partial" + assert snap.project.missing_sessions == 1 and snap.project.known_sessions == 2 + assert _input(snap.project_usage) == 130 # s2 from its inventory checkpoint + + +def test_a_deleted_inventory_makes_completeness_unknown_never_zero_missing(world): + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + world.store.inventory_path(world.root, DAY).unlink() + _wipe_cache(world) + snap = world.refresh("s1") + assert snap.project.completeness == "unknown" + assert _input(snap.project_usage) == 100 + + +def test_first_use_with_other_sessions_today_is_unknown_and_alone_is_complete(world): + world.start("s1") + world.add("s1", 100) + assert world.refresh("s1").project.completeness == "complete" + other = World(world.tmp / "w2") + other.start("a") + other.add("a", 5) + other.start("b") + other.add("b", 6) + snap = other.refresh("b") + assert snap.project.completeness == "unknown" + assert _input(snap.project_usage) == 11 + + +def test_a_new_day_is_complete(world): + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + world.add("s1", 7, ts="2026-10-08T01:00:00Z") + snap = world.refresh("s1", now=datetime(2026, 10, 8, 2, tzinfo=UTC).timestamp()) + assert snap.day == "2026-10-08" and snap.project.completeness == "complete" + assert _input(snap.project_usage) == 7 and _input(snap.session_usage) == 107 + + +# --- days ---------------------------------------------------------------------------------- + + +def test_a_rebuild_after_midnight_keeps_old_usage_in_its_own_day(world): + world.start("s1", ts="2026-10-06T22:00:00Z") + world.add("s1", 100, ts="2026-10-06T23:30:00Z") + world.refresh("s1", now=datetime(2026, 10, 6, 23, 45, tzinfo=UTC).timestamp()) + _wipe_cache(world) + world.add("s1", 5, ts=None) # no timestamp: nearest earlier is 10-06 + world.add("s1", 9, ts="2026-10-07T00:10:00Z") + snap = world.refresh("s1", now=datetime(2026, 10, 7, 0, 20, tzinfo=UTC).timestamp()) + assert snap.day == "2026-10-07" + assert _input(snap.project_usage) == 9 + assert _input(snap.record.usage_on("2026-10-06")) == 105 + assert _input(snap.session_usage) == 114 + + +def test_the_day_follows_the_policy_time_zone(world): + world.start("s1") + world.add("s1", 100, ts="2026-10-07T03:00:00Z") # 22:00 on 10-06 in Chicago + chicago = ZoneInfo("America/Chicago") + snap = world.refresh("s1", now=datetime(2026, 10, 7, 4, tzinfo=UTC).timestamp(), + zone=chicago) + assert snap.day == "2026-10-06" and _input(snap.project_usage) == 100 + + +def test_usage_with_no_recorded_time_counts_for_the_session_only(world): + p = world.path("s1") + world.start("s1", ts=None) + world.add("s1", 60, ts=None, path=p) + snap = world.refresh("s1") + assert _input(snap.record.buckets()[UNKNOWN_DAY]) == 60 + assert _input(snap.session_usage) == 60 and _input(snap.project_usage) == 0 + + +# --- odds and ends ------------------------------------------------------------------------------ + + +def test_a_session_outside_any_project_has_a_session_total_only(world): + world.start("s1", cwd=world.tmp / "elsewhere") + world.add("s1", 10) + snap = world.refresh("s1") + assert snap.record.project_root is None and snap.project is None + assert _input(snap.session_usage) == 10 + + +def test_a_missing_log_keeps_the_last_totals(world): + world.start("s1") + world.add("s1", 10) + world.refresh("s1") + loc = world.loc("s1") + world.path("s1").unlink() + snap = refresh(world.store, world.logs, loc, UTC_ZONE, now=NOW) + assert snap.record.log_missing and _input(snap.session_usage) == 10 + + +def test_unsafe_session_ids_are_hashed_into_file_names(world): + path = world.store.record_path("claude", "../../etc/passwd") + assert path.parent == world.store.cache / "sessions" / "claude" + assert path.name.startswith("h-") + + +def test_merge_never_replaces_a_newer_entry(): + pd = ProjectDay("/p", DAY, {"a/1": {"revision": 5, "usage": {"m": {"input": 5}}}}) + assert not BudgetStore.merge_entries(pd, {"a/1": {"revision": 4, "usage": {}}}) + assert not BudgetStore.merge_entries(pd, {"a/1": {"revision": 5, "usage": {}}}) + assert BudgetStore.merge_entries(pd, {"a/1": {"revision": 6, "usage": {"m": {"input": 1}}}, + "b/2": {"revision": 1, "usage": {}}}) + assert pd.entries["a/1"]["revision"] == 6 and "b/2" in pd.entries diff --git a/tests/test_budget_evaluate.py b/tests/test_budget_evaluate.py new file mode 100644 index 0000000..aa4e6e0 --- /dev/null +++ b/tests/test_budget_evaluate.py @@ -0,0 +1,130 @@ +"""allow / warn / deny from a usage snapshot (design v1.0 §6, §11).""" + +import pytest + +from aegis_core.budget import price_table +from aegis_core.budget.accounting import UNKNOWN_DAY, ProjectDay, SessionRecord, Snapshot +from aegis_core.budget.evaluate import ALLOW, DENY, WARN, evaluate +from aegis_core.budget.policy import BudgetPolicy, Limit + +DAY = "2026-10-07" +OPUS = "claude-opus-5-5" # $4 in / $20 out per million + + +def policy(**kw) -> BudgetPolicy: + base = dict(path="budget.yaml", principal="admin", unit="usd", session=Limit(10.0), + project_day=Limit(50.0), tz="UTC", agents=("claude", "codex", "copilot"), + on_unknown_log="allow", on_unknown_price="estimate", pricing={}, + copilot_premium_requests=Limit(20)) + base.update(kw) + return BudgetPolicy(**base) + + +def snap(session: dict, project: dict | None = None, agent: str = "claude", *, + completeness: str = "complete", unknown_day: dict | None = None, **rec_kw) -> Snapshot: + buckets = {DAY: session} + if unknown_day: + buckets[UNKNOWN_DAY] = unknown_day + rec = SessionRecord(agent, "s1", {"agent": agent, "session_id": "s1", "path": "x"}, "UTC", + sources={"main": {"buckets": buckets}}, **rec_kw) + project = session if project is None else project + pd = ProjectDay("/p", DAY, {f"{agent}/s1": {"revision": 1, "usage": project}}, + completeness=completeness) + return Snapshot(rec, DAY, pd, rec.usage_total(), pd.usage(), {}) + + +def dollars(d: float) -> dict: + """Opus input tokens costing ``d`` dollars.""" + return {OPUS: {"input": d / 4 * 1_000_000}} + + +def test_under_the_warning_threshold_is_allowed_silently(): + p = policy() + v = evaluate(p, price_table(p), snap(dollars(5))) + assert v.decision == ALLOW and v.messages == [] and v.reasons == [] + assert [m.name for m in v.measures] == ["session", "project_day"] + assert v.measures[0].used == pytest.approx(5.0) + + +def test_crossing_warn_at_warns_with_what_is_left(): + p = policy() + v = evaluate(p, price_table(p), snap(dollars(8.5))) + assert v.decision == WARN and v.reasons == ["session-warn"] + assert "85%" in v.messages[0] and "$1.50 left" in v.messages[0] + + +def test_reaching_a_limit_denies_and_reports_overshoot(): + p = policy() + v = evaluate(p, price_table(p), snap(dollars(12))) + assert v.decision == DENY and "session-over" in v.reasons + assert v.measures[0].overshoot == pytest.approx(2.0) + assert "overshoot $2.00" in v.messages[0] + + +def test_the_project_day_limit_counts_every_session(): + p = policy() + v = evaluate(p, price_table(p), snap(dollars(1), project=dollars(60))) + assert v.decision == DENY and v.reasons == ["project_day-over"] + + +def test_token_mode_counts_all_four_categories(): + p = policy(unit="tokens", session=Limit(1000), project_day=None) + usage = {OPUS: {"input": 100, "output": 200, "cache_read": 300, "cache_write": 400}} + v = evaluate(p, price_table(p), snap(usage)) + assert v.measures[0].used == 1000 and v.decision == DENY + + +def test_unpriced_models_are_estimated_or_denied(): + usage = {"claude-opus-9": {"input": 1000}} + est = policy() + v = evaluate(est, price_table(est), snap(usage)) + assert v.decision == WARN and "estimated-price" in v.reasons + assert v.estimated == {"claude-opus-9": ["input"]} + strict = policy(on_unknown_price="deny") + v = evaluate(strict, price_table(strict), snap(usage)) + assert v.decision == DENY and "unpriced-model" in v.reasons + assert "add it under 'pricing:'" in v.messages[-1] + priced = policy(on_unknown_price="deny", pricing={"claude-opus-9": { + "input": 1.0, "output": 1.0, "cache_read": 0.1, "cache_write": 1.0}}) + assert evaluate(priced, price_table(priced), snap(usage)).decision == ALLOW + + +@pytest.mark.parametrize("kw,fragment", [ + ({"log_missing": True}, "the session's log is missing"), + ({"problem_total": 3}, "3 unreadable log record(s)"), + ({"unknown_day": {OPUS: {"input": 10}}}, "no recorded time"), + ({"completeness": "unknown"}, "today's project total is unknown"), + ({"completeness": "partial"}, "today's project total is partial"), +]) +def test_unverifiable_usage_follows_on_unknown_log(kw, fragment): + lenient = policy() + v = evaluate(lenient, price_table(lenient), snap(dollars(1), **kw)) + assert v.decision == WARN and "unknown-log-warn" in v.reasons + assert fragment in v.messages[-1] + strict = policy(on_unknown_log="deny") + v = evaluate(strict, price_table(strict), snap(dollars(1), **kw)) + assert v.decision == DENY and "unknown-log" in v.reasons + + +def test_copilot_is_measured_in_premium_requests(): + p = policy() + s = snap({"copilot-premium-requests": {"requests": 21}}, agent="copilot") + v = evaluate(p, price_table(p), s) + assert [m.name for m in v.measures] == ["copilot_premium_requests"] + assert v.decision == DENY and "21.00 premium requests" in v.messages[0] + + +def test_unlisted_or_unmeasured_agents_are_not_evaluated(): + p = policy(agents=("claude",)) + assert evaluate(p, price_table(p), snap(dollars(99), agent="codex")).decision == ALLOW + q = policy(copilot_premium_requests=None) + s = snap({"copilot-premium-requests": {"requests": 99}}, agent="copilot") + assert evaluate(q, price_table(q), s).decision == ALLOW + + +def test_a_session_outside_any_project_has_only_the_session_limit(): + p = policy() + s = snap(dollars(1)) + s.project = None + v = evaluate(p, price_table(p), s) + assert [m.name for m in v.measures] == ["session"] diff --git a/tests/test_budget_logs.py b/tests/test_budget_logs.py new file mode 100644 index 0000000..486884c --- /dev/null +++ b/tests/test_budget_logs.py @@ -0,0 +1,365 @@ +"""Reading usage from each agent's logs (design v1.0 §4), against redacted +samples of real logs in tests/fixtures/budget/.""" + +import json +import os +import shutil +import sqlite3 +from collections import defaultdict +from pathlib import Path + +import pytest + +from aegis_core.budget.logs import ( + ClaudeLog, + CodexLog, + CopilotLog, + GeminiLog, + LogMissing, + OpenCodeLog, + parse_ts, + tail_jsonl, +) + +FIX = Path(__file__).parent / "fixtures" / "budget" +CLAUDE_SID = "11111111-2222-3333-4444-555555555555" +CODEX_PARENT = "01a00000-0000-7000-8000-000000000001" +CODEX_CHILD = "01a00000-0000-7000-8000-000000000002" +GEMINI_SID = "77777777-aaaa-bbbb-cccc-000000000001" +COPILOT_SID = "cccccccc-0000-0000-0000-000000000001" + + +def _copy(src: Path, dest: Path, project: Path) -> Path: + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_text(src.read_text().replace("{PROJECT}", str(project))) + return dest + + +def _records(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + + +def build_opencode_db(db_path: Path, project: Path) -> dict: + spec = json.loads((FIX / "opencode" / "db.json").read_text().replace("{PROJECT}", + str(project))) + db_path.parent.mkdir(parents=True, exist_ok=True) + db = sqlite3.connect(db_path) + db.executescript("PRAGMA foreign_keys=OFF;") + for sql in spec["schema"].values(): + db.execute(sql) + for s in spec["session"]: + db.execute("INSERT INTO session (id, project_id, parent_id, slug, directory, title, " + "version, time_created, time_updated) VALUES (?,?,?,?,?,?,?,?,?)", + (s["id"], s["project_id"], s["parent_id"], s["slug"], s["directory"], + s["title"], s["version"], s["time_created"], s["time_updated"])) + for m in spec["message"]: + db.execute("INSERT INTO message (id, session_id, time_created, time_updated, data) " + "VALUES (?,?,?,?,?)", (m["id"], m["session_id"], m["time_created"], + m["time_updated"], json.dumps(m["data"]))) + db.commit() + db.close() + return spec + + +@pytest.fixture +def home(tmp_path): + project = tmp_path / "project" + (project / ".aegis").mkdir(parents=True) + h = tmp_path / "home" + claude_main = _copy(FIX / "claude" / "session.jsonl", + h / ".claude" / "projects" / "-project" / f"{CLAUDE_SID}.jsonl", project) + _copy(FIX / "claude" / "subagent.jsonl", claude_main.with_suffix("") / "subagents" / + "agent-a0000000000000001.jsonl", project) + day = h / ".codex" / "sessions" / "2026" / "10" / "01" + _copy(FIX / "codex" / "parent.jsonl", day / f"rollout-2026-10-01T00-29-49-{CODEX_PARENT}.jsonl", + project) + _copy(FIX / "codex" / "child.jsonl", day / f"rollout-2026-10-01T00-31-02-{CODEX_CHILD}.jsonl", + project) + chats = h / ".gemini" / "tmp" / "project" / "chats" + _copy(FIX / "gemini" / "session.jsonl", chats / "session-2026-09-27T20-58-77777777.jsonl", + project) + (chats.parent / ".project_root").write_text(f"{project}\n") + _copy(FIX / "copilot" / "events.jsonl", + h / ".copilot" / "session-state" / COPILOT_SID / "events.jsonl", project) + build_opencode_db(h / ".local" / "share" / "opencode" / "opencode.db", project) + return {"home": h, "project": project} + + +def _totals(reads) -> dict[str, float]: + out: dict[str, float] = defaultdict(float) + keyed: dict[str, dict] = {} + for sr in reads: + for ev in sr.events: + if ev.key is not None: + keyed[ev.key] = ev.usage + else: + for k, v in ev.usage.items(): + out[k] += v + for usage in keyed.values(): + for k, v in usage.items(): + out[k] += v + return dict(out) + + +def _cursors(reads): + return {sr.source: sr.cursor for sr in reads} + + +# --- Claude Code ------------------------------------------------------------------------ + + +def _claude_expected(*paths: Path) -> dict[str, float]: + seen, out = set(), defaultdict(float) + for path in paths: + for r in _records(path): + if r.get("type") != "assistant": + continue + m = r["message"] + key = (m["id"], r.get("requestId")) + if key in seen: + continue + seen.add(key) + u = m["usage"] + out["input"] += u.get("input_tokens", 0) + out["output"] += u.get("output_tokens", 0) + out["cache_read"] += u.get("cache_read_input_tokens", 0) + out["cache_write"] += u.get("cache_creation_input_tokens", 0) + return dict(out) + + +def test_claude_counts_each_request_once_and_reads_linked_subagents(home): + log = ClaudeLog(home["home"] / ".claude") + loc = log.locate(CLAUDE_SID) + reads = log.read(loc, {}) + assert [sr.source for sr in reads] == ["main", "sub:agent-a0000000000000001.jsonl"] + main = Path(loc.path) + sub = main.with_suffix("") / "subagents" / "agent-a0000000000000001.jsonl" + assert _totals(reads) == _claude_expected(main, sub) + # the fixture really does repeat usage per content block + assistant = [r for r in _records(main) if r["type"] == "assistant"] + assert len(assistant) > len({r["message"]["id"] for r in assistant}) + assert all(not sr.problems for sr in reads) + assert all(ev.ts is not None and ev.provider == "anthropic" for sr in reads + for ev in sr.events) + + +def test_claude_info_and_discovery(home): + log = ClaudeLog(home["home"] / ".claude") + loc = log.locate(CLAUDE_SID) + info = log.info(loc) + assert info.cwd == str(home["project"]) and info.parent is None + assert info.start == parse_ts(_records(Path(loc.path))[0]["timestamp"]) + found = log.discover(0) + assert [x.session_id for x in found] == [CLAUDE_SID] # sub-agent logs are not sessions + assert log.locate("nope") is None + assert log.locate(CLAUDE_SID, transcript_path=loc.path).path == loc.path + + +def test_claude_incremental_read_finds_nothing_new(home): + log = ClaudeLog(home["home"] / ".claude") + loc = log.locate(CLAUDE_SID) + first = log.read(loc, {}) + again = log.read(loc, _cursors(first)) + assert _totals(again) == {} and not any(sr.reset for sr in again) + + +# --- the JSONL tail ------------------------------------------------------------------ + + +def test_tail_leaves_a_half_written_line_for_next_time(tmp_path): + p = tmp_path / "log.jsonl" + p.write_text('{"a": 1}\n{"b": ') + recs, problems, cur, reset = tail_jsonl(p, None) + assert recs == [{"a": 1}] and not problems and not reset + assert cur["offset"] == len('{"a": 1}\n') + with open(p, "a") as f: + f.write('2}\n{"c": 3}\n') + recs, _, cur, reset = tail_jsonl(p, cur) + assert recs == [{"b": 2}, {"c": 3}] and not reset + + +def test_tail_detects_replacement_and_truncation(tmp_path): + p = tmp_path / "log.jsonl" + p.write_text('{"a": 1}\n{"b": 2}\n') + _, _, cur, _ = tail_jsonl(p, None) + p.write_text('{"a": 1}\n') # truncated: shorter than the offset + recs, _, cur2, reset = tail_jsonl(p, cur) + assert reset and recs == [{"a": 1}] + replacement = tmp_path / "new.jsonl" + replacement.write_text('{"z": 9}\n') + os.replace(replacement, p) # replaced: another inode, other head + recs, _, _, reset = tail_jsonl(p, cur2) + assert reset and recs == [{"z": 9}] + + +def test_tail_reports_unreadable_lines_and_missing_files(tmp_path): + p = tmp_path / "log.jsonl" + p.write_text('{"a": 1}\nnot json\n') + recs, problems, _, _ = tail_jsonl(p, None) + assert recs == [{"a": 1}] and len(problems) == 1 + with pytest.raises(LogMissing): + tail_jsonl(tmp_path / "gone.jsonl", None) + + +def test_records_without_a_timestamp_take_the_nearest_earlier_one(tmp_path): + log = ClaudeLog(tmp_path) + p = tmp_path / "projects" / "x" / "s.jsonl" + p.parent.mkdir(parents=True) + usage = {"input_tokens": 5, "output_tokens": 1} + lines = [{"type": "user", "timestamp": "2026-10-01T10:00:00Z", "cwd": "/w"}, + {"type": "assistant", "requestId": "r1", + "message": {"id": "m1", "model": "claude-opus-5-5", "usage": usage}}] + p.write_text("".join(json.dumps(x) + "\n" for x in lines)) + [main] = log.read(log.locate("s"), {}) + assert main.events[0].ts == parse_ts("2026-10-01T10:00:00Z") + # with no earlier timestamp in the source at all, the time stays unknown + p.write_text(json.dumps(lines[1]) + "\n") + [main] = log.read(log.locate("s"), {}) + assert main.events[0].ts is None + + +# --- Codex CLI ------------------------------------------------------------------------ + + +def _codex_final_total(path: Path) -> dict[str, float]: + total = None + for r in _records(path): + p = r.get("payload") or {} + if r["type"] == "event_msg" and p.get("type") == "token_count" and p.get("info"): + total = p["info"]["total_token_usage"] + cached = total.get("cached_input_tokens", 0) + return {"input": total["input_tokens"] - cached, "output": total["output_tokens"], + "cache_read": cached, "cache_write": total.get("cache_write_input_tokens", 0)} + + +def test_codex_sums_differences_of_running_totals_and_links_subagents(home): + log = CodexLog(home["home"] / ".codex") + loc = log.locate(CODEX_PARENT) + reads = log.read(loc, {}) + assert [sr.source for sr in reads][0] == "main" + assert any(sr.source.startswith("child:") for sr in reads) + parent_file = Path(loc.path) + child_file = parent_file.with_name(f"rollout-2026-10-01T00-31-02-{CODEX_CHILD}.jsonl") + expected = defaultdict(float) + for f in (parent_file, child_file): + for k, v in _codex_final_total(f).items(): + expected[k] += v + got = _totals(reads) + assert {k: v for k, v in got.items() if v} == {k: v for k, v in expected.items() if v} + models = {ev.model for sr in reads for ev in sr.events} + assert "unknown" not in models + # the classification of other rollout files is remembered + assert reads[0].cursor["children"] + assert _totals(log.read(loc, _cursors(reads))) == {} + + +def test_codex_subagent_names_its_parent_and_is_not_discovered(home): + log = CodexLog(home["home"] / ".codex") + child = log.locate(CODEX_CHILD) + assert log.info(child).parent == CODEX_PARENT + assert [x.session_id for x in log.discover(0)] == [CODEX_PARENT] + assert log.info(log.locate(CODEX_PARENT)).cwd == str(home["project"]) + + +# --- Gemini CLI ----------------------------------------------------------------------- + + +def test_gemini_keys_messages_by_id_and_reads_the_project_root(home): + log = GeminiLog(home["home"] / ".gemini") + loc = log.locate(GEMINI_SID) + [main] = log.read(loc, {}) + assert all(ev.key for ev in main.events) + keys = [ev.key for ev in main.events] + assert len(keys) > len(set(keys)) # Gemini rewrites messages; each is counted once + expected = defaultdict(float) + seen = {} + for r in _records(Path(loc.path)): + if r.get("type") == "gemini" and "tokens" in r: + seen[r["id"]] = r["tokens"] + for t in seen.values(): + expected["input"] += t["input"] - t.get("cached", 0) + t.get("tool", 0) + expected["output"] += t["output"] + t.get("thoughts", 0) + expected["cache_read"] += t.get("cached", 0) + assert {k: v for k, v in _totals([main]).items() if v} == \ + {k: v for k, v in expected.items() if v} + assert log.info(loc).cwd == str(home["project"]) + assert [x.session_id for x in log.discover(0)] == [GEMINI_SID] + + +# --- OpenCode ------------------------------------------------------------------------- + + +def test_opencode_reads_the_session_and_its_child_sessions(home): + db = home["home"] / ".local" / "share" / "opencode" / "opencode.db" + spec = json.loads((FIX / "opencode" / "db.json").read_text()) + parent = next(s["id"] for s in spec["session"] if s["parent_id"] is None) + log = OpenCodeLog(db.parent) + loc = log.locate(parent) + [src] = log.read(loc, {}) + expected = defaultdict(float) + for m in spec["message"]: + t = m["data"].get("tokens") + if m["data"].get("role") != "assistant" or not t: + continue + expected["input"] += t["input"] + expected["output"] += t["output"] + t.get("reasoning", 0) + expected["cache_read"] += t["cache"]["read"] + expected["cache_write"] += t["cache"]["write"] + assert {k: v for k, v in _totals([src]).items() if v} == \ + {k: v for k, v in expected.items() if v} + assert [x.session_id for x in log.discover(0)] == [parent] + child = next(s["id"] for s in spec["session"] if s["parent_id"] == parent) + assert log.info(log.locate(child)).parent == parent + + +def test_opencode_rewritten_rows_replace_their_earlier_usage(home): + db_path = home["home"] / ".local" / "share" / "opencode" / "opencode.db" + spec = json.loads((FIX / "opencode" / "db.json").read_text()) + parent = next(s["id"] for s in spec["session"] if s["parent_id"] is None) + log = OpenCodeLog(db_path.parent) + loc = log.locate(parent) + [first] = log.read(loc, {}) + msg = next(m for m in spec["message"] if m["data"].get("role") == "assistant" + and m["session_id"] == parent) + data = dict(msg["data"], tokens={**msg["data"]["tokens"], "output": 99999}) + db = sqlite3.connect(db_path) + # OpenCode stamps a rewritten row with the time of the rewrite + db.execute("UPDATE message SET data = ?, time_updated = ? WHERE id = ?", + (json.dumps(data), first.cursor["since"] + 1, msg["id"])) + db.commit() + db.close() + [second] = log.read(loc, {"db": first.cursor}) + assert [ev.key for ev in second.events] == [msg["id"]] + assert second.events[0].usage["output"] == 99999 + data["tokens"].get("reasoning", 0) + + +def test_opencode_missing_database_or_session(tmp_path, home): + assert OpenCodeLog(tmp_path / "none").locate("x") is None + db = home["home"] / ".local" / "share" / "opencode" + from aegis_core.budget.logs import Locator + with pytest.raises(LogMissing): + OpenCodeLog(db).read(Locator("opencode", "ses-gone", str(db / "opencode.db")), {}) + + +# --- Copilot CLI ---------------------------------------------------------------------- + + +def test_copilot_counts_premium_requests_once(home): + log = CopilotLog(home["home"] / ".copilot") + loc = log.locate(COPILOT_SID) + [main] = log.read(loc, {}) + totals = [r["data"]["totalPremiumRequests"] for r in _records(Path(loc.path)) + if r["type"] in ("session.usage_checkpoint", "session.shutdown")] + assert _totals([main]) == {"requests": max(totals)} + assert len(totals) > len(set(totals)) # repeated totals in the fixture are not recounted + assert log.info(loc).cwd == str(home["project"]) + assert [x.session_id for x in log.discover(0)] == [COPILOT_SID] + + +def test_discovery_honours_the_modified_since_filter(home): + log = ClaudeLog(home["home"] / ".claude") + loc = log.locate(CLAUDE_SID) + os.utime(loc.path, (1_000_000, 1_000_000)) + assert log.discover(2_000_000) == [] + shutil.rmtree(home["home"] / ".claude") + assert log.discover(0) == [] diff --git a/tests/test_budget_policy.py b/tests/test_budget_policy.py new file mode 100644 index 0000000..51ce950 --- /dev/null +++ b/tests/test_budget_policy.py @@ -0,0 +1,241 @@ +"""budget.yaml loading (design v1.0 §3) and the price table (§5).""" + +from pathlib import Path + +import pytest +import yaml + +from aegis_core.budget.policy import AGENTS, load_budget_policy +from aegis_core.budget.pricing import ( + CATEGORIES, + PriceTable, + load_builtin_table, + normalize_model, +) +from aegis_core.signing import SignatureError, sign_file + +AUTHORITY = {"admin": {"budget", "identity"}, "developer": {"configuration"}} + +FULL = { + "version": 1, + "principal": "admin", + "unit": "usd", + "session": {"limit": 20, "warn_at": 0.75}, + "project_day": {"limit": 100, "tz": "America/Chicago"}, + "agents": ["claude", "codex", "copilot"], + "on_unknown_log": "deny", + "on_unknown_price": "estimate", + "pricing": {"claude-opus-5-5": {"input": 4, "output": 20, "cache_read": 0.2, + "cache_write": 8}}, + "copilot": {"premium_requests": 50}, +} + + +def _write(tmp_path: Path, doc) -> Path: + path = tmp_path / "budget.yaml" + path.write_text(yaml.safe_dump(doc) if not isinstance(doc, str) else doc) + return path + + +def _load(tmp_path, doc, **kw): + return load_budget_policy(_write(tmp_path, doc), authority_map=AUTHORITY, insecure=True, + **kw) + + +def _with(**changes): + doc = {**FULL, **changes} + return {k: v for k, v in doc.items() if v is not None} + + +def test_full_policy_loads(tmp_path): + p = _load(tmp_path, FULL) + assert (p.unit, p.principal, p.tz) == ("usd", "admin", "America/Chicago") + assert (p.session.limit, p.session.warn_at) == (20, 0.75) + assert (p.project_day.limit, p.project_day.warn_at) == (100, 0.8) + assert p.agents == ("claude", "codex", "copilot") + assert p.copilot_premium_requests.limit == 50 + assert p.pricing["claude-opus-5-5"]["output"] == 20.0 + assert p.warnings == () + assert p.measures("claude") and p.measures("copilot") and not p.measures("gemini") + + +def test_defaults(tmp_path): + p = _load(tmp_path, {"principal": "admin", "unit": "tokens", "session": {"limit": 5000}}) + assert p.agents == AGENTS + assert (p.on_unknown_log, p.on_unknown_price, p.tz) == ("allow", "estimate", "UTC") + assert p.project_day is None and p.session.warn_at == 0.8 + # copilot is listed by default but has no premium-request limit: not measured, said so + assert not p.measures("copilot") + assert any(w.startswith("copilot-unmeasured") for w in p.warnings) + + +@pytest.mark.parametrize("doc,match", [ + (_with(unit=None), "'unit' is required"), + (_with(unit="eur"), "unit must be one of"), + (_with(extra=1), r"unknown field\(s\) \['extra'\]"), + (_with(version=2), "unsupported version"), + (_with(principal=None), "'principal' is required"), + (_with(principal="developer"), "does not hold the 'budget' class"), + (_with(session={"limit": 0}), "session.limit: must be greater than 0"), + (_with(session={"limit": -3}), "must be greater than 0"), + (_with(session={"limit": "ten"}), "must be a number"), + (_with(session={"limit": 5, "warn_at": 1}), "warn_at: must be a fraction"), + (_with(session={"limit": 5, "warn_at": 0}), "warn_at: must be a fraction"), + (_with(session={"limit": 5, "hard": True}), r"session: unknown field"), + (_with(session={"warn_at": 0.5}), "'limit' is required"), + (_with(project_day={"limit": 5, "tz": "Mars/Olympus"}), "unknown time zone"), + (_with(agents=["claude", "cursor"]), r"unknown agent\(s\) \['cursor'\]"), + (_with(agents=["claude", "claude"]), "listed twice"), + (_with(agents=[]), "non-empty list"), + (_with(on_unknown_log="warn"), "on_unknown_log must be one of"), + (_with(on_unknown_price="ignore"), "on_unknown_price must be one of"), + (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1}}), + r"missing price field\(s\) \['cache_write'\]"), + (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1, "cache_write": -1}}), + "must be a number >= 0"), + (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1, "cache_write": 1, + "batch": 1}}), "unknown price field"), + (_with(copilot={"limit": 5}), "'premium_requests'"), + (_with(session=None, project_day=None, copilot=None), "set at least one limit"), +]) +def test_shape_errors_are_load_errors(tmp_path, doc, match): + with pytest.raises(ValueError, match=match): + _load(tmp_path, doc) + + +def test_token_limits_are_whole_numbers(tmp_path): + with pytest.raises(ValueError, match="whole number of tokens"): + _load(tmp_path, {"principal": "admin", "unit": "tokens", "session": {"limit": 10.5}}) + assert _load(tmp_path, {"principal": "admin", "unit": "usd", + "session": {"limit": 10.5}}).session.limit == 10.5 + + +def test_authority_map_is_required(tmp_path): + with pytest.raises(ValueError, match="authority map is required"): + load_budget_policy(_write(tmp_path, FULL), authority_map=None, insecure=True) + + +def test_token_mode_warns_that_prices_are_ignored(tmp_path): + p = _load(tmp_path, _with(unit="tokens", session={"limit": 1000}, + project_day={"limit": 5000})) + assert "ignored: on_unknown_price has no effect with unit: tokens" in p.warnings + assert "ignored: pricing has no effect with unit: tokens" in p.warnings + + +def test_signed_policy_verifies_and_tampering_is_refused(tmp_path): + key = bytes(range(32)) + path = _write(tmp_path, FULL) + sign_file(path, key) + assert load_budget_policy(path, authority_map=AUTHORITY, key=key).session.limit == 20 + path.write_text(path.read_text().replace("limit: 20", "limit: 2000")) + with pytest.raises(SignatureError): + load_budget_policy(path, authority_map=AUTHORITY, key=key) + + +def test_unsigned_without_key_is_recorded(tmp_path): + path = _write(tmp_path, FULL) + p = load_budget_policy(path, authority_map=AUTHORITY) + assert p.warnings[0] == f"unsigned: {path}" + + +# --- prices --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def table(): + return load_builtin_table() + + +@pytest.mark.parametrize("raw,key", [ + ("claude-opus-5-5", "claude-opus-5-5"), + ("anthropic/claude-sonnet-4.5", "claude-sonnet-4-5"), + ("claude-opus-4-6[1m]", "claude-opus-4-6"), + ("claude-haiku-4-5-20251001", "claude-haiku-4-5"), + ("GPT-5.5", "gpt-5-5"), + ("gemini-2.5-pro-2026-01-15", "gemini-2-5-pro"), +]) +def test_normalize_model(raw, key): + assert normalize_model(raw) == key + + +def test_every_builtin_entry_is_complete_or_documented(table): + prices = PriceTable(table) + for vendor, spec in table["vendors"].items(): + for model, entry in spec["models"].items(): + assert prices.vendor_of(model) == vendor + missing = [c for c in CATEGORIES if c not in entry] + assert missing in ([], ["cache_read"]), (model, missing) + assert table["as_of"] + + +def test_exact_and_override_prices(table): + prices = PriceTable(table) + p = prices.price("claude-opus-5-5") + assert (p.vendor, p.estimated) == ("anthropic", frozenset()) + assert p.rates == {"input": 4.0, "output": 20.0, "cache_read": 0.2, "cache_write": 8.0} + over = PriceTable(table, {"claude-opus-5-5": {"input": 1, "output": 2, "cache_read": 3, + "cache_write": 4}}) + assert over.price("claude-opus-5-5").rates["cache_read"] == 3 + + +def test_unknown_model_of_a_known_vendor_takes_the_highest_rate_per_category(table): + prices = PriceTable(table) + p = prices.price("claude-opus-9") + anthropic = table["vendors"]["anthropic"]["models"].values() + for cat in CATEGORIES: + assert p.rates[cat] == max(m[cat] for m in anthropic if cat in m) + assert p.estimated == frozenset(CATEGORIES) + assert p.rates == {"input": 15.0, "output": 75.0, "cache_read": 1.5, "cache_write": 30.0} + + +def test_fallback_rates_are_taken_per_category_not_from_one_model(): + table = {"version": 1, "vendors": {"v": {"prefixes": ["v-"], "models": { + "v-a": {"input": 9, "output": 1, "cache_read": 1, "cache_write": 1}, + "v-b": {"input": 1, "output": 9, "cache_read": 1, "cache_write": 1}, + "v-c": {"input": 1, "output": 1, "cache_read": 9, "cache_write": 9}}}}} + p = PriceTable(table).price("v-new") + assert p.rates == {"input": 9, "output": 9, "cache_read": 9, "cache_write": 9} + + +def test_unknown_vendor_takes_the_highest_rate_in_the_whole_table(table): + prices = PriceTable(table) + p = prices.price("mystery-model-7") + assert p.vendor is None + every = [m for v in table["vendors"].values() for m in v["models"].values()] + for cat in CATEGORIES: + assert p.rates[cat] == max(m[cat] for m in every if cat in m) + + +def test_missing_category_falls_back_for_that_category_only(table): + prices = PriceTable(table) + p = prices.price("gpt-5.5-pro") + assert p.estimated == frozenset({"cache_read"}) + assert p.rates["input"] == 30.0 + openai = table["vendors"]["openai"]["models"].values() + assert p.rates["cache_read"] == max(m["cache_read"] for m in openai if "cache_read" in m) + + +def test_provider_field_identifies_the_vendor(table): + prices = PriceTable(table) + assert prices.vendor_of("some-new-model", "anthropic") == "anthropic" + assert prices.vendor_of("claude-opus-5-5", "github-copilot") == "anthropic" + assert prices.vendor_of("Qwen3-1.7B", "omlx") is None + + +def test_cost_reports_estimates_only_for_used_categories(table): + prices = PriceTable(table) + usage = {"gpt-5.5-pro": {"input": 1_000_000, "output": 0, "cache_read": 0}} + cost = prices.cost(usage) + assert cost.dollars == pytest.approx(30.0) + assert cost.estimated == {} + usage["gpt-5.5-pro"]["cache_read"] = 1_000_000 + assert prices.cost(usage).estimated == {"gpt-5.5-pro": ["cache_read"]} + + +def test_an_override_reprices_earlier_estimated_usage(table): + usage = {"claude-new-1": {"input": 2_000_000, "output": 1_000_000}} + before = PriceTable(table).cost(usage) + assert before.estimated == {"claude-new-1": ["input", "output"]} + 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) From 13362d67fe0ce5e5e72e8c80bcc8fedf628babaf Mon Sep 17 00:00:00 2001 From: moneytools Date: Thu, 8 Oct 2026 17:42:10 -0500 Subject: [PATCH 2/2] =?UTF-8?q?fix(budget):=20review=20=E2=80=94=20unverif?= =?UTF-8?q?iable=20sessions=20reach=20project=20totals,=20damaged=20invent?= =?UTF-8?q?ories=20are=20unknown,=20non-finite=20numbers=20refused?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - project entries carry each session's verification status (log missing, unreadable records, usage with no recorded time); evaluation of any session in the project sees the others', so strict mode fails closed - inventory reads report ok / missing / unreadable / corrupt (malformed, foreign or truncated lines); damaged history or a corrupt day index is 'unknown', valid checkpoints still used; inventory writes that fail no longer crash the hook - NaN and +/-inf rejected in limits, warn_at, price overrides and the table --- src/aegis_core/budget/accounting.py | 143 ++++++++++++++++++++-------- src/aegis_core/budget/evaluate.py | 14 ++- src/aegis_core/budget/policy.py | 5 +- src/aegis_core/budget/pricing.py | 7 +- tests/test_budget_accounting.py | 86 +++++++++++++++++ tests/test_budget_policy.py | 27 +++++- 6 files changed, 237 insertions(+), 45 deletions(-) diff --git a/src/aegis_core/budget/accounting.py b/src/aegis_core/budget/accounting.py index 343a948..d084142 100644 --- a/src/aegis_core/budget/accounting.py +++ b/src/aegis_core/budget/accounting.py @@ -184,6 +184,17 @@ def usage_on(self, day: str) -> ByModel: def usage_total(self) -> ByModel: return sum_by_model(*self.buckets().values()) + def unverified(self) -> list[str]: + """Why this session's usage is only a lower bound (empty if it is not).""" + reasons = [] + if self.log_missing: + reasons.append("log-missing") + if self.problem_total: + reasons.append("unreadable-records") + if self.buckets().get(UNKNOWN_DAY): + reasons.append("unknown-time") + return reasons + def to_json(self) -> str: return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) @@ -203,6 +214,7 @@ class ProjectDay: completeness: str = "complete" # complete | partial | unknown missing_sessions: int = 0 known_sessions: int = 0 + history: str = "ok" # the inventory read at the last rebuild: see read_inventory def usage(self) -> ByModel: return sum_by_model(*(e.get("usage", {}) for e in self.entries.values())) @@ -213,6 +225,12 @@ def providers(self) -> dict[str, str]: out.update(e.get("providers", {})) return out + def unverified(self, exclude: str | None = None) -> dict[str, list[str]]: + """Contributing sessions whose usage is only a lower bound, with why + (each entry carries its session's verification status).""" + return {k: list(e["unverified"]) for k, e in self.entries.items() + if e.get("unverified") and k != exclude} + def to_json(self) -> str: return json.dumps({"version": 1, **self.__dict__}, sort_keys=True) @@ -345,36 +363,76 @@ def merge_entries(pd: ProjectDay, incoming: Mapping[str, dict[str, Any]]) -> boo # the inventory (caller holds the project lock to write) @staticmethod - def _inventory_lines(path: Path) -> list[dict[str, Any]]: - out = [] + def read_inventory(path: Path) -> tuple[list[dict[str, Any]], str]: + """The valid lines of an inventory file and how the read went: + ``ok``; ``missing`` (no file); ``unreadable`` (it exists but cannot + be read); ``corrupt`` (some lines are not inventory records, + including a truncated last line). Valid lines are returned whatever + the status, so surviving checkpoints are still used.""" + if not path.exists(): + return [], "missing" try: - with open(path) as f: - for line in f: - with contextlib.suppress(json.JSONDecodeError): - d = json.loads(line) - if isinstance(d, dict): - out.append(d) + data = path.read_bytes() except OSError: - pass - return out + return [], "unreadable" + out: list[dict[str, Any]] = [] + status = "ok" + if data and not data.endswith(b"\n"): + status = "corrupt" # a truncated last line + for raw in data.splitlines(): + if not raw.strip(): + continue + try: + d = json.loads(raw) + except (json.JSONDecodeError, UnicodeDecodeError): + status = "corrupt" + continue + if (isinstance(d, dict) and isinstance(d.get("session"), str) + and d.get("event") in ("seen", "checkpoint")): + out.append(d) + else: + status = "corrupt" + return out, status - def inventory(self, root: str, day: str) -> list[dict[str, Any]] | None: - path = self.inventory_path(root, day) - return self._inventory_lines(path) if path.exists() else None - - def append_inventory(self, root: str, day: str, line: dict[str, Any]) -> None: + @classmethod + def _inventory_lines(cls, path: Path) -> list[dict[str, Any]]: + return cls.read_inventory(path)[0] + + def inventory(self, root: str, day: str) -> tuple[list[dict[str, Any]], str]: + return self.read_inventory(self.inventory_path(root, day)) + + def read_index(self, root: str) -> tuple[dict[str, Any] | None, str]: + """The project's list of days with an inventory: ``ok``, ``missing`` + or ``corrupt`` (exists but unreadable or not the expected shape).""" + path = self.index_path(root) + if not path.exists(): + return None, "missing" + d = _read_json(path) + if not isinstance(d, dict) or not isinstance(d.get("days"), list): + return None, "corrupt" + return d, "ok" + + def append_inventory(self, root: str, day: str, line: dict[str, Any]) -> bool: + """Appends one line; returns False if the inventory cannot be written + (the caller tries again on a later hook; a rebuild then reports the + damaged history as unknown).""" path = self.inventory_path(root, day) - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "a") as f: - f.write(json.dumps(line, sort_keys=True) + "\n") - f.flush() - os.fsync(f.fileno()) + try: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a") as f: + f.write(json.dumps(line, sort_keys=True) + "\n") + f.flush() + os.fsync(f.fileno()) + except OSError: + return False index = _read_json(self.index_path(root)) or {} days = set(index.get("days", [])) if day not in days: days.add(day) - _atomic_write(self.index_path(root), - json.dumps({"root": root, "days": sorted(days)}, sort_keys=True)) + with contextlib.suppress(OSError): + _atomic_write(self.index_path(root), + json.dumps({"root": root, "days": sorted(days)}, sort_keys=True)) + return True # --- reading a session (steps 1-3 of the update protocol) -------------------------------- @@ -460,7 +518,8 @@ def update_session(store: BudgetStore, log: AgentLog, loc: Locator, zone: ZoneIn 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}} + "providers": {m: p for m, p in rec.providers.items() if m in usage}, + "unverified": rec.unverified()} # --- rebuild ------------------------------------------------------------------------- @@ -484,10 +543,10 @@ def rebuild(store: BudgetStore, logs: Mapping[str, AgentLog], root: str, day: st info = log.info(loc) if info and info.parent is None and find_project_root(info.cwd) == root: candidates[key] = loc - inventory = store.inventory(root, day) + inventory, history = store.inventory(root, day) checkpoints: dict[str, dict[str, Any]] = {} in_inventory: set[str] = set() - for line in inventory or []: + for line in inventory: key = line.get("session") if not isinstance(key, str): continue @@ -506,23 +565,31 @@ def rebuild(store: BudgetStore, logs: Mapping[str, AgentLog], root: str, day: st continue with store.session_lock(loc.agent, loc.session_id): rec = update_session(store, logs[loc.agent], loc, zone) + if rec.log_missing: + missing += 1 if rec.log_missing and not rec.sources: # Neither the log nor a cached record holds this session's usage: # the inventory checkpoint is all that is left. - missing += 1 cp = checkpoints.get(key) if cp: incoming[key] = {"revision": int(cp.get("revision", 0)), - "usage": cp.get("usage", {}), "providers": {}} + "usage": cp.get("usage", {}), "providers": {}, + "unverified": ["log-missing", "from-checkpoint"]} continue + # A missing log with a cached record keeps its usage as a lower bound; + # the entry says so (its "unverified" reasons). incoming[key] = _entry(rec, day) - if inventory is not None: + if history in ("unreadable", "corrupt"): + completeness = "unknown" # the record of which sessions existed is damaged + elif history == "ok": completeness = "partial" if missing else "complete" else: - index = _read_json(store.index_path(root)) - if isinstance(index, dict): - completeness = "unknown" if day in index.get("days", []) else "complete" + index, index_status = store.read_index(root) + if index_status == "corrupt": + completeness = "unknown" + elif index is not None: + completeness = "unknown" if day in index["days"] else "complete" else: others = [k for k in candidates if k != exclude] completeness = "unknown" if others else "complete" @@ -533,6 +600,7 @@ def rebuild(store: BudgetStore, logs: Mapping[str, AgentLog], root: str, day: st pd.completeness = completeness pd.missing_sessions = missing pd.known_sessions = len(in_inventory) + pd.history = history store.save_day(pd) return pd @@ -589,15 +657,14 @@ def _keep_inventory(store: BudgetStore, rec: SessionRecord, day: str, now: float root = rec.project_root assert root changed = False - if day not in rec.inventory_days: - store.append_inventory(root, day, {"event": "seen", "session": rec.key, - "locator": rec.locator, "start": rec.start}) + if day not in rec.inventory_days and store.append_inventory( + root, day, {"event": "seen", "session": rec.key, "locator": rec.locator, + "start": rec.start}): rec.inventory_days = sorted({*rec.inventory_days, day})[-7:] changed = True - if now - rec.checkpoint_at >= CHECKPOINT_EVERY: - store.append_inventory(root, day, {"event": "checkpoint", "session": rec.key, - "revision": rec.revision, - "usage": rec.usage_on(day), "at": now}) + if now - rec.checkpoint_at >= CHECKPOINT_EVERY and store.append_inventory( + root, day, {"event": "checkpoint", "session": rec.key, "revision": rec.revision, + "usage": rec.usage_on(day), "at": now}): rec.checkpoint_at = now changed = True return changed diff --git a/src/aegis_core/budget/evaluate.py b/src/aegis_core/budget/evaluate.py index a222ec1..5bcfe94 100644 --- a/src/aegis_core/budget/evaluate.py +++ b/src/aegis_core/budget/evaluate.py @@ -139,9 +139,17 @@ def measure(name: str, unit: str, used: float, lim: Limit) -> None: 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)") - if snap.project is not None and snap.project.completeness != "complete": - unknown.append(f"today's project total is {snap.project.completeness} " - f"({snap.project.missing_sessions} session(s) without a log)") + 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) diff --git a/src/aegis_core/budget/policy.py b/src/aegis_core/budget/policy.py index 55db13f..41766b1 100644 --- a/src/aegis_core/budget/policy.py +++ b/src/aegis_core/budget/policy.py @@ -8,6 +8,7 @@ class in ``authority.yaml``, so raising a limit is a policy change an agent from __future__ import annotations +import math from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -81,6 +82,8 @@ def lim(x: Limit | None) -> dict[str, float] | None: def _number(where: str, value: Any, *, integer: bool) -> float: if isinstance(value, bool) or not isinstance(value, (int, float)): raise ValueError(f"{where}: must be a number") + if not math.isfinite(value): + raise ValueError(f"{where}: must be a finite number") if integer and value != int(value): raise ValueError(f"{where}: must be a whole number of tokens") if value <= 0: @@ -102,7 +105,7 @@ def _limit(path: Path, name: str, raw: Any, *, integer: bool, limit = _number(f"{where}.limit", raw["limit"], integer=integer) warn_at = raw.get("warn_at", DEFAULT_WARN_AT) if isinstance(warn_at, bool) or not isinstance(warn_at, (int, float)) or not ( - 0 < warn_at < 1 + math.isfinite(warn_at) and 0 < warn_at < 1 ): raise ValueError(f"{where}.warn_at: must be a fraction between 0 and 1 (e.g. 0.8)") return Limit(limit, float(warn_at)), {k: raw[k] for k in extra if k in raw} diff --git a/src/aegis_core/budget/pricing.py b/src/aegis_core/budget/pricing.py index 8f64343..9561dee 100644 --- a/src/aegis_core/budget/pricing.py +++ b/src/aegis_core/budget/pricing.py @@ -14,6 +14,7 @@ from __future__ import annotations +import math import re from collections.abc import Mapping from dataclasses import dataclass, field @@ -88,8 +89,10 @@ def _check_rates(where: str, entry: Any, *, require_all: bool) -> dict[str, floa raise ValueError(f"{where}: missing price field(s) {missing}") rates: dict[str, float] = {} for cat, value in entry.items(): - if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: - raise ValueError(f"{where}.{cat}: must be a number >= 0 (USD per million tokens)") + if (isinstance(value, bool) or not isinstance(value, (int, float)) + or not math.isfinite(value) or value < 0): + raise ValueError(f"{where}.{cat}: must be a finite number >= 0 " + "(USD per million tokens)") rates[cat] = float(value) return rates diff --git a/tests/test_budget_accounting.py b/tests/test_budget_accounting.py index d14b00a..6756d15 100644 --- a/tests/test_budget_accounting.py +++ b/tests/test_budget_accounting.py @@ -404,3 +404,89 @@ def test_merge_never_replaces_a_newer_entry(): assert BudgetStore.merge_entries(pd, {"a/1": {"revision": 6, "usage": {"m": {"input": 1}}}, "b/2": {"revision": 1, "usage": {}}}) assert pd.entries["a/1"]["revision"] == 6 and "b/2" in pd.entries + + +# --- verification status reaches the project (review, PR #36) --------------------------------- + + +def _strict_verdict(snap): + from aegis_core.budget import price_table + from aegis_core.budget.evaluate import evaluate + from aegis_core.budget.policy import BudgetPolicy, Limit + p = BudgetPolicy("budget.yaml", "admin", "tokens", Limit(10**9), Limit(10**9), "UTC", + ("claude",), "deny", "estimate") + return evaluate(p, price_table(p), snap) + + +def test_another_sessions_missing_log_with_a_cached_record_fails_strict_closed(world): + """Reviewer's reproduction: s2's log is deleted but its session cache + survives; the project/day map is deleted; s1 refreshes.""" + for sid, n in (("s1", 100), ("s2", 30)): + world.start(sid) + world.add(sid, n) + world.refresh(sid) + world.path("s2").unlink() + world.store.day_path(world.root, DAY).unlink() + snap = world.refresh("s1") + assert _input(snap.project_usage) == 130 # cached usage kept as a lower bound + assert snap.project.unverified() == {"claude/s2": ["log-missing"]} + assert snap.project.completeness == "partial" + v = _strict_verdict(snap) + assert v.decision == "deny" and "unknown-log" in v.reasons + + +@pytest.mark.parametrize("via_rebuild", [False, True]) +def test_a_malformed_record_in_another_session_fails_strict_closed(world, via_rebuild): + for sid in ("s1", "s2"): + world.start(sid) + world.add(sid, 10) + world.refresh(sid) + with open(world.path("s2"), "a") as f: + f.write("{not json\n") + if via_rebuild: + world.store.day_path(world.root, DAY).unlink() + else: + world.refresh("s2") # s2's own hook records it + snap = world.refresh("s1") + assert snap.project.unverified(exclude="claude/s1") == {"claude/s2": ["unreadable-records"]} + assert _strict_verdict(snap).decision == "deny" + assert not snap.record.unverified() # s1 itself is fine + + +@pytest.mark.parametrize("damage", ["malformed", "truncated", "foreign", "unreadable"]) +def test_a_damaged_inventory_is_unknown_history_not_an_empty_one(world, damage): + """Reviewer's reproduction: s2's log is gone, the day's inventory is + damaged, the cache is deleted, s1 refreshes.""" + for sid, n in (("s1", 100), ("s2", 30)): + world.start(sid) + world.add(sid, n) + world.refresh(sid) + world.path("s2").unlink() + inv = world.store.inventory_path(world.root, DAY) + good = inv.read_text() + if damage == "malformed": + inv.write_text("{this is not json\n" + good) + elif damage == "truncated": + inv.write_text(good[:-20]) + elif damage == "foreign": + inv.write_text('{"hello": "world"}\n' + good) + else: + inv.unlink() + inv.mkdir() # exists, cannot be read as a file + _wipe_cache(world) + snap = world.refresh("s1") + assert snap.project.completeness == "unknown" + assert snap.project.history in ("corrupt", "unreadable") + assert _strict_verdict(snap).decision == "deny" + if damage in ("malformed", "foreign"): + assert _input(snap.project_usage) == 130 # valid checkpoints still used + + +def test_a_corrupt_day_index_is_unknown(world): + world.start("s1") + world.add("s1", 100) + world.refresh("s1") + world.store.inventory_path(world.root, DAY).unlink() + world.store.index_path(world.root).write_text("garbage") + _wipe_cache(world) + assert world.refresh("s1").project.completeness == "unknown" diff --git a/tests/test_budget_policy.py b/tests/test_budget_policy.py index 51ce950..20a14d7 100644 --- a/tests/test_budget_policy.py +++ b/tests/test_budget_policy.py @@ -92,7 +92,7 @@ def test_defaults(tmp_path): (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1}}), r"missing price field\(s\) \['cache_write'\]"), (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1, "cache_write": -1}}), - "must be a number >= 0"), + "must be a finite number >= 0"), (_with(pricing={"m": {"input": 1, "output": 1, "cache_read": 1, "cache_write": 1, "batch": 1}}), "unknown price field"), (_with(copilot={"limit": 5}), "'premium_requests'"), @@ -103,6 +103,31 @@ def test_shape_errors_are_load_errors(tmp_path, doc, match): _load(tmp_path, doc) +@pytest.mark.parametrize("value", [".nan", ".inf", "-.inf"]) +def test_non_finite_numbers_are_load_errors(tmp_path, value): + """YAML .nan compares false with everything, so an unchecked NaN price or + limit would make every over-budget comparison false (review, PR #36).""" + base = "principal: admin\nunit: usd\n" + cases = [ + (base + f"session: {{limit: {value}}}\n", "finite number"), + (base + f"session: {{limit: 5, warn_at: {value}}}\n", "warn_at: must be a fraction"), + (base + f"project_day: {{limit: {value}}}\n", "finite number"), + (base + f"copilot: {{premium_requests: {value}}}\n", "finite number"), + (base + "session: {limit: 5}\npricing:\n m: {input: " + value + + ", output: 1, cache_read: 1, cache_write: 1}\n", "finite number >= 0"), + ] + for doc, match in cases: + with pytest.raises(ValueError, match=match): + _load(tmp_path, doc) + + +def test_non_finite_prices_in_a_table_are_refused(): + table = {"version": 1, "vendors": {"v": {"prefixes": ["v-"], "models": { + "v-a": {"input": float("nan"), "output": 1}}}}} + with pytest.raises(ValueError, match="finite number"): + PriceTable(table) + + def test_token_limits_are_whole_numbers(tmp_path): with pytest.raises(ValueError, match="whole number of tokens"): _load(tmp_path, {"principal": "admin", "unit": "tokens", "session": {"limit": 10.5}})