feat: add review effort accounting and budget governor
This commit is contained in:
@@ -27,6 +27,11 @@ pair to rotate instead of two. Tags carry the finer cut:
|
||||
`provider:headroom`, `model:<bare>`, `engine:opencode`, `repo:<owner/name>`,
|
||||
`lens:<id>` 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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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 <key of cost_model.PRICES> — per-repo override
|
||||
cost_target <key of cost_model.PRICES> — 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
|
||||
|
||||
|
||||
|
||||
+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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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):")
|
||||
|
||||
Reference in New Issue
Block a user