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:
Marcos
2026-08-18 04:15:11 +00:00
parent 76b6752f48
commit 087834565d
7 changed files with 522 additions and 46 deletions
+113 -5
View File
@@ -134,14 +134,16 @@ def parse_text_blocks(content: list) -> str:
return "\n".join(out).strip()
def format_review_body(findings: str, model: str, sha: str, summary: str = "") -> str:
def format_review_body(findings: str, model: str, sha: str, summary: str = "", usage_section: str = "") -> str:
"""Format the posted review summary body.
`findings` is the bullet text for findings that could NOT be anchored inline
(or, on the legacy/no-inline path, the whole review). Empty -> "No issues
found.". `summary` (optional, opencode engine) is rendered as a "Summary"
section right under the header. The hidden sha marker is always appended
for the dedupe pass.
section right under the header. `usage_section` (optional, shown only when
the PR carries the `AI-USAGE` label) is rendered between the summary and the
findings bullets. The hidden sha marker is always appended for the dedupe
pass.
"""
header = REVIEW_HEADER.format(model=model, sha=sha[:8] if sha else "unknown")
findings = (findings or "").strip()
@@ -151,6 +153,8 @@ def format_review_body(findings: str, model: str, sha: str, summary: str = "") -
parts = [header]
if summary:
parts.append(summary.strip())
if usage_section:
parts.append(usage_section.strip())
parts.append(findings)
body = "\n\n".join(parts)
if marker:
@@ -158,6 +162,88 @@ def format_review_body(findings: str, model: str, sha: str, summary: str = "") -
return body
def _finding_weight(f: dict) -> int:
"""Body-weight used to attribute output tokens to a finding (char length of
its rendered problem + fix + suggestion). One model pass produces all
findings, so per-finding tokens can't be measured directly — we split the
measured output total by this weight as an honest attribution."""
return (
len(f.get("problem") or "")
+ len(f.get("fix") or "")
+ len(f.get("suggestion") or "")
)
def compute_attribution(findings: list[dict], output_tokens: int) -> None:
"""Stash `_tok_attrib` (attributed output tokens) and `_tok_pct` (0..1) on
each finding, splitting `output_tokens` by each finding's body weight.
Mutates in place. No-op when there are no findings or no output budget."""
if not findings or not output_tokens:
return
weights = [_finding_weight(f) for f in findings]
total_w = sum(weights)
if total_w <= 0:
# All-zero weights (no prose): split evenly.
share = output_tokens / len(findings)
for f in findings:
f["_tok_attrib"] = int(round(share))
f["_tok_pct"] = 1.0 / len(findings)
return
for f, w in zip(findings, weights):
f["_tok_attrib"] = int(round(output_tokens * w / total_w))
f["_tok_pct"] = w / total_w
def format_usage_section(usage: dict | None, findings: list[dict], model: str) -> str:
"""Render the `## 🔋 AI usage` block for the review body.
Only called when the PR carries the `AI-USAGE` label (and the opencode
engine produced a usage dict). Reports the MEASURED total
(input/output/reasoning/cache/cost/steps/duration) plus an ATTRIBUTED
per-finding table — one model pass generates all findings, so per-comment
counts are an estimate (output split by body weight), clearly labelled.
Returns "" if `usage` is None.
"""
if not usage:
return ""
dur = usage.get("duration_s")
dur_s = f"{dur}s" if dur is not None else "?"
cost = usage.get("cost") or 0.0
cost_s = f"${cost:.4f}" if cost else "$0.00"
cost_note = (
"(on-network glm-5.2:cloud via headroom — no per-token charge)"
if not cost else "(billed by provider)"
)
lines = [
"## 🔋 AI usage",
"",
f"- model: `{model}` · engine: opencode · agent steps: {usage.get('steps', 0)} · duration: {dur_s}",
(
f"- tokens: {usage.get('input', 0)} in · {usage.get('output', 0)} out · "
f"{usage.get('reasoning', 0)} reasoning · cache "
f"{usage.get('cache_read', 0)} read / {usage.get('cache_write', 0)} write "
f"{usage.get('total', 0)} total"
),
f"- est. cost: {cost_s} {cost_note}",
"- scope: whole-repo checkout at head sha (agent can read any file + run linters, not just the diff) — input tokens include files read beyond the diff",
"- per-comment tokens are *attributed* (one model pass produces all findings; output split by each finding's body weight)",
]
# Per-finding attribution table.
rows = [f for f in findings if f.get("_tok_attrib") is not None]
if rows:
lines.append("")
lines.append("| severity | location | ≈out tok | % |")
lines.append("|---|---|---:|---:|")
for f in rows:
loc = f"{f['path']}:{f['line']}" if f.get("line") else f.get("path", "?")
pct = f.get("_tok_pct", 0.0) * 100
lines.append(
f"| {f.get('severity', '').upper()} | `{loc}` | "
f"{f.get('_tok_attrib', 0)} | {pct:.0f}% |"
)
return "\n".join(lines)
def build_user_prompt(
title: str,
body: str,
@@ -483,6 +569,10 @@ def inline_comment_body(f: dict) -> str:
ref = f.get("reference", "")
if ref:
body += f"\n\n📎 ref: {ref}"
tok = f.get("_tok_attrib")
if tok is not None:
pct = (f.get("_tok_pct", 0.0) or 0.0) * 100
body += f"\n\n🪙 ~{tok} tok ({pct:.0f}% · attributed output)"
return body
@@ -713,6 +803,7 @@ def review_pr(
model: str,
max_tokens: int = 8000,
max_chars: int = 150000,
report_usage: bool = False,
) -> bool:
"""Run one review and post it as `pragent-bot`.
@@ -722,6 +813,11 @@ def review_pr(
review with inline comments + suggestions (unanchored findings → summary
bullets).
`report_usage`: when True (PR carries the `AI-USAGE` label), the opencode
engine's measured token/cost usage is rendered as a `## 🔋 AI usage` section
on the review body and an attributed `🪙 ~N tok` line on each inline
comment. No-op on the ollama fallback (no usage available).
Returns True on success (including a deliberate skip), False on failure
(failure note posted when possible). Never raises — fail-open by design.
Both the CI `run()` entry point and the central webhook server call this.
@@ -752,7 +848,7 @@ def review_pr(
# `model` here is the bare id (OLLAMA_MODEL). OPENCODE_MODEL overrides
# with the full ref; otherwise we prefix the configured provider.
oc_model = os.environ.get("OPENCODE_MODEL") or f"headroom/{model}"
stdout = opencode_review.run(
stdout, usage = opencode_review.run(
api=api, repo=repo, index=index, sha=sha, token=token,
title=title, body=body, diff=diff, config=config,
prior_reviews=prior, model=oc_model,
@@ -767,6 +863,15 @@ def review_pr(
user_prompt = build_user_prompt(title, body, diff, config, prior)
raw_findings = call_model(ollama_url, model, SYSTEM_PROMPT, user_prompt, max_tokens)
findings = parse_findings(raw_findings)
usage = None
# Attribute output tokens to each finding (mutates finding dicts) so
# inline comments + the usage table can show a per-comment estimate.
# Only meaningful when we have measured usage AND the PR asked for it.
usage_section = ""
if report_usage and usage and usage.get("output"):
compute_attribution(findings, usage["output"])
usage_section = format_usage_section(usage, findings, model)
anchors = parse_diff_anchors(diff)
anchored, unanchored = split_findings(findings, anchors)
@@ -782,7 +887,10 @@ def review_pr(
summary_parts.append(bullets)
if not summary_parts:
summary_parts.append("No issues found.")
summary_body = format_review_body("\n\n".join(summary_parts), model, sha, summary=review_summary)
summary_body = format_review_body(
"\n\n".join(summary_parts), model, sha,
summary=review_summary, usage_section=usage_section,
)
post_inline_review(api, repo, index, token, summary_body, anchored)
print(
+132 -15
View File
@@ -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 (13 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)
+34 -15
View File
@@ -37,13 +37,17 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from ai_review import review_pr
# Pull-request webhook `action` values worth reviewing on. Gitea emits
# GitHub-style payload `action` names (`labeled`, `synchronize`) even though the
# `X-Gitea-Event-Type` header uses `label_updated` / `synchronized` — accept both
# so the gate is robust to either. The label gate below means a non-AI-REVIEW
# label update is a no-op.
REVIEW_ACTIONS = {"opened", "reopened", "synchronize", "synchronized", "labeled", "label_updated"}
# Pull-request webhook `action` values. We fire on EVERY pull_request action
# except `closed` (no point reviewing a closed/merged PR) — the AI-REVIEW label
# gate + sha dedupe downstream make broadening safe: a same-sha re-fire (title
# edit, assignee, milestone, label toggle of another label…) is skipped by
# `review_pr`'s dedupe, and an `unlabeled` event that removed AI-REVIEW fails
# the label gate (payload `labels` reflect current state). Gitea emits
# GitHub-style `action` names (`labeled`, `synchronize`) even though the
# `X-Gitea-Event-Type` header uses `label_updated` / `synchronized`.
SKIP_ACTIONS = {"closed"}
AI_REVIEW_LABEL = "AI-REVIEW"
AI_USAGE_LABEL = "AI-USAGE"
GITEA_API = os.environ.get("GITEA_API", "http://gitea-http.gitea.svc.cluster.local:3000")
BOT_TOKEN = os.environ.get("PRAGENT_BOT_TOKEN", "")
@@ -55,17 +59,23 @@ WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET", "").encode()
PORT = int(os.environ.get("WEBHOOK_PORT", "8080"))
def _labels_have_ai_review(labels) -> bool:
def _labels_have(labels, name: str) -> bool:
"""True if the Gitea PR `labels` list (dicts with `name`, or bare strings)
contains `name`."""
if not isinstance(labels, list):
return False
for lab in labels:
if isinstance(lab, dict) and lab.get("name") == AI_REVIEW_LABEL:
if isinstance(lab, dict) and lab.get("name") == name:
return True
if isinstance(lab, str) and lab == AI_REVIEW_LABEL:
if isinstance(lab, str) and lab == name:
return True
return False
def _labels_have_ai_review(labels) -> bool:
return _labels_have(labels, AI_REVIEW_LABEL)
def _verify_signature(raw_body: bytes, headers) -> bool:
if not WEBHOOK_SECRET:
return False # refuse to run without a configured secret
@@ -87,12 +97,13 @@ def _handle_pull_request(payload: dict) -> tuple[int, str]:
repo_obj = payload.get("repository") or {}
repo = repo_obj.get("full_name") or ""
if action not in REVIEW_ACTIONS:
if action in SKIP_ACTIONS:
return 200, f"ignore action={action}"
if not repo:
return 400, "no repository.full_name"
if not _labels_have_ai_review(pr.get("labels")):
labels = pr.get("labels")
if not _labels_have_ai_review(labels):
return 200, f"ignore (no {AI_REVIEW_LABEL} label) action={action}"
index = pr.get("number")
@@ -106,15 +117,22 @@ def _handle_pull_request(payload: dict) -> tuple[int, str]:
if not BOT_TOKEN:
return 500, "PRAGENT_BOT_TOKEN not set"
# AI-USAGE label (opt-in) → append the token-usage section + per-comment 🪙
# lines to the review. PRAGENT_USAGE_ALWAYS forces it on for testing / a
# future default-on.
report_usage = _labels_have(labels, AI_USAGE_LABEL) or bool(
os.environ.get("PRAGENT_USAGE_ALWAYS")
)
threading.Thread(
target=_run_review,
args=(repo, str(index), title, body, sha),
args=(repo, str(index), title, body, sha, report_usage),
daemon=True,
).start()
return 202, f"reviewing {repo}#{index} action={action} sha={sha[:8]}"
return 202, f"reviewing {repo}#{index} action={action} sha={sha[:8]} usage={report_usage}"
def _run_review(repo: str, index: str, title: str, body: str, sha: str) -> None:
def _run_review(repo: str, index: str, title: str, body: str, sha: str, report_usage: bool) -> None:
try:
ok = review_pr(
api=GITEA_API,
@@ -128,8 +146,9 @@ def _run_review(repo: str, index: str, title: str, body: str, sha: str) -> None:
model=OLLAMA_MODEL,
max_tokens=OLLAMA_MAX_TOKENS,
max_chars=DIFF_MAX_CHARS,
report_usage=report_usage,
)
print(f"pragent-webhook: reviewed {repo}#{index} sha={sha[:8]} ok={ok}", flush=True)
print(f"pragent-webhook: reviewed {repo}#{index} sha={sha[:8]} ok={ok} usage={report_usage}", flush=True)
except Exception as e: # review_pr is fail-open, but guard the thread anyway
print(f"pragent-webhook: thread crashed for {repo}#{index}: {e}", flush=True)