feat: add review effort accounting and budget governor

pilot/review/budget.py

@@ -0,0 +1,142 @@
+"""Trusted review budget policy and thread-safe accounting."""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+from dataclasses import dataclass
+
+
+DEFAULTS = {
+    "max_steps": 20,
+    "max_total_tokens": 120_000,
+    "max_output_tokens": 20_000,
+    "max_duration_seconds": 480,
+}
+
+
+@dataclass(frozen=True)
+class Budget:
+    max_steps: int = DEFAULTS["max_steps"]
+    max_total_tokens: int = DEFAULTS["max_total_tokens"]
+    max_output_tokens: int = DEFAULTS["max_output_tokens"]
+    max_duration_seconds: int = DEFAULTS["max_duration_seconds"]
+    max_lenses: int = 4
+    max_equivalent_cost_usd: float | None = None
+    price_target: str = "claude-sonnet-5"
+
+    @classmethod
+    def from_config(cls, config: dict | None) -> "Budget":
+        values = dict(DEFAULTS)
+        values["max_lenses"] = _env_int("PRAGENT_MAX_PARALLEL_LENSES", 4)
+        env_map = {
+            "max_steps": "PRAGENT_MAX_REVIEW_STEPS",
+            "max_total_tokens": "PRAGENT_MAX_REVIEW_TOKENS",
+            "max_output_tokens": "PRAGENT_MAX_REVIEW_OUTPUT_TOKENS",
+            "max_duration_seconds": "PRAGENT_REVIEW_TIMEOUT",
+            "max_lenses": "PRAGENT_MAX_PARALLEL_LENSES",
+        }
+        for key, env_key in env_map.items():
+            if env_key in os.environ:
+                value = _env_int(env_key, values[key])
+                if value > 0:
+                    values[key] = value
+        raw = (config or {}).get("budget")
+        if isinstance(raw, dict):
+            for key in set(DEFAULTS) | {"max_lenses", "max_equivalent_cost_usd"}:
+                if key in raw:
+                    values[key] = raw[key]
+        cost = values.get("max_equivalent_cost_usd")
+        price_target = str(
+            (config or {}).get("cost_target")
+            or os.environ.get("PRAGENT_PRICE_TARGET")
+            or "claude-sonnet-5"
+        )
+        return cls(
+            max_steps=int(values["max_steps"]),
+            max_total_tokens=int(values["max_total_tokens"]),
+            max_output_tokens=int(values["max_output_tokens"]),
+            max_duration_seconds=int(values["max_duration_seconds"]),
+            max_lenses=int(values["max_lenses"]),
+            max_equivalent_cost_usd=float(cost) if cost is not None else None,
+            price_target=price_target,
+        )
+
+
+def _env_int(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, default))
+    except (TypeError, ValueError):
+        return default
+
+
+class BudgetState:
+    """Cumulative accounting shared by all subprocesses in one review."""
+
+    def __init__(self, budget: Budget):
+        self.budget = budget
+        self.started = time.monotonic()
+        self.steps = 0
+        self.total_tokens = 0
+        self.output_tokens = 0
+        self.equivalent_cost_usd = 0.0
+        self.cap_reason = ""
+        self._lock = threading.RLock()
+
+    def record(self, usage: dict, equivalent_cost_usd: float = 0.0) -> str:
+        with self._lock:
+            self.steps += int(usage.get("steps") or 0)
+            self.total_tokens += int(usage.get("total") or 0)
+            self.output_tokens += int(usage.get("output") or 0)
+            self.equivalent_cost_usd += equivalent_cost_usd
+            reason = self.reason()
+            if reason:
+                self.cap_reason = reason
+            return reason
+
+    def reason(self) -> str:
+        with self._lock:
+            if self.steps >= self.budget.max_steps:
+                return "max_steps"
+            if self.total_tokens >= self.budget.max_total_tokens:
+                return "max_total_tokens"
+            if self.output_tokens >= self.budget.max_output_tokens:
+                return "max_output_tokens"
+            if time.monotonic() - self.started >= self.budget.max_duration_seconds:
+                return "max_duration_seconds"
+            if (
+                self.budget.max_equivalent_cost_usd is not None
+                and self.equivalent_cost_usd >= self.budget.max_equivalent_cost_usd
+            ):
+                return "max_equivalent_cost_usd"
+            return ""
+
+    def snapshot(self) -> dict:
+        with self._lock:
+            return {
+                "steps": self.steps,
+                "total_tokens": self.total_tokens,
+                "output_tokens": self.output_tokens,
+                "equivalent_cost_usd": round(self.equivalent_cost_usd, 6),
+                "cap_hit": bool(self.cap_reason),
+                "cap_reason": self.cap_reason or None,
+            }
+
+
+def equivalent_cost(usage: dict, model: str, price_target: str = "") -> float:
+    """Estimate comparison cost for one completed iteration."""
+    try:
+        from cost_model import PRICES, Usage, cost
+        target = price_target or os.environ.get("PRAGENT_PRICE_TARGET", "claude-sonnet-5")
+        price = PRICES.get(target)
+        if price is None:
+            return 0.0
+        return cost(Usage(
+            uncached_input=max(0, int(usage.get("input") or 0) - int(usage.get("cache_read") or 0)),
+            cached_input=int(usage.get("cache_read") or 0),
+            cache_writes=int(usage.get("cache_write") or 0),
+            output=int(usage.get("output") or 0),
+        ), price)
+    except Exception:
+        return 0.0