feat: add review effort accounting and budget governor

This commit is contained in:
Claude
2026-09-01 11:41:08 +00:00
parent a202bd3598
commit 3a10deef06
15 changed files with 522 additions and 23 deletions
+60 -11
View File
@@ -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: