feat(pilot): token-usage reporting gated by AI-USAGE label
Add per-review + per-comment token accounting, surfaced only when a PR carries the new AI-USAGE label (on top of the existing AI-REVIEW trigger). opencode_review: - run_opencode now uses `--format json`; parse_opencode_events reconstructs the assistant text from `text` events and sums tokens/cost/steps from every `step_finish` event (tolerant of noise / missing fields). - run() measures duration_s around the opencode call and returns (text, usage). - changed_files(diff) extracts the `+++ b/` paths; the brief now lists them under a "Changed files" focus block so the agent grounds findings in the diff's neighbourhood instead of unbounded whole-repo walks. ai_review: - format_usage_section renders a `## AI usage` block: measured totals (in/out/reasoning/cache/cost/steps/duration), the whole-repo scope note, and an attributed per-finding table. Per-comment counts are output tokens split by each finding's body weight — labelled "attributed" since one model pass produces all findings. - inline_comment_body appends `🪙 ~N tok (X% · attributed output)` when attribution is present. - review_pr gains report_usage; compute_attribution stashes _tok_attrib/_tok_pct. - format_review_body inserts the usage section between summary and findings. webhook_server: - Fire on every pull_request action except `closed` (denylist, was an allowlist) — the AI-REVIEW gate + sha dedupe keep this safe. - AI-USAGE label detection + PRAGENT_USAGE_ALWAYS env drive report_usage. .opencode factory + review-methodology skill: new "Ground findings in context" step — read callers/imports/sibling functions per changed file (1-3 files per finding), no unbounded walks. Tests: parse_opencode_events (text+usage sum, malformed tolerance, none-usage), changed_files, compute_attribution math, inline 🪙 line, format_usage_section totals/table/cost, format_review_body ordering. 68 passing. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+132
-15
@@ -34,10 +34,12 @@ Env:
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import tarfile
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
@@ -143,6 +145,34 @@ def _extract_tar_strip_one(blob: bytes, dest: str) -> None:
|
||||
|
||||
BRIEF_PATH = ".pragent/brief.md"
|
||||
|
||||
# Matches unified-diff new-file path headers: `+++ b/path` (and `+++ /dev/null`
|
||||
# for deletions, which we skip). Captures the path after the `b/` prefix.
|
||||
_NEW_FILE_HEADER_RE = re.compile(r"^\+\+\+ b/(.+?)\s*$")
|
||||
|
||||
|
||||
def changed_files(diff: str) -> list[str]:
|
||||
"""Extract the sorted list of changed file paths from a unified diff.
|
||||
|
||||
Pulled from `+++ b/<path>` headers (the post-change side). Deletions
|
||||
(`+++ /dev/null`) are excluded. Used to give the agent a clean focus list
|
||||
for context research, so it reads callers/imports of the actually-changed
|
||||
files instead of re-deriving them from the raw diff.
|
||||
"""
|
||||
out = []
|
||||
seen = set()
|
||||
for line in (diff or "").splitlines():
|
||||
if not line.startswith("+++ b/"):
|
||||
continue
|
||||
m = _NEW_FILE_HEADER_RE.match(line)
|
||||
if not m:
|
||||
continue
|
||||
path = m.group(1).strip()
|
||||
if path and path not in seen:
|
||||
seen.add(path)
|
||||
out.append(path)
|
||||
return sorted(out)
|
||||
|
||||
|
||||
_BRIEF_TEMPLATE = """\
|
||||
# pragent review brief
|
||||
|
||||
@@ -156,6 +186,14 @@ _BRIEF_TEMPLATE = """\
|
||||
## Description
|
||||
{description}
|
||||
|
||||
## Changed files (focus your context research here)
|
||||
{changed_files}
|
||||
|
||||
For each changed file, read its callers, imports, sibling functions, and type
|
||||
definitions so findings reflect how the change is actually used — don't flag a
|
||||
hunk in isolation. Stop once a finding is grounded (1–3 related files per
|
||||
finding; avoid runaway whole-repo walks).
|
||||
|
||||
## Repo review config (.pr-review.json)
|
||||
{config}
|
||||
|
||||
@@ -198,12 +236,15 @@ def write_brief(
|
||||
prior = "\n\n---\n\n".join(prior_reviews)
|
||||
if len(prior) > 8000:
|
||||
prior = prior[:8000] + "\n…[prior reviews truncated]"
|
||||
files = changed_files(diff)
|
||||
files_block = "\n".join(f"- `{p}`" for p in files) if files else "_(none)_"
|
||||
content = _BRIEF_TEMPLATE.format(
|
||||
repo=repo or "?",
|
||||
index=index or "?",
|
||||
sha=sha or "?",
|
||||
title=title or "(none)",
|
||||
description=description.strip() or "_(none)_",
|
||||
changed_files=files_block,
|
||||
config=cfg,
|
||||
prior=prior,
|
||||
diff=diff or "_(empty)_",
|
||||
@@ -233,6 +274,66 @@ def drop_factory(workdir: str) -> None:
|
||||
# opencode invocation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _new_usage() -> dict:
|
||||
return {
|
||||
"input": 0, "output": 0, "reasoning": 0,
|
||||
"cache_read": 0, "cache_write": 0, "total": 0,
|
||||
"cost": 0.0, "steps": 0,
|
||||
}
|
||||
|
||||
|
||||
def parse_opencode_events(stdout: str) -> tuple[str, dict | None]:
|
||||
"""Parse `opencode run --format json` NDJSON stdout into (text, usage).
|
||||
|
||||
- assistant text: concatenation of every `{"type":"text","part":{"text":…}}`
|
||||
event, in order → the agent's full message (prose + the findings ```json
|
||||
block). This is what `ai_review.parse_review_output` then extracts the
|
||||
findings JSON from.
|
||||
- usage: summed across every `{"type":"step_finish","part":{"tokens":…,
|
||||
"cost":…}}` event (one per model turn). Returns a dict with input/output/
|
||||
reasoning/cache_read/cache_write/total/cost/steps, or None if no
|
||||
step_finish was seen (e.g. empty/failed run).
|
||||
|
||||
Tolerant: non-JSON lines, missing fields, or non-dict events are skipped
|
||||
(warm-up / log noise / tool events we don't care about). Never raises.
|
||||
"""
|
||||
text_parts: list[str] = []
|
||||
usage = _new_usage()
|
||||
saw_step = False
|
||||
for line in (stdout or "").splitlines():
|
||||
line = line.strip()
|
||||
if not line or not line.startswith("{"):
|
||||
continue
|
||||
try:
|
||||
ev = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if not isinstance(ev, dict):
|
||||
continue
|
||||
etype = ev.get("type")
|
||||
part = ev.get("part") or {}
|
||||
if etype == "text" and isinstance(part, dict):
|
||||
t = part.get("text")
|
||||
if isinstance(t, str):
|
||||
text_parts.append(t)
|
||||
elif etype == "step_finish" and isinstance(part, dict):
|
||||
tok = part.get("tokens") or {}
|
||||
if isinstance(tok, dict):
|
||||
saw_step = True
|
||||
usage["steps"] += 1
|
||||
usage["input"] += int(tok.get("input") or 0)
|
||||
usage["output"] += int(tok.get("output") or 0)
|
||||
usage["reasoning"] += int(tok.get("reasoning") or 0)
|
||||
cache = tok.get("cache") or {}
|
||||
if isinstance(cache, dict):
|
||||
usage["cache_read"] += int(cache.get("read") or 0)
|
||||
usage["cache_write"] += int(cache.get("write") or 0)
|
||||
usage["total"] += int(tok.get("total") or 0)
|
||||
cost = part.get("cost")
|
||||
if isinstance(cost, (int, float)):
|
||||
usage["cost"] += float(cost)
|
||||
return "".join(text_parts), (usage if saw_step else None)
|
||||
|
||||
_PROMPT = (
|
||||
"Read .pragent/brief.md and review this pull request as pragent. "
|
||||
"Load the review-methodology and findings-schema skills, inspect the "
|
||||
@@ -333,8 +434,12 @@ def _warm_opencode(home: str, model: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def run_opencode(workdir: str, model: str, timeout: int | None = None) -> str:
|
||||
"""Run the pragent agent headlessly in `workdir`. Returns the agent's stdout.
|
||||
def run_opencode(workdir: str, model: str, timeout: int | None = None) -> tuple[str, dict | None]:
|
||||
"""Run the pragent agent headlessly in `workdir`. Returns `(text, usage)`.
|
||||
|
||||
`text` is the reconstructed assistant message (prose + findings JSON) from
|
||||
the `--format json` event stream; `usage` is the summed token/cost usage
|
||||
across all model turns (or None if no step_finish event was seen).
|
||||
|
||||
Isolates from the host user's global opencode config by pointing HOME at a
|
||||
shared temp dir (so ~/.config/opencode is not merged) and passing --pure
|
||||
@@ -342,7 +447,11 @@ def run_opencode(workdir: str, model: str, timeout: int | None = None) -> str:
|
||||
drop_factory) are the only project config discovered; the shared home's
|
||||
global opencode.json supplies the provider/model/permission. PATH prepends
|
||||
the rtk dir so the agent's bash tool can call `rtk`. Warms the HOME first
|
||||
(cold runs produce no output) and retries once on empty stdout.
|
||||
(cold runs produce no output) and retries once on empty text.
|
||||
|
||||
`--format json` makes opencode emit NDJSON events (text + step_finish with
|
||||
token usage) instead of formatted stdout — `parse_opencode_events` turns
|
||||
that into the assistant text + a usage dict.
|
||||
|
||||
stdin=DEVNULL is critical: opencode blocks on stdin (permission prompt /
|
||||
interactive input) when run headlessly via subprocess, hanging until timeout.
|
||||
@@ -356,6 +465,7 @@ def run_opencode(workdir: str, model: str, timeout: int | None = None) -> str:
|
||||
bin_,
|
||||
"run",
|
||||
"--pure",
|
||||
"--format", "json",
|
||||
"--agent", "pragent",
|
||||
"--dir", workdir,
|
||||
"--model", model,
|
||||
@@ -371,10 +481,13 @@ def run_opencode(workdir: str, model: str, timeout: int | None = None) -> str:
|
||||
except subprocess.TimeoutExpired as e:
|
||||
last_err = f"opencode timed out after {e.timeout}s"
|
||||
continue
|
||||
out = (proc.stdout or "").strip()
|
||||
if out:
|
||||
return proc.stdout
|
||||
last_err = f"opencode empty stdout (rc={proc.returncode}); stderr: {(proc.stderr or '')[-1500:]}"
|
||||
text, usage = parse_opencode_events(proc.stdout or "")
|
||||
if text.strip():
|
||||
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")
|
||||
|
||||
|
||||
@@ -396,16 +509,18 @@ def run(
|
||||
config: dict | None,
|
||||
prior_reviews: list[str] | None,
|
||||
model: str,
|
||||
) -> str:
|
||||
"""End-to-end: checkout archive → brief → drop factory → opencode → stdout.
|
||||
) -> tuple[str, dict | None]:
|
||||
"""End-to-end: checkout archive → brief → drop factory → opencode → (text, usage).
|
||||
|
||||
Returns the raw opencode stdout (summary + findings JSON). Raises on any
|
||||
failure; the caller (review_pr) fails open. The workdir is removed unless
|
||||
PRAGENT_KEEP_WORK is set.
|
||||
Returns the reconstructed opencode assistant text (summary + findings JSON)
|
||||
and a usage dict (token/cost totals + `duration_s`), or `(text, None)` when
|
||||
no usage events were seen. Raises on any failure; the caller (`review_pr`)
|
||||
fails open. The workdir is removed unless PRAGENT_KEEP_WORK is set.
|
||||
"""
|
||||
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"))
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
fetch_archive(api, repo, sha, token, workdir)
|
||||
write_brief(
|
||||
@@ -414,10 +529,12 @@ def run(
|
||||
diff=diff, config=config, prior_reviews=prior_reviews,
|
||||
)
|
||||
drop_factory(workdir)
|
||||
stdout = run_opencode(workdir, model)
|
||||
if not stdout.strip():
|
||||
text, usage = run_opencode(workdir, model)
|
||||
if not text.strip():
|
||||
raise RuntimeError("opencode produced no output")
|
||||
return stdout
|
||||
if usage is not None:
|
||||
usage["duration_s"] = round(time.monotonic() - t0, 1)
|
||||
return text, usage
|
||||
finally:
|
||||
if not keep:
|
||||
shutil.rmtree(workdir, ignore_errors=True)
|
||||
Reference in New Issue
Block a user