From 3a10deef068263a4d3d532f7a1afa17bd20a3ecd Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 11:41:08 +0000 Subject: [PATCH 1/4] feat: add review effort accounting and budget governor --- docs/architecture.md | 5 + pilot/README-langfuse.md | 5 + pilot/README-webhook.md | 18 +++ pilot/README.md | 3 + pilot/observability/langfuse.py | 17 +++ pilot/review/budget.py | 142 ++++++++++++++++++ pilot/review/configuration.py | 46 ++++++ pilot/review/opencode.py | 71 +++++++-- pilot/review/opencode_lenses.py | 17 ++- pilot/review/opencode_runtime.py | 98 +++++++++++- pilot/review/output.py | 11 ++ .../observability_tests/langfuse_test.py | 19 +++ tests/pilot/review_tests/ai_review_test.py | 36 ++++- tests/pilot/review_tests/budget_test.py | 49 ++++++ .../review_tests/opencode_workspace_test.py | 8 +- 15 files changed, 522 insertions(+), 23 deletions(-) create mode 100644 pilot/review/budget.py create mode 100644 tests/pilot/review_tests/budget_test.py diff --git a/docs/architecture.md b/docs/architecture.md index 0f207a2..7eca1ed 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -18,6 +18,7 @@ review_pr facade/orchestrator │ ├── opencode_workspace archive, sanitization, brief, factory │ ├── opencode_lens_config reviewer configuration and selection │ └── opencode_synthesis normalization, deduplication, summaries + ├── review/budget trusted limits and cumulative accounting ├── review parsing normalize findings + validate anchors ├── feedback persist reactions and derive scores └── langfuse_trace usage, cost, evaluation telemetry @@ -50,6 +51,10 @@ The internal seams are deliberately narrower: transport works. - `langfuse_trace` is an optional sink. It is fail-open and cannot change the review result. +- `review/budget` owns deployment and repository resource limits. The runner + streams model events and terminates the subprocess after a completed step + crosses a step, token, duration, or equivalent-cost limit. Cap status is + retained in review output and Langfuse metadata. ## Trust model diff --git a/pilot/README-langfuse.md b/pilot/README-langfuse.md index 4f24741..dbcc396 100644 --- a/pilot/README-langfuse.md +++ b/pilot/README-langfuse.md @@ -27,6 +27,11 @@ pair to rotate instead of two. Tags carry the finer cut: `provider:headroom`, `model:`, `engine:opencode`, `repo:`, `lens:` per fan-out lens. +Each trace metadata record also includes the completed iteration count, tool +calls, per-iteration token details, and whether a configured budget stopped the +run. A capped run is tagged in the review body and can be filtered in Langfuse +with `cap_hit` / `cap_reason` metadata. + To split into two *projects* later, point `LANGFUSE_PUBLIC_KEY` / `LANGFUSE_SECRET_KEY` at the second project on whichever deployment runs the Claude path. Nothing in the code needs to change. diff --git a/pilot/README-webhook.md b/pilot/README-webhook.md index a7e5bc0..2e946af 100644 --- a/pilot/README-webhook.md +++ b/pilot/README-webhook.md @@ -272,6 +272,14 @@ Set `.pr-review.json: "reviewers": []` to opt out (single-primary fallback). { "id": "my-lens", "agent_file": ".opencode/agents/my-lens.md", "model": "headroom/glm-5.2:cloud" } ], "triage": { "enabled": true, "max_lenses": 4 }, + "budget": { + "max_steps": 20, + "max_total_tokens": 120000, + "max_output_tokens": 20000, + "max_duration_seconds": 480, + "max_lenses": 4, + "max_equivalent_cost_usd": 1.00 + }, "max_findings": 7 } ``` @@ -285,9 +293,19 @@ files. Fail-open: if triage errors, all lenses run. | var | default | effect | |-----|---------|--------| | `PRAGENT_MAX_PARALLEL_LENSES` | 4 | cap concurrency | +| `PRAGENT_MAX_REVIEW_STEPS` | 20 | maximum completed model iterations per review | +| `PRAGENT_MAX_REVIEW_TOKENS` | 120000 | maximum cumulative tokens per review | +| `PRAGENT_MAX_REVIEW_OUTPUT_TOKENS` | 20000 | maximum generated tokens per review | +| `PRAGENT_REVIEW_TIMEOUT` | 480 | maximum review duration (s) | | `PRAGENT_LENS_TIMEOUT` | 540 | per-lens subprocess timeout (s) | | `PRAGENT_REVIEWERS` | unset | force multi-lens fan-out even without `reviewers[]` | +Budget limits can also be set per repository in the trusted base-branch +`.pr-review.json`. The process terminates after a completed iteration crosses a +limit, preserves any output already emitted, and records the cap reason in the +review body and Langfuse metadata. Environment variables provide deployment-wide +defaults; repository budget values override them. + **Cross-lens dedup:** synthesiser drops duplicates by `sha256[:16](path|line|severity|problem[:80])` (matches the feedback DB's `posthash`), then promotes multi-lens agreement by one severity step diff --git a/pilot/README.md b/pilot/README.md index cbbddb4..39c300d 100644 --- a/pilot/README.md +++ b/pilot/README.md @@ -74,6 +74,9 @@ security, and webhook registration details. | `PRAGENT_ENGINE` | `opencode` | `opencode` or legacy model path | | `DIFF_MAX_CHARS` | `150000` | Diff input cap | | `PRAGENT_MAX_CONCURRENT_REVIEWS` | `2` | Process concurrency bound | +| `PRAGENT_MAX_REVIEW_STEPS` | `20` | Maximum completed model iterations per review | +| `PRAGENT_MAX_REVIEW_TOKENS` | `120000` | Maximum cumulative tokens per review | +| `PRAGENT_MAX_REVIEW_OUTPUT_TOKENS` | `20000` | Maximum generated tokens per review | | `LANGFUSE_HOST` + keys | unset | Enables telemetry; unset is a no-op | ## Tests diff --git a/pilot/observability/langfuse.py b/pilot/observability/langfuse.py index 34d761c..73807a3 100644 --- a/pilot/observability/langfuse.py +++ b/pilot/observability/langfuse.py @@ -249,6 +249,11 @@ def build_batch( tags.append(f"tier:{tier}") for lens in lenses or []: tags.append(f"lens:{lens}") + if usage.get("budget_cap_hit"): + tags.extend([ + "budget:capped", + f"budget:{usage.get('budget_cap_reason', 'unknown')}", + ]) costs, cost_basis = _cost_details(usage, model, price_target) if usage else ({}, "") if cost_basis: @@ -267,7 +272,19 @@ def build_batch( "severities": _severity_counts(findings), "provider_cost_usd": usage.get("cost"), "cost_basis": cost_basis or None, + "iterations": usage.get("steps"), + "tool_calls": usage.get("tool_calls"), + "cap_hit": usage.get("budget_cap_hit"), + "cap_reason": usage.get("budget_cap_reason"), + "tokens_per_finding": round( + float(usage.get("total") or 0) / max(1, len(findings or [])), 2 + ), + "steps_per_finding": round( + float(usage.get("steps") or 0) / max(1, len(findings or [])), 2 + ), } + if usage.get("iterations"): + metadata["iteration_usage"] = usage["iterations"][:50] if lenses: metadata["lenses"] = lenses if tier: diff --git a/pilot/review/budget.py b/pilot/review/budget.py new file mode 100644 index 0000000..709e93e --- /dev/null +++ b/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 diff --git a/pilot/review/configuration.py b/pilot/review/configuration.py index daa2be9..b50f3ec 100644 --- a/pilot/review/configuration.py +++ b/pilot/review/configuration.py @@ -25,6 +25,11 @@ CONFIG_MAX_INSTRUCTIONS_CHARS = 4000 CONFIG_MAX_PATTERNS_ITEMS = 16 # allow + deny separately, total 32 entries CONFIG_MAX_FINDINGS = 30 CONFIG_MAX_STATIC_MESSAGE_CHARS = 400 # free-text banner, mirror of instructions +MAX_BUDGET_STEPS = 100 +MAX_BUDGET_TOKENS = 2_000_000 +MAX_BUDGET_SECONDS = 3_600 +MAX_BUDGET_LENSES = 8 +MAX_BUDGET_COST_USD = 100.0 STYLES = frozenset(STYLE_DEFAULTS) SEVERITY_VALUES = frozenset(SEVERITIES) @@ -49,6 +54,8 @@ def parse_repo_config(raw: str) -> dict: patterns {allow:[…], deny:[…]} — post-filter globs model — per-repo override cost_target — see equivalent_cost + budget {max_steps, max_total_tokens, max_output_tokens, + max_duration_seconds, max_lenses, max_equivalent_cost_usd} additional_context_urls list[str] (≤ 8) — see fetch_additional_context """ if not raw: @@ -193,6 +200,45 @@ def parse_repo_config(raw: str) -> dict: if cleaned: out["compare_against"] = cleaned[:12] + budget = _parse_budget(data.get("budget")) + if budget: + out["budget"] = budget + + return out + + +def _parse_budget(raw) -> dict: + """Sanitize optional per-review resource limits from trusted config.""" + if not isinstance(raw, dict): + return {} + out: dict = {} + integer_limits = { + "max_steps": (1, MAX_BUDGET_STEPS), + "max_total_tokens": (1, MAX_BUDGET_TOKENS), + "max_output_tokens": (1, MAX_BUDGET_TOKENS), + "max_duration_seconds": (1, MAX_BUDGET_SECONDS), + "max_lenses": (1, MAX_BUDGET_LENSES), + } + for key, (lo, hi) in integer_limits.items(): + value = raw.get(key) + if isinstance(value, int) and not isinstance(value, bool): + if lo <= value <= hi: + out[key] = value + elif isinstance(value, str) and value.strip().isdigit(): + number = int(value.strip()) + if lo <= number <= hi: + out[key] = number + cost = raw.get("max_equivalent_cost_usd") + if isinstance(cost, (int, float)) and not isinstance(cost, bool): + if 0 < float(cost) <= MAX_BUDGET_COST_USD: + out["max_equivalent_cost_usd"] = float(cost) + elif isinstance(cost, str): + try: + number = float(cost.strip()) + except ValueError: + number = 0 + if 0 < number <= MAX_BUDGET_COST_USD: + out["max_equivalent_cost_usd"] = number return out diff --git a/pilot/review/opencode.py b/pilot/review/opencode.py index c994be5..7c54dce 100644 --- a/pilot/review/opencode.py +++ b/pilot/review/opencode.py @@ -73,6 +73,7 @@ from .opencode_lenses import ( filter_by_skip_if, intersect_with_triage, merge_usage, run_lenses, ) from . import opencode_runtime as _runtime +from .budget import Budget, BudgetState _filter_by_skip_if = filter_by_skip_if _intersect_with_triage = intersect_with_triage @@ -112,7 +113,7 @@ def _new_usage() -> dict: return { "input": 0, "output": 0, "reasoning": 0, "cache_read": 0, "cache_write": 0, "total": 0, - "cost": 0.0, "steps": 0, + "cost": 0.0, "steps": 0, "tool_calls": 0, "iterations": [], } @@ -146,6 +147,8 @@ def parse_opencode_events(stdout: str) -> tuple[str, dict | None]: continue etype = ev.get("type") part = ev.get("part") or {} + if etype in ("tool_use", "tool_result", "tool_call"): + usage["tool_calls"] += 1 if etype == "text" and isinstance(part, dict): t = part.get("text") if isinstance(t, str): @@ -166,6 +169,18 @@ def parse_opencode_events(stdout: str) -> tuple[str, dict | None]: cost = part.get("cost") if isinstance(cost, (int, float)): usage["cost"] += float(cost) + usage["iterations"].append({ + "step": usage["steps"], + "input": int(tok.get("input") or 0), + "output": int(tok.get("output") or 0), + "reasoning": int(tok.get("reasoning") or 0), + "cache_read": int((cache or {}).get("read") or 0) + if isinstance(cache, dict) else 0, + "cache_write": int((cache or {}).get("write") or 0) + if isinstance(cache, dict) else 0, + "total": int(tok.get("total") or 0), + "cost": float(cost) if isinstance(cost, (int, float)) else 0.0, + }) return "".join(text_parts), (usage if saw_step else None) _PROMPT = ( @@ -299,6 +314,17 @@ def _balanced_jsons(text: str): start = None +def _run_process( + cmd, *, cwd, env, timeout, parse_events, budget=None, budget_state=None, + model="", +): + return _runtime._run_process( + cmd, cwd=cwd, env=env, timeout=timeout, parse_events=parse_events, + budget=budget, budget_state=budget_state, model=model, + runner=subprocess.run, + ) + + def triage( @@ -307,6 +333,8 @@ def triage( reviewers: list[ReviewerSpec], default_model: str, factory_root: str, + budget: Budget | None = None, + budget_state: BudgetState | None = None, ) -> list[str] | None: """Run the triage agent. Returns the lens subset with surface. @@ -344,9 +372,10 @@ def triage( prompt, ] try: - proc = subprocess.run( - cmd, cwd=workdir, env=env, capture_output=True, text=True, - stdin=subprocess.DEVNULL, timeout=120, + proc = _run_process( + cmd, cwd=workdir, env=env, timeout=min(120, budget.max_duration_seconds) + if budget else 120, parse_events=parse_opencode_events, + budget=budget, budget_state=budget_state, model=default_model, ) except (subprocess.TimeoutExpired, Exception) as e: print(f"pragent: triage crashed: {e}; falling back to all lenses", flush=True) @@ -401,6 +430,8 @@ def run_lenses_review( model: str, compression_note: str = "", additional_context: str = "", + budget: Budget | None = None, + budget_state: BudgetState | None = None, ) -> tuple[str, dict | None]: """Fan-out + synthesize path. Returns (merged-text, merged-usage). @@ -409,6 +440,8 @@ def run_lenses_review( always has: prose + a final ```json fence with the legacy schema). """ os.makedirs(WORK_ROOT, exist_ok=True) + budget = budget or Budget.from_config(config) + budget_state = budget_state or BudgetState(budget) workdir = tempfile.mkdtemp(prefix=f"{repo.replace('/', '_')}-{sha[:8]}-", dir=WORK_ROOT) keep = bool(os.environ.get("PRAGENT_KEEP_WORK")) t0 = time.monotonic() @@ -429,14 +462,16 @@ def run_lenses_review( # Edge case: reviewers[] present but every entry had activation:off. # Fall back to single-primary. return _fallback_single_primary( - workdir=workdir, model=model, + workdir=workdir, model=model, budget=budget, + budget_state=budget_state, ) triage_cfg = parse_triage_config((config or {}).get("triage")) changed_paths = changed_files(diff) - reviewers = _filter_by_skip_if(reviewers, changed_paths) + reviewers = _filter_by_skip_if(reviewers, changed_paths)[:budget.max_lenses] selected = triage( workdir, triage_cfg, reviewers, model, _factory_dir(), + budget=budget, budget_state=budget_state, ) if selected is not None: if not selected: @@ -453,12 +488,15 @@ def run_lenses_review( return _no_surface_response(repo, index, sha, 0) factory_root = _factory_dir() - results = run_lenses(workdir, reviewers, model, factory_root) + results = run_lenses( + workdir, reviewers, model, factory_root, budget, budget_state, + ) # Merge findings + usage across lenses findings_per_lens = {lid: r[0] for lid, r in results.items()} merged = synthesize(findings_per_lens, reviewers) merged_usage = merge_usage([r[1] for r in results.values()]) + merged_usage.update({f"budget_{k}": v for k, v in budget_state.snapshot().items()}) # Build a synthetic text response that ai_review.parse_review_output # can consume (prose summary + final ```json fence with legacy schema). @@ -543,10 +581,15 @@ def _no_surface_response( return text, None -def _fallback_single_primary(workdir: str, model: str) -> tuple[str, dict | None]: +def _fallback_single_primary( + workdir: str, model: str, budget: Budget | None = None, + budget_state: BudgetState | None = None, +) -> tuple[str, dict | None]: """Used when reviewers[] resolves to empty (all activation:off).""" try: - text, usage = run_opencode(workdir, model) + text, usage = run_opencode( + workdir, model, budget=budget, budget_state=budget_state, + ) return text, usage except Exception as e: print(f"pragent: fallback single-primary failed: {e}", flush=True) @@ -576,12 +619,14 @@ def _warm_opencode(home: str, model: str) -> None: def run_opencode( workdir: str, model: str, timeout: int | None = None, + budget: Budget | None = None, budget_state: BudgetState | None = None, ) -> tuple[str, dict | None]: return _runtime.run_opencode( workdir, model, opencode_bin=_opencode_bin(), shared_home_fn=_shared_home, warm_fn=_warm_opencode, build_environment=_build_env, parse_events=parse_opencode_events, - prompt=_PROMPT, timeout=timeout or TIMEOUT, runner=subprocess.run, + prompt=_PROMPT, timeout=timeout or TIMEOUT, budget=budget, + budget_state=budget_state, runner=subprocess.run, ) @@ -642,6 +687,8 @@ def run( additional_context=additional_context, ) + budget = Budget.from_config(config) + budget_state = BudgetState(budget) os.makedirs(WORK_ROOT, exist_ok=True) workdir = tempfile.mkdtemp(prefix=f"{repo.replace('/', '_')}-{sha[:8]}-", dir=WORK_ROOT) keep = bool(os.environ.get("PRAGENT_KEEP_WORK")) @@ -663,7 +710,9 @@ def run( additional_context=additional_context, ) drop_factory(workdir) - text, usage = run_opencode(workdir, model) + text, usage = run_opencode( + workdir, model, budget=budget, budget_state=budget_state, + ) if not text.strip(): raise RuntimeError("opencode produced no output") if usage is not None: diff --git a/pilot/review/opencode_lenses.py b/pilot/review/opencode_lenses.py index a5ed4e1..e7c21ef 100644 --- a/pilot/review/opencode_lenses.py +++ b/pilot/review/opencode_lenses.py @@ -7,7 +7,7 @@ from .opencode_synthesis import _normalize_lens_finding LENS_TIMEOUT_S = 540 -def _run_one_lens(workdir, spec, model, factory_root): +def _run_one_lens(workdir, spec, model, factory_root, budget=None, budget_state=None): """Run one lens through the compatibility module's runtime seam.""" from . import opencode as oc @@ -28,9 +28,12 @@ def _run_one_lens(workdir, spec, model, factory_root): "--agent", spec.id, "--dir", workdir, "--model", model, prompt, ] try: - proc = oc.subprocess.run( - cmd, cwd=workdir, env=env, capture_output=True, text=True, - stdin=oc.subprocess.DEVNULL, timeout=LENS_TIMEOUT_S, + proc = oc._run_process( + cmd, cwd=workdir, env=env, + timeout=min(LENS_TIMEOUT_S, budget.max_duration_seconds) + if budget else LENS_TIMEOUT_S, + parse_events=oc.parse_opencode_events, budget=budget, + budget_state=budget_state, model=model, ) except oc.subprocess.TimeoutExpired: print(f"pragent: lens {spec.id} timed out after {LENS_TIMEOUT_S}s", flush=True) @@ -69,7 +72,7 @@ def _run_one_lens(workdir, spec, model, factory_root): return normalized, usage, spec.id -def run_lenses(workdir, reviewers, default_model, factory_root): +def run_lenses(workdir, reviewers, default_model, factory_root, budget=None, budget_state=None): """Run configured lenses in parallel and return results by lens id.""" if not reviewers: return {} @@ -81,7 +84,7 @@ def run_lenses(workdir, reviewers, default_model, factory_root): futures = { ex.submit( _run_one_lens, workdir, spec, - spec.model or default_model, factory_root, + spec.model or default_model, factory_root, budget, budget_state, ): spec for spec in reviewers } @@ -131,5 +134,5 @@ def merge_usage(parts): for key in base: if isinstance(base[key], (int, float)): base[key] += usage.get(key, 0) or 0 + base["iterations"].extend(usage.get("iterations", []) or []) return base - diff --git a/pilot/review/opencode_runtime.py b/pilot/review/opencode_runtime.py index c9cc3af..534f2d8 100644 --- a/pilot/review/opencode_runtime.py +++ b/pilot/review/opencode_runtime.py @@ -2,6 +2,9 @@ import os import subprocess +import time + +from .budget import Budget, BudgetState, equivalent_cost _ENV_ALLOW = frozenset({ @@ -66,6 +69,7 @@ def warm_opencode( def run_opencode( workdir, model, *, opencode_bin, shared_home_fn, warm_fn, build_environment, parse_events, prompt, timeout, + budget: Budget | None = None, budget_state: BudgetState | None = None, runner=subprocess.run, ): home = shared_home_fn() @@ -78,19 +82,107 @@ def run_opencode( last_err = "" for _ in range(2): try: - proc = runner( - cmd, cwd=workdir, env=env, capture_output=True, text=True, - stdin=subprocess.DEVNULL, timeout=timeout, + proc = _run_process( + cmd, cwd=workdir, env=env, timeout=timeout, + parse_events=parse_events, budget=budget, + budget_state=budget_state, model=model, runner=runner, ) except subprocess.TimeoutExpired as exc: last_err = f"opencode timed out after {exc.timeout}s" continue text, usage = parse_events(proc.stdout or "") + if usage and budget_state: + usage.update({f"budget_{k}": v for k, v in budget_state.snapshot().items()}) if text.strip(): return text, usage + if budget_state and budget_state.cap_reason: + reason = budget_state.cap_reason + empty_usage = usage or { + "input": 0, "output": 0, "reasoning": 0, + "cache_read": 0, "cache_write": 0, "total": 0, + "cost": 0.0, "steps": 0, "tool_calls": 0, "iterations": [], + } + empty_usage.update({ + f"budget_{k}": v for k, v in budget_state.snapshot().items() + }) + return ( + "Review stopped before a complete response was produced " + f"because the budget reached {reason}.\n\n" + "```json\n{\"summary\": \"Review budget reached\", " + "\"findings\": []}\n```\n", + empty_usage, + ) + if usage and usage.get("budget_cap_hit"): + return text, usage last_err = ( f"opencode empty text (rc={proc.returncode}); " f"stderr: {(proc.stderr or '')[-1500:]}" ) raise RuntimeError(last_err or "opencode produced no output") + +def _run_process( + cmd, *, cwd, env, timeout, parse_events, budget, budget_state, model, runner, +): + """Run a process, terminating it after a completed event exceeds budget.""" + if budget is None or budget_state is None: + return runner( + cmd, cwd=cwd, env=env, capture_output=True, text=True, + stdin=subprocess.DEVNULL, timeout=timeout, + ) + existing_reason = budget_state.reason() + if existing_reason: + budget_state.cap_reason = existing_reason + return subprocess.CompletedProcess(cmd, 0, "", "") + proc = subprocess.Popen( + cmd, cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, text=True, + ) + output: list[str] = [] + previous = {"steps": 0, "total": 0, "output": 0, "cost": 0.0} + cap_reason = "" + started = time.monotonic() + try: + assert proc.stdout is not None + for line in proc.stdout: + output.append(line) + _, usage = parse_events("".join(output)) + if usage: + delta = { + "steps": usage.get("steps", 0) - previous["steps"], + "input": usage.get("input", 0) - previous.get("input", 0), + "cache_read": usage.get("cache_read", 0) - previous.get("cache_read", 0), + "cache_write": usage.get("cache_write", 0) - previous.get("cache_write", 0), + "total": usage.get("total", 0) - previous["total"], + "output": usage.get("output", 0) - previous["output"], + } + cost = float(usage.get("cost", 0.0)) - previous["cost"] + previous.update({ + "steps": usage.get("steps", 0), + "input": usage.get("input", 0), + "cache_read": usage.get("cache_read", 0), + "cache_write": usage.get("cache_write", 0), + "total": usage.get("total", 0), + "output": usage.get("output", 0), + "cost": float(usage.get("cost", 0.0)), + }) + cap_reason = budget_state.record( + delta, equivalent_cost(delta, model, budget.price_target), + ) + if cap_reason or time.monotonic() - started >= budget.max_duration_seconds: + cap_reason = cap_reason or "max_duration_seconds" + budget_state.cap_reason = cap_reason + proc.terminate() + break + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + finally: + if proc.stdout: + proc.stdout.close() + stderr = proc.stderr.read() if proc.stderr else "" + return subprocess.CompletedProcess( + cmd, proc.returncode, "".join(output), stderr, + ) diff --git a/pilot/review/output.py b/pilot/review/output.py index 15726ab..c92d2c8 100644 --- a/pilot/review/output.py +++ b/pilot/review/output.py @@ -741,6 +741,17 @@ def _render_collapsible_usage(usage: dict | None, model: str, config: dict | Non f"- **Actual**: {actual_s}{actual_note}", f"- **Scope**: {scope}", ] + if usage.get("budget_cap_hit"): + lines.append( + f"- **Budget**: capped at `{usage.get('budget_cap_reason', 'configured limit')}`" + ) + budget = (config or {}).get("budget") or {} + if budget: + limits = ", ".join( + f"{key.removeprefix('max_')}={value}" + for key, value in budget.items() + ) + lines.append(f"- **Budget limits**: {limits}") if eq_rows: lines.append("") lines.append("- **Equivalent cost on paid providers** (this run's tokens):") diff --git a/tests/pilot/observability_tests/langfuse_test.py b/tests/pilot/observability_tests/langfuse_test.py index e4c79ee..c9d0148 100644 --- a/tests/pilot/observability_tests/langfuse_test.py +++ b/tests/pilot/observability_tests/langfuse_test.py @@ -159,6 +159,25 @@ def test_batch_without_usage_has_no_generation(): assert types[0] == "trace-create" +def test_batch_exposes_iteration_and_budget_metadata(): + usage = { + **USAGE, + "tool_calls": 6, + "budget_cap_hit": True, + "budget_cap_reason": "max_steps", + "iterations": [{"step": 1, "total": 100}], + } + batch = lt.build_batch( + model="headroom/glm-5.2:cloud", **{**BASE, "usage": usage} + ) + metadata = batch[0]["body"]["metadata"] + assert metadata["iterations"] == 28 + assert metadata["tool_calls"] == 6 + assert metadata["cap_hit"] is True + assert metadata["cap_reason"] == "max_steps" + assert metadata["iteration_usage"] == [{"step": 1, "total": 100}] + + def test_trace_carries_repo_pr_session_and_severity_counts(): batch = lt.build_batch(model="headroom/glm-5.2:cloud", **BASE) body = batch[0]["body"] diff --git a/tests/pilot/review_tests/ai_review_test.py b/tests/pilot/review_tests/ai_review_test.py index 2a936ab..c25be82 100644 --- a/tests/pilot/review_tests/ai_review_test.py +++ b/tests/pilot/review_tests/ai_review_test.py @@ -429,6 +429,41 @@ def test_parse_repo_config_static_message_ignores_blank(): assert "static_message" not in parse_repo_config(json.dumps({"static_message": 42})) +def test_parse_repo_config_sanitizes_budget_limits(): + cfg = parse_repo_config(json.dumps({ + "budget": { + "max_steps": "20", + "max_total_tokens": 120000, + "max_output_tokens": 20000, + "max_duration_seconds": 480, + "max_lenses": 4, + "max_equivalent_cost_usd": "1.25", + "unknown": 99, + } + })) + assert cfg["budget"] == { + "max_steps": 20, + "max_total_tokens": 120000, + "max_output_tokens": 20000, + "max_duration_seconds": 480, + "max_lenses": 4, + "max_equivalent_cost_usd": 1.25, + } + + +def test_parse_repo_config_drops_invalid_budget_values(): + cfg = parse_repo_config(json.dumps({ + "budget": { + "max_steps": 0, + "max_total_tokens": 999999999, + "max_duration_seconds": -1, + "max_lenses": 99, + "max_equivalent_cost_usd": 0, + } + })) + assert "budget" not in cfg + + def test_parse_repo_config_reads_model_override(): # Per-repo override is validated against cost_model.PRICES. Only keys # the cost model knows about can override the review engine. @@ -2148,4 +2183,3 @@ def test_format_review_body_confidence_clamps_out_of_range(): assert "Merge confidence: 5/5 🟢" in body_hi body_lo = format_review_body("- x", "glm-5.2:cloud", "abcdef1234567890", confidence=0) assert "Merge confidence: 1/5 🔴" in body_lo - diff --git a/tests/pilot/review_tests/budget_test.py b/tests/pilot/review_tests/budget_test.py new file mode 100644 index 0000000..58bb5d1 --- /dev/null +++ b/tests/pilot/review_tests/budget_test.py @@ -0,0 +1,49 @@ +"""Budget policy and accounting tests.""" + +import os +import sys +import json + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +sys.path.insert(0, os.path.join(ROOT, "pilot")) + +from review.budget import Budget, BudgetState # noqa: E402 +from review import opencode_runtime # noqa: E402 +from review.opencode import parse_opencode_events # noqa: E402 + + +def test_budget_reads_config_over_environment(monkeypatch): + monkeypatch.setenv("PRAGENT_MAX_REVIEW_STEPS", "3") + budget = Budget.from_config({"budget": {"max_steps": 7}}) + assert budget.max_steps == 7 + + +def test_budget_state_stops_at_token_limit(): + state = BudgetState(Budget(max_steps=20, max_total_tokens=100)) + assert state.record({"steps": 1, "total": 60, "output": 10}) == "" + assert state.record({"steps": 1, "total": 40, "output": 10}) == "max_total_tokens" + assert state.snapshot()["cap_hit"] is True + + +def test_budget_state_tracks_cost_cap(): + state = BudgetState(Budget(max_equivalent_cost_usd=1.0)) + assert state.record({"steps": 1, "total": 1}, 0.75) == "" + assert state.record({"steps": 1, "total": 1}, 0.25) == "max_equivalent_cost_usd" + + +def test_process_terminates_after_step_budget(): + code = ( + "import json,time; " + "print(json.dumps({'type':'step_finish','part':{'tokens':{" + "'input':1,'output':1,'total':2}}}), flush=True); " + "time.sleep(30)" + ) + budget = Budget(max_steps=1, max_duration_seconds=10) + state = BudgetState(budget) + proc = opencode_runtime._run_process( + [sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(), + timeout=10, parse_events=parse_opencode_events, budget=budget, + budget_state=state, model="glm-5.2:cloud", runner=None, + ) + assert state.snapshot()["cap_reason"] == "max_steps" + assert proc.stdout.count("step_finish") == 1 diff --git a/tests/pilot/review_tests/opencode_workspace_test.py b/tests/pilot/review_tests/opencode_workspace_test.py index f95a753..9f59de2 100644 --- a/tests/pilot/review_tests/opencode_workspace_test.py +++ b/tests/pilot/review_tests/opencode_workspace_test.py @@ -201,6 +201,13 @@ def test_parse_events_text_and_usage_summed(): assert usage["cache_write"] == 1 assert usage["total"] == 150 assert abs(usage["cost"] - 0.01) < 1e-9 + assert usage["tool_calls"] == 0 + assert usage["iterations"] == [ + {"step": 1, "input": 90, "output": 10, "reasoning": 0, + "cache_read": 5, "cache_write": 0, "total": 100, "cost": 0.0}, + {"step": 2, "input": 40, "output": 10, "reasoning": 2, + "cache_read": 0, "cache_write": 1, "total": 50, "cost": 0.01}, + ] def test_parse_events_no_step_finish_returns_none_usage(): @@ -481,4 +488,3 @@ def test_committed_config_has_no_private_address(): # --------------------------------------------------------------------------- # Multi-lens orchestration # --------------------------------------------------------------------------- - -- 2.52.0 From f9eef969e6352ed60d68c3db0d278e4ca3833501 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 11:42:59 +0000 Subject: [PATCH 2/4] docs: document review budget defaults --- README.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/README.md b/README.md index b81f9d2..2caa290 100644 --- a/README.md +++ b/README.md @@ -150,10 +150,21 @@ rates, calibrated against runs measured through the usage telemetry (`OBSERVED_RUNS` in that file — append to it, don't guess). Tokens are summed from opencode `step_finish` events per review. +Each review is governed by hard limits: 20 completed steps, 120,000 total +tokens, 20,000 output tokens, and 480 seconds by default. Limits can be +overridden deployment-wide with `PRAGENT_MAX_REVIEW_STEPS`, +`PRAGENT_MAX_REVIEW_TOKENS`, `PRAGENT_MAX_REVIEW_OUTPUT_TOKENS`, and +`PRAGENT_REVIEW_TIMEOUT`, or per repository in the trusted base-branch +`.pr-review.json`; repository values win. + Two measured reviews of a ~1100-line PR in this repo: 28 and 31 agent steps, ~2.1M input tokens each, **zero cache reads or writes**. The demo repo's PR, same tier: 126K tokens. +The measurements above are historical uncapped runs. A capped run preserves +completed output, reports the cap reason in the review, and records it in +Langfuse. + | Model | this repo, ~1100-line PR | demo repo PR | |---|---:|---:| | Claude Opus 5 | ~$10.79 | ~$0.71 | -- 2.52.0 From 4d51e14a1abad488c8c69eddae5bbfe4270acfd4 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 11:50:46 +0000 Subject: [PATCH 3/4] feat: adapt review budget to diff size --- README.md | 5 ++++ pilot/README-webhook.md | 5 ++++ pilot/review/budget.py | 35 +++++++++++++++++++++++++ pilot/review/opencode.py | 4 +-- tests/pilot/review_tests/budget_test.py | 13 +++++++++ 5 files changed, 60 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 2caa290..6dded1b 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,11 @@ overridden deployment-wide with `PRAGENT_MAX_REVIEW_STEPS`, `PRAGENT_REVIEW_TIMEOUT`, or per repository in the trusted base-branch `.pr-review.json`; repository values win. +For repositories with broad diffs, the pilot automatically raises headroom to +40/400K, 60/800K, or 80/1.2M steps/tokens as changed lines cross 200, 800, or +2,000. Explicit repository budgets always take precedence, and the global +hard ceilings remain in force. + Two measured reviews of a ~1100-line PR in this repo: 28 and 31 agent steps, ~2.1M input tokens each, **zero cache reads or writes**. The demo repo's PR, same tier: 126K tokens. diff --git a/pilot/README-webhook.md b/pilot/README-webhook.md index 2e946af..9749fe1 100644 --- a/pilot/README-webhook.md +++ b/pilot/README-webhook.md @@ -306,6 +306,11 @@ limit, preserves any output already emitted, and records the cap reason in the review body and Langfuse metadata. Environment variables provide deployment-wide defaults; repository budget values override them. +Without an explicit repository budget, changed diffs receive adaptive headroom: +the default 20-step/120K-token budget grows to 40/400K, 60/800K, or 80/1.2M for +diffs over 200, 800, or 2,000 changed lines. This keeps focused PRs inexpensive +while allowing broad TypeScript/Go reviews to finish. Hard ceilings still apply. + **Cross-lens dedup:** synthesiser drops duplicates by `sha256[:16](path|line|severity|problem[:80])` (matches the feedback DB's `posthash`), then promotes multi-lens agreement by one severity step diff --git a/pilot/review/budget.py b/pilot/review/budget.py index 709e93e..59b1f82 100644 --- a/pilot/review/budget.py +++ b/pilot/review/budget.py @@ -15,6 +15,13 @@ DEFAULTS = { "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: @@ -63,6 +70,34 @@ class Budget: 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: diff --git a/pilot/review/opencode.py b/pilot/review/opencode.py index 7c54dce..ce6f5cf 100644 --- a/pilot/review/opencode.py +++ b/pilot/review/opencode.py @@ -440,7 +440,7 @@ def run_lenses_review( always has: prose + a final ```json fence with the legacy schema). """ os.makedirs(WORK_ROOT, exist_ok=True) - budget = budget or Budget.from_config(config) + budget = budget or Budget.for_review(config, diff) budget_state = budget_state or BudgetState(budget) workdir = tempfile.mkdtemp(prefix=f"{repo.replace('/', '_')}-{sha[:8]}-", dir=WORK_ROOT) keep = bool(os.environ.get("PRAGENT_KEEP_WORK")) @@ -687,7 +687,7 @@ def run( additional_context=additional_context, ) - budget = Budget.from_config(config) + budget = Budget.for_review(config, diff) budget_state = BudgetState(budget) os.makedirs(WORK_ROOT, exist_ok=True) workdir = tempfile.mkdtemp(prefix=f"{repo.replace('/', '_')}-{sha[:8]}-", dir=WORK_ROOT) diff --git a/tests/pilot/review_tests/budget_test.py b/tests/pilot/review_tests/budget_test.py index 58bb5d1..bb5e6da 100644 --- a/tests/pilot/review_tests/budget_test.py +++ b/tests/pilot/review_tests/budget_test.py @@ -18,6 +18,19 @@ def test_budget_reads_config_over_environment(monkeypatch): assert budget.max_steps == 7 +def test_budget_scales_for_broad_diff(): + diff = "".join("+changed\n" for _ in range(850)) + budget = Budget.for_review({}, diff) + assert budget.max_steps == 60 + assert budget.max_total_tokens == 800_000 + + +def test_explicit_budget_wins_over_diff_profile(): + diff = "".join("+changed\n" for _ in range(2_100)) + budget = Budget.for_review({"budget": {"max_steps": 9}}, diff) + assert budget.max_steps == 9 + + def test_budget_state_stops_at_token_limit(): state = BudgetState(Budget(max_steps=20, max_total_tokens=100)) assert state.record({"steps": 1, "total": 60, "output": 10}) == "" -- 2.52.0 From 592747d98dd9bf862fb43c30a95be7e3b743a235 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 1 Sep 2026 12:14:18 +0000 Subject: [PATCH 4/4] fix: resolve review budget feedback --- pilot/review/budget.py | 22 +++++++++++----------- pilot/review/opencode_runtime.py | 24 +++++++++++++++++++++++- tests/pilot/review_tests/budget_test.py | 16 +++++++++++++++- 3 files changed, 49 insertions(+), 13 deletions(-) diff --git a/pilot/review/budget.py b/pilot/review/budget.py index 59b1f82..f30780b 100644 --- a/pilot/review/budget.py +++ b/pilot/review/budget.py @@ -163,15 +163,15 @@ 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: + except (ImportError, ModuleNotFoundError): return 0.0 + 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) diff --git a/pilot/review/opencode_runtime.py b/pilot/review/opencode_runtime.py index 534f2d8..f1de8b5 100644 --- a/pilot/review/opencode_runtime.py +++ b/pilot/review/opencode_runtime.py @@ -1,6 +1,7 @@ """Isolated opencode process runtime.""" import os +import selectors import subprocess import time @@ -130,6 +131,8 @@ def _run_process( cmd, cwd=cwd, env=env, capture_output=True, text=True, stdin=subprocess.DEVNULL, timeout=timeout, ) + if runner not in (None, subprocess.run): + raise ValueError("custom runners are unsupported for budgeted streaming") existing_reason = budget_state.reason() if existing_reason: budget_state.cap_reason = existing_reason @@ -142,9 +145,27 @@ def _run_process( previous = {"steps": 0, "total": 0, "output": 0, "cost": 0.0} cap_reason = "" started = time.monotonic() + selector = selectors.DefaultSelector() try: assert proc.stdout is not None - for line in proc.stdout: + selector.register(proc.stdout, selectors.EVENT_READ) + while selector.get_map(): + remaining = budget.max_duration_seconds - (time.monotonic() - started) + if remaining <= 0: + cap_reason = "max_duration_seconds" + budget_state.cap_reason = cap_reason + proc.terminate() + break + events = selector.select(timeout=remaining) + if not events: + cap_reason = "max_duration_seconds" + budget_state.cap_reason = cap_reason + proc.terminate() + break + line = proc.stdout.readline() + if not line: + selector.unregister(proc.stdout) + break output.append(line) _, usage = parse_events("".join(output)) if usage: @@ -180,6 +201,7 @@ def _run_process( proc.kill() proc.wait() finally: + selector.close() if proc.stdout: proc.stdout.close() stderr = proc.stderr.read() if proc.stderr else "" diff --git a/tests/pilot/review_tests/budget_test.py b/tests/pilot/review_tests/budget_test.py index bb5e6da..91369a6 100644 --- a/tests/pilot/review_tests/budget_test.py +++ b/tests/pilot/review_tests/budget_test.py @@ -3,6 +3,7 @@ import os import sys import json +import subprocess ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) sys.path.insert(0, os.path.join(ROOT, "pilot")) @@ -56,7 +57,20 @@ def test_process_terminates_after_step_budget(): proc = opencode_runtime._run_process( [sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(), timeout=10, parse_events=parse_opencode_events, budget=budget, - budget_state=state, model="glm-5.2:cloud", runner=None, + budget_state=state, model="glm-5.2:cloud", runner=subprocess.run, ) assert state.snapshot()["cap_reason"] == "max_steps" assert proc.stdout.count("step_finish") == 1 + + +def test_process_terminates_silent_child_at_duration_budget(): + code = "import time; time.sleep(30)" + budget = Budget(max_steps=20, max_duration_seconds=1) + state = BudgetState(budget) + proc = opencode_runtime._run_process( + [sys.executable, "-u", "-c", code], cwd=".", env=os.environ.copy(), + timeout=10, parse_events=parse_opencode_events, budget=budget, + budget_state=state, model="glm-5.2:cloud", runner=subprocess.run, + ) + assert state.snapshot()["cap_reason"] == "max_duration_seconds" + assert proc.stdout == "" -- 2.52.0