pragent/pilot/review/budget.py

"""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,
}

PROFILES = (
    # (changed lines threshold, steps, total tokens, output tokens, seconds)
    (2_000, 80, 1_200_000, 80_000, 1_800),
    (800, 60, 800_000, 60_000, 1_200),
    (200, 40, 400_000, 40_000, 900),
)


@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,
        )

    @classmethod
    def for_review(cls, config: dict | None, diff: str) -> "Budget":
        """Choose safe headroom from diff size, unless config is explicit."""
        if isinstance((config or {}).get("budget"), dict):
            return cls.from_config(config)
        changed_lines = _changed_line_count(diff)
        for threshold, steps, tokens, output, seconds in PROFILES:
            if changed_lines >= threshold:
                return cls.from_config({
                    **(config or {}),
                    "budget": {
                        "max_steps": steps,
                        "max_total_tokens": tokens,
                        "max_output_tokens": output,
                        "max_duration_seconds": seconds,
                    },
                })
        return cls.from_config(config)


def _changed_line_count(diff: str) -> int:
    """Count changed lines without treating hunk headers as additions."""
    return sum(
        1 for line in diff.splitlines()
        if (line.startswith("+") and not line.startswith("+++"))
        or (line.startswith("-") and not line.startswith("---"))
    )


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