pilot/ai_review.py

@@ -66,6 +66,8 @@ SHA_MARKER = "<!-- pragent:sha={sha} -->"
 _SHA_MARKER_RE = re.compile(r"<!-- pragent:sha=([0-9a-f]{7,40}) -->")
 
 AI_REVIEW_LABEL = "AI-REVIEW"
+# Opt-in label for the token-usage block. Read at render time — see pr_has_label.
+AI_USAGE_LABEL = "AI-USAGE"
 SEVERITIES = ("critical", "high", "medium", "low")
 # Severity rank — higher = more severe. Used by `apply_repo_config` to drop
 # findings below `severity_threshold`. Critical=3, high=2, medium=1, low=0.
@@ -156,6 +158,33 @@ def parse_text_blocks(content: list) -> str:
     return "\n".join(out).strip()
 
 
+def pr_has_label(api: str, repo: str, index: str, token: str, label: str) -> bool:
+    """True if the PR currently carries `label`. False on any failure.
+
+    Read at RENDER time, not at review start. A reviewer labels AI-REVIEW and
+    AI-USAGE seconds apart; the review claims on the first event and the
+    second is dropped by the in-flight dedupe, so the trigger payload never
+    sees the opt-in. Re-reading when the review begins is no better — that is
+    still milliseconds after the first click. Only a read taken once the
+    review has finished (a minute or more later) reliably sees the label.
+    """
+    try:
+        code, raw = gitea_get(api, repo, f"issues/{index}/labels", token)
+        if code >= 300:
+            return False
+        data = json.loads(raw.decode() or "[]")
+    except Exception as e:
+        print(f"pragent: could not re-read labels for {repo}#{index}: {e}",
+              file=sys.stderr, flush=True)
+        return False
+    if not isinstance(data, list):
+        return False
+    return any(
+        (isinstance(x, dict) and x.get("name") == label) or x == label
+        for x in data
+    )
+
+
 def _int_env(name: str, default: int) -> int:
     """Read an int from the environment, falling back on anything unparseable.
 
@@ -1896,6 +1925,13 @@ def review_pr(
                     file=sys.stderr, flush=True,
                 )
                 salvaged = salvage_summary(stdout)
+                # The AI-USAGE opt-in is re-checked HERE, at render time: the label
+                # is usually applied moments after AI-REVIEW, long after this
+                # review was claimed and its trigger payload frozen.
+                if not report_usage:
+                    report_usage = pr_has_label(api, repo, index, token, AI_USAGE_LABEL)
+                    if report_usage and usage and usage.get('output'):
+                        compute_attribution(findings, usage['output'])
                 usage_section = _render_collapsible_usage(usage, model, config=config) if report_usage else ""
                 post_review(api, repo, index, token, format_review_body(
                     salvaged or "AI review produced no parseable output.",
@@ -1936,6 +1972,13 @@ def review_pr(
         # for it.
         if report_usage and usage and usage.get("output"):
             compute_attribution(findings, usage["output"])
+        # The AI-USAGE opt-in is re-checked HERE, at render time: the label
+        # is usually applied moments after AI-REVIEW, long after this
+        # review was claimed and its trigger payload frozen.
+        if not report_usage:
+            report_usage = pr_has_label(api, repo, index, token, AI_USAGE_LABEL)
+            if report_usage and usage and usage.get('output'):
+                compute_attribution(findings, usage['output'])
         usage_section = _render_collapsible_usage(usage, model, config=config) if report_usage else ""
 
         # Anchor against the RAW diff, never the compressed one. Compression

pilot/webhook_server.py

@@ -39,7 +39,6 @@ import hmac
 import json
 import os
 import threading
-import urllib.request
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 
 from ai_review import review_pr
@@ -177,48 +176,10 @@ def _release(key: tuple[str, str, str]) -> None:
         _inflight.discard(key)
 
 
-def _fetch_current_labels(repo: str, index: str) -> list:
-    """Re-read the PR's labels from the API. [] on any failure.
-
-    The webhook payload is a snapshot from the instant the event fired. When a
-    PR is labelled AI-REVIEW first and AI-USAGE a moment later, the review is
-    already claimed and running, so the second event is deduped and the usage
-    opt-in is lost — the review posts without its usage block. Reading the
-    labels again at review start closes that window.
-    """
-    # GITEA_API is the bare host (no /api/v1) — every caller appends the
-    # version prefix itself; see ai_review._get / _post.
-    url = f"{GITEA_API}/api/v1/repos/{repo}/issues/{index}/labels"
-    req = urllib.request.Request(url, headers={
-        "Authorization": f"token {BOT_TOKEN}",
-        "Accept": "application/json",
-    })
-    try:
-        with urllib.request.urlopen(req, timeout=15) as resp:
-            return json.loads(resp.read().decode() or "[]")
-    except Exception as e:
-        print(f"pragent-webhook: could not re-read labels for {repo}#{index}: {e}", flush=True)
-        return []
-
-
 def _run_review(
     key: tuple[str, str, str], title: str, body: str, report_usage: bool, base_ref: str
 ) -> None:
     repo, index, sha = key
-    # The AI-USAGE opt-in may have been applied *after* the label event that
-    # triggered this review (labelling is two events, the review claims on the
-    # first). Re-read the labels now so the later opt-in still counts.
-    if not report_usage:
-        report_usage = _labels_have(
-            _fetch_current_labels(repo, index), AI_USAGE_LABEL
-        )
-        if report_usage:
-            print(
-                f"pragent-webhook: {repo}#{index} AI-USAGE found on re-read "
-                f"(added after the trigger event)",
-                flush=True,
-            )
-
     try:
         with _review_slots:
             ok = review_pr(

tests/pilot/test_ai_review.py

@@ -616,6 +616,35 @@ def test_reference_non_url_renders_as_plain_text():
     assert "](CVE-" not in body
 
 
+def test_pr_has_label_reads_the_live_labels(monkeypatch):
+    # The AI-USAGE opt-in is read at render time, not from the trigger
+    # payload: labelling AI-REVIEW then AI-USAGE is two events, the review
+    # claims on the first, and the second is dropped by the in-flight dedupe.
+    seen = {}
+
+    def _get(api, repo, path, token, accept="application/json"):
+        seen["path"] = path
+        return 200, b'[{"name": "AI-REVIEW"}, {"name": "AI-USAGE"}]'
+
+    monkeypatch.setattr(ai_review, "gitea_get", _get)
+    assert ai_review.pr_has_label("http://api", "o/r", "9", "t", "AI-USAGE") is True
+    assert seen["path"] == "issues/9/labels"
+    assert ai_review.pr_has_label("http://api", "o/r", "9", "t", "NOPE") is False
+
+
+def test_pr_has_label_survives_a_broken_api(monkeypatch):
+    def _boom(*a, **k):
+        raise RuntimeError("gitea down")
+    monkeypatch.setattr(ai_review, "gitea_get", _boom)
+    assert ai_review.pr_has_label("http://api", "o/r", "9", "t", "AI-USAGE") is False
+
+    monkeypatch.setattr(ai_review, "gitea_get", lambda *a, **k: (404, b"nope"))
+    assert ai_review.pr_has_label("http://api", "o/r", "9", "t", "AI-USAGE") is False
+
+    monkeypatch.setattr(ai_review, "gitea_get", lambda *a, **k: (200, b'{"not": "a list"}'))
+    assert ai_review.pr_has_label("http://api", "o/r", "9", "t", "AI-USAGE") is False
+
+
 def test_int_env_falls_back_on_garbage(monkeypatch, capsys):
     monkeypatch.setenv("PRAGENT_DIFF_CONTEXT", "two")
     assert ai_review._int_env("PRAGENT_DIFF_CONTEXT", 1) == 1

tests/pilot/test_webhook_server.py

@@ -134,64 +134,3 @@ def test_missing_label_is_ignored(monkeypatch):
 def test_review_slots_bound_matches_config():
     assert ws.MAX_CONCURRENT >= 1
     assert ws._review_slots._value <= ws.MAX_CONCURRENT
-
-
-# ---------------------------------------------------------------------------
-# AI-USAGE opt-in applied after the triggering label event
-# ---------------------------------------------------------------------------
-
-
-def test_run_review_rereads_labels_for_late_ai_usage(monkeypatch):
-    # Labelling a PR AI-REVIEW then AI-USAGE fires two events; the review
-    # claims on the first, so the second is deduped and the payload snapshot
-    # never sees AI-USAGE. The re-read at review start must recover it.
-    seen = {}
-    monkeypatch.setattr(ws, "_fetch_current_labels",
-                        lambda repo, index: [{"name": "AI-REVIEW"}, {"name": "AI-USAGE"}])
-    monkeypatch.setattr(ws, "review_pr", lambda **kw: seen.update(kw) or True)
-    ws._run_review(("o/r", "9", "abc1234"), "t", "b", False, "main")
-    assert seen["report_usage"] is True
-
-
-def test_run_review_keeps_usage_off_when_label_absent(monkeypatch):
-    seen = {}
-    monkeypatch.setattr(ws, "_fetch_current_labels",
-                        lambda repo, index: [{"name": "AI-REVIEW"}])
-    monkeypatch.setattr(ws, "review_pr", lambda **kw: seen.update(kw) or True)
-    ws._run_review(("o/r", "9", "abc1234"), "t", "b", False, "main")
-    assert seen["report_usage"] is False
-
-
-def test_fetch_current_labels_hits_the_versioned_api_path(monkeypatch):
-    # GITEA_API is the bare host; the /api/v1 prefix is the caller's job.
-    # Getting this wrong 404s silently and the opt-in is lost — which is
-    # exactly what shipped the first time.
-    seen = {}
-
-    class _Resp:
-        def read(self): return b'[{"name": "AI-USAGE"}]'
-        def __enter__(self): return self
-        def __exit__(self, *a): return False
-
-    def _urlopen(req, timeout=None):
-        seen["url"] = req.full_url
-        return _Resp()
-
-    monkeypatch.setattr(ws.urllib.request, "urlopen", _urlopen)
-    out = ws._fetch_current_labels("o/r", "9")
-    assert out == [{"name": "AI-USAGE"}]
-    assert seen["url"] == f"{ws.GITEA_API}/api/v1/repos/o/r/issues/9/labels"
-
-
-def test_label_reread_failure_is_not_fatal(monkeypatch):
-    # A dead API must not take the review down with it — the re-read is a
-    # best-effort upgrade of an opt-in flag, nothing more.
-    def _boom(*a, **k):
-        raise OSError("api down")
-    monkeypatch.setattr(ws.urllib.request, "urlopen", _boom)
-    assert ws._fetch_current_labels("o/r", "9") == []
-
-    seen = {}
-    monkeypatch.setattr(ws, "review_pr", lambda **kw: seen.update(kw) or True)
-    ws._run_review(("o/r", "9", "abc1234"), "t", "b", False, "main")
-    assert seen["report_usage"] is False