diff --git a/pilot/ai_review.py b/pilot/ai_review.py index fec53ea..9dfee33 100644 --- a/pilot/ai_review.py +++ b/pilot/ai_review.py @@ -66,6 +66,8 @@ SHA_MARKER = "" _SHA_MARKER_RE = re.compile(r"") 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 diff --git a/tests/pilot/test_ai_review.py b/tests/pilot/test_ai_review.py index 2ab4d08..591bb5e 100644 --- a/tests/pilot/test_ai_review.py +++ b/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