feat: add review effort accounting and budget governor
This commit is contained in:
+60
-11
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user