diff --git a/.github/workflows/ai_pr_review.yml b/.github/workflows/ai_pr_review.yml index 3a79ae92..af1b7c76 100644 --- a/.github/workflows/ai_pr_review.yml +++ b/.github/workflows/ai_pr_review.yml @@ -269,6 +269,16 @@ jobs: echo "Base SHA does not contain tools/notebook_md_extract.py; tutorial prose extraction will be skipped (one-shot bootstrap)." fi + # Stage the diff budget filter from BASE_SHA for the same reason. + # Bootstrap-safe: absent on base, the unified diff is appended + # untrimmed (the pre-filter behaviour). + if git show "${BASE_SHA}":tools/ai_review_diff_budget.py > /tmp/ai_review_diff_budget.py 2>/dev/null; then + echo "Diff budget filter staged from base SHA ${BASE_SHA}." + else + rm -f /tmp/ai_review_diff_budget.py + echo "Base SHA does not contain tools/ai_review_diff_budget.py; the unified diff will be appended untrimmed (one-shot bootstrap)." + fi + # Sanitize untrusted text so hostile content can't close the # wrapper tags and inject instructions to the reviewer. # Case- and whitespace-tolerant; PR_TITLE / PR_BODY / PREV_REVIEW @@ -338,19 +348,10 @@ jobs: echo "" echo "Changed files:" git --no-pager diff --name-status "$BASE_SHA" "$HEAD_SHA" - echo "" - echo "Unified diff (context=5):" - # Exclude large generated/data files from the full diff to stay - # within the model's input limit. The --name-status above still - # lists them. Narrowed to real-data assets, notebook outputs, and - # the qte golden fixture (generated single-line JSON, ~700KB - a - # regeneration blew the input limit at PR #688; the reviewable - # artifact is benchmarks/R/generate_qte_golden.R). - git --no-pager diff --unified=5 "$BASE_SHA" "$HEAD_SHA" \ - -- . ':!benchmarks/data/real/*.json' ':!benchmarks/data/real/*.csv' \ - ':!benchmarks/data/qte_golden.json' \ - ':!docs/tutorials/*.ipynb' } >> "$PROMPT" + # The unified diff is appended LAST (after the notebook prose + # below) so its budget can be computed from the final size of + # everything else in the prompt. # Tutorial notebook prose extraction: substitute a markdown view # for the .ipynb JSON that the unified diff excludes. The extractor @@ -512,6 +513,42 @@ jobs: } >> "$PROMPT" fi + # Unified diff, fitted to the remaining prompt budget. + # Exclude large generated/data files from the full diff to stay + # within the model's input limit. The --name-status above still + # lists them. Narrowed to real-data assets, notebook outputs, and + # the qte golden fixture (generated single-line JSON, ~700KB - a + # regeneration blew the input limit at PR #688; the reviewable + # artifact is benchmarks/R/generate_qte_golden.R). + git --no-pager diff --unified=5 "$BASE_SHA" "$HEAD_SHA" \ + -- . ':!benchmarks/data/real/*.json' ':!benchmarks/data/real/*.csv' \ + ':!benchmarks/data/qte_golden.json' \ + ':!docs/tutorials/*.ipynb' > /tmp/ai-review-diff.patch + # Codex rejects a prompt above 1,048,576 characters outright + # (input_too_large; PR #861 compiled to 1.72M). PROMPT_BUDGET is the + # whole-prompt target in BYTES (a bound on characters for UTF-8), + # leaving ~48K headroom under the hard cap; the diff gets whatever + # the rules, PR text, previous review, file list and notebook prose + # left. The filter (staged from BASE_SHA above) passes an in-budget + # patch through byte-for-byte and otherwise replaces whole per-file + # blocks with one-line stubs, largest and lowest-value paths first, + # each carrying the git command that shows that file's hunks. + PROMPT_BUDGET=1000000 + { + echo "" + echo "Unified diff (context=5):" + } >> "$PROMPT" + if [ -f /tmp/ai_review_diff_budget.py ]; then + DIFF_BUDGET=$(( PROMPT_BUDGET - $(wc -c < "$PROMPT") )) + if [ "$DIFF_BUDGET" -lt 0 ]; then DIFF_BUDGET=0; fi + python3 /tmp/ai_review_diff_budget.py --diff /tmp/ai-review-diff.patch \ + --budget "$DIFF_BUDGET" --base-sha "$BASE_SHA" --head-sha "$HEAD_SHA" >> "$PROMPT" + else + cat /tmp/ai-review-diff.patch >> "$PROMPT" + fi + rm -f /tmp/ai-review-diff.patch + echo "Compiled prompt size: $(wc -c < "$PROMPT") bytes (target ${PROMPT_BUDGET})." + - name: Run Codex id: run_codex if: steps.pr.outputs.is_fork == 'false' diff --git a/TODO.md b/TODO.md index 6f36089a..e84743fe 100644 --- a/TODO.md +++ b/TODO.md @@ -86,6 +86,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| +| Local codex review (`openai_review.py --backend codex`) has no diff budget - an oversized PR needs hand-trimmed `:(exclude)` pathspecs (#861 compiled to 1.35 MB locally); reuse `tools/ai_review_diff_budget.py` on the codex-backend patch the way the CI workflow does | `.claude/scripts/openai_review.py` | #861 | Quick | Low | | Re-validate the plan-review dual engine on `gpt-6.1-sol` with `tools/plan-review-eval`: re-pin every arm's composite `cli_version` to the live Claude Code + codex CLI pair, add a gpt-6.1-sol arm, run the harness against the gpt-5.6-sol baseline (the 2026-10 swap carried the model over from the code-review campaign 2 verdict without a plan-review campaign). | `tools/plan-review-eval/config/configs.json` | reviewer-eval gpt61sol PR-2 | Heavy | Low | | TWFE-weights tutorial section: `attgt_weights` (all three `type=` values side by side) + `decompose_twfe_weights` + `plot_twfe_weights` on `mpdta`, either as a section of `02_staggered_did.ipynb` or a new notebook registered in `docs/tutorials/index.rst`; follow-on PR by decision, the feature ships without it | `docs/tutorials/02_staggered_did.ipynb` | #812 | Mid | Medium | | Geo-experiments tutorial: the SDiD fits on the `generate_factor_data(...)` panel (conversions ~1,500, sigma-hat ~310) hit the Frank-Wolfe iteration cap - the placebo fit warns once and the bootstrap cross-check warns on 100/100 draws - yet the narrative presents the bootstrap SE as a paper-faithful robustness check. Pre-existing on `main` under BOTH backends (surfaced when the notebook outputs were refreshed for the shape-only fit diagnostic); the committed April outputs predate the aggregated non-convergence warning. Either rescale / re-tune the DGP (or raise `max_iter` / loosen `min_decrease`) so the refits converge and re-execute, or label the bootstrap result unreliable and drop it from the "Validity evidence" claim. | `docs/tutorials/18_geo_experiments.ipynb` | local codex review of the SDiD shape-only fit-RMSE PR | Mid | Low | diff --git a/changelog.d/20261019-ai-review-diff-budget.md b/changelog.d/20261019-ai-review-diff-budget.md new file mode 100644 index 00000000..a4ea8e7f --- /dev/null +++ b/changelog.d/20261019-ai-review-diff-budget.md @@ -0,0 +1,13 @@ +### Internal +- **CI AI review fits oversized diffs into the codex input cap.** The compiled review + prompt of PR #861 reached 1.72M characters against codex's hard 1,048,576-character + limit (`input_too_large`), so the review never ran; the workflow capped notebook prose + but never the unified diff. `.github/workflows/ai_pr_review.yml` now appends the diff + last, under a whole-prompt byte budget (1,000,000) derived from everything else already + in the prompt, through the new `tools/ai_review_diff_budget.py` (staged from `BASE_SHA` + like the notebook extractor; untrimmed one-shot fallback when absent on base). An + in-budget patch passes through byte-for-byte; an oversized one has whole per-file blocks + replaced by one-line stubs - `tests/` / `benchmarks/` / `tools/` first, `diff_diff/` / + `rust/` only when unavoidable, largest first - each naming the `git diff` command the + reviewer can run for that file. Pinned by `tests/test_ai_review_diff_budget.py`; the + staging/ordering contract joins `TestWorkflowDoesNotExecutePRHeadCode`. diff --git a/tests/test_ai_review_diff_budget.py b/tests/test_ai_review_diff_budget.py new file mode 100644 index 00000000..0836f3e1 --- /dev/null +++ b/tests/test_ai_review_diff_budget.py @@ -0,0 +1,365 @@ +"""Tests for `tools/ai_review_diff_budget.py`. + +`.github/workflows/ai_pr_review.yml` runs this filter (staged from BASE_SHA) +over the unified diff so the compiled Codex prompt stays under the +1,048,576-character input cap that killed the review of PR #861 (1.72M +compiled). Inline fixtures only; the skip-guard covers `rust-test.yml`'s +isolated-install matrix, which copies just `tests/` and has no `tools/`. +""" + +from __future__ import annotations + +import importlib.util +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT_PATH = REPO_ROOT / "tools" / "ai_review_diff_budget.py" + +if not SCRIPT_PATH.exists(): + pytest.skip( + "tools/ai_review_diff_budget.py not present in working tree", + allow_module_level=True, + ) + + +def _load(): + spec = importlib.util.spec_from_file_location("ai_review_diff_budget", SCRIPT_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _file_block(path: str, body_lines: int, *, old_path: str | None = None) -> bytes: + """A realistic per-file unified-diff block of roughly ``body_lines`` lines.""" + a = old_path or path + head = ( + f"diff --git a/{a} b/{path}\n" + f"index 0123456..89abcde 100644\n" + f"--- a/{a}\n" + f"+++ b/{path}\n" + f"@@ -1,{body_lines} +1,{body_lines} @@\n" + ) + body = "".join(f"+line {i} of {path}\n" for i in range(body_lines)) + return (head + body).encode() + + +@pytest.fixture(scope="module") +def mod(): + return _load() + + +class TestSplitBlocks: + def test_splits_on_column_zero_header_only(self, mod): + # A body line that CONTAINS "diff --git" carries a prefix char and must + # not start a new block. + trap = b"+diff --git a/trap b/trap\n diff --git a/trap2 b/trap2\n" + patch = _file_block("diff_diff/a.py", 2) + trap + _file_block("tests/b.py", 1) + preamble, blocks = mod.split_blocks(patch) + assert preamble == b"" + assert [b.path for b in blocks] == ["diff_diff/a.py", "tests/b.py"] + assert trap in blocks[0].data + assert b"".join(b.data for b in blocks) == patch + + def test_preamble_kept(self, mod): + patch = b"some preamble\n" + _file_block("x.py", 1) + preamble, blocks = mod.split_blocks(patch) + assert preamble == b"some preamble\n" + assert len(blocks) == 1 + + def test_deleted_file_uses_a_side_path(self, mod): + block = ( + b"diff --git a/gone.py b/gone.py\n" + b"deleted file mode 100644\n" + b"--- a/gone.py\n" + b"+++ /dev/null\n" + b"@@ -1 +0,0 @@\n-x\n" + ) + _, blocks = mod.split_blocks(block) + assert blocks[0].path == "gone.py" + + def test_binary_block_path_comes_from_the_header(self, mod): + block = b"diff --git a/img.png b/img.png\nBinary files a/img.png and b/img.png differ\n" + _, blocks = mod.split_blocks(block) + assert blocks[0].path == "img.png" + + def test_mode_only_block_with_space_b_slash_in_path(self, mod): + block = b"diff --git a/x b/y.py b/x b/y.py\nold mode 100644\nnew mode 100755\n" + _, blocks = mod.split_blocks(block) + assert blocks[0].path == "x b/y.py" + + def test_pure_rename_uses_rename_to_and_keeps_source_tier(self, mod): + block = ( + b"diff --git a/diff_diff/old.py b/diff_diff/new.py\n" + b"similarity index 100%\nrename from diff_diff/old.py\nrename to diff_diff/new.py\n" + ) + _, blocks = mod.split_blocks(block) + assert blocks[0].path == "diff_diff/new.py" + assert mod.tier(blocks[0].path) == 2 + + def test_c_quoted_paths_are_unquoted(self, mod): + new_file = ( + b'diff --git "a/caf\\303\\251 \\"x\\".py" "b/caf\\303\\251 \\"x\\".py"\n' + b"new file mode 100644\n--- /dev/null\n" + b'+++ "b/caf\\303\\251 \\"x\\".py"\n@@ -0,0 +1 @@\n+x\n' + ) + binary = b'diff --git "a/tests/caf\\303\\251.png" "b/tests/caf\\303\\251.png"\nBinary files differ\n' + _, blocks = mod.split_blocks(new_file + binary) + assert [b.path for b in blocks] == ['caf\u00e9 "x".py', "tests/caf\u00e9.png"] + assert mod.tier(blocks[1].path) == 0 + + def test_quoted_header_with_escaped_quote_then_space_in_path(self, mod): + # Path `diff_diff/end" x.py`: the a-side token contains `\" ` so a + # naive split on `" "` picks the wrong boundary. + mode_only = ( + b'diff --git "a/diff_diff/end\\" x.py" "b/diff_diff/end\\" x.py"\n' + b"old mode 100644\nnew mode 100755\n" + ) + binary = b'diff --git "a/diff_diff/end\\" x.png" "b/diff_diff/end\\" x.png"\nBinary files differ\n' + _, blocks = mod.split_blocks(mode_only + binary) + assert [b.path for b in blocks] == ['diff_diff/end" x.py', 'diff_diff/end" x.png'] + assert [mod.tier(b.path) for b in blocks] == [2, 2] + out = mod.stub(blocks[0], "a" * 40, "b" * 40) + assert b"""-- ':(literal)diff_diff/end" x.py'""" in out + + def test_c_unquote_escapes_and_octal(self, mod): + assert mod._c_unquote(b'"a\\tb\\\\c\\"d\\303\\251"') == b'a\tb\\c"d\xc3\xa9' + assert mod._c_unquote(b"plain") == b"plain" + + def test_empty_patch(self, mod): + assert mod.split_blocks(b"") == (b"", []) + + +class TestTiering: + @pytest.mark.parametrize( + "path, expected", + [ + ("tests/test_x.py", 0), + ("tests/fixtures/big.json", 0), + ("benchmarks/python/bench.py", 0), + ("tools/foo.py", 0), + (".claude/scripts/x.py", 0), + ("docs/methodology/REGISTRY.md", 1), + ("docs/v4-deprecations.yaml", 1), + (".github/workflows/ci.yml", 1), + ("README.md", 1), + ("diff_diff/staggered.py", 2), + ("rust/src/lib.rs", 2), + ], + ) + def test_tier(self, mod, path, expected): + assert mod.tier(path) == expected + + +class TestFitToBudget: + def test_in_budget_passthrough_is_byte_identical(self, mod): + patch = _file_block("diff_diff/a.py", 50) + _file_block("tests/b.py", 50) + assert mod.fit_to_budget(patch, len(patch)) == patch + assert mod.fit_to_budget(patch, len(patch) + 1) == patch + + def test_drops_lowest_tier_largest_first(self, mod): + src = _file_block("diff_diff/core.py", 400) # tier 2, biggest + doc = _file_block("docs/methodology/REGISTRY.md", 300) # tier 1 + big_test = _file_block("tests/test_big.py", 200) # tier 0, bigger + small_test = _file_block("tests/test_small.py", 20) # tier 0, smaller + patch = src + doc + big_test + small_test + budget = len(patch) - len(big_test) // 2 # dropping big_test alone suffices + out = mod.fit_to_budget(patch, budget, "aaaaaaa", "bbbbbbb") + assert len(out) <= budget + assert src in out and doc in out and small_test in out + assert big_test not in out + assert b"[diff omitted for prompt budget] tests/test_big.py (" in out + assert ( + b"git --no-pager diff --unified=5 aaaaaaa bbbbbbb -- ':(literal)tests/test_big.py'" + in out + ) + assert out.startswith(b"[prompt budget] 1 of 4 file diffs are omitted") + + def test_source_dropped_only_when_unavoidable(self, mod): + src = _file_block("diff_diff/core.py", 400) + test = _file_block("tests/test_big.py", 200) + patch = src + test + # Budget smaller than src alone: both must go. + out = mod.fit_to_budget(patch, len(src) - 1, "aaaaaaa", "bbbbbbb") + assert len(out) <= len(src) - 1 + assert b"tests/test_big.py (" in out and b"diff_diff/core.py (" in out + # Budget that fits src + stub but not src + test: only the test goes. + out2 = mod.fit_to_budget(patch, len(src) + 600, "aaaaaaa", "bbbbbbb") + assert src in out2 and test not in out2 + + def test_order_of_kept_and_stubbed_blocks_is_preserved(self, mod): + blocks = [ + _file_block("diff_diff/a.py", 10), + _file_block("tests/t1.py", 300), + _file_block("diff_diff/b.py", 10), + ] + patch = b"".join(blocks) + out = mod.fit_to_budget(patch, len(patch) - 10) + ia = out.index(b"diff --git a/diff_diff/a.py") + it = out.index(b"[diff omitted for prompt budget] tests/t1.py") + ib = out.index(b"diff --git a/diff_diff/b.py") + assert ia < it < ib + + def test_tiny_budget_emits_bounded_summary_only(self, mod): + patch = _file_block("diff_diff/a.py", 10) + _file_block("tests/t.py", 10) + out = mod.fit_to_budget(patch, 300) + assert len(out) <= 300 + assert out.startswith(b"[prompt budget] the unified diff (") + assert b"does not fit the remaining budget of 300 bytes" in out + assert b"diff --git" not in out + # The summary itself is truncated to the budget, down to empty. + for budget in (0, 1, 17, 60): + short = mod.fit_to_budget(patch, budget) + assert len(short) <= budget + assert short.startswith(b"[prompt budget] the unified diff ("[:budget]) + assert mod.fit_to_budget(patch, 0) == b"" + + def test_no_headers_hard_truncates_within_budget(self, mod): + patch = b"x" * 1000 + out = mod.fit_to_budget(patch, 100) + assert len(out) <= 100 + assert out.endswith(b"[prompt budget] diff truncated here.\n") + for budget in (0, 1, 10, 37, 38, 39): + assert len(mod.fit_to_budget(patch, budget)) <= budget + + def test_non_shrinking_stubs_are_never_substituted(self, mod): + """A stub larger than the hunks it replaces frees nothing; ten tiny test + diffs must stay visible while the one large source diff is stubbed.""" + small = [_file_block(f"tests/t{i}.py", 1) for i in range(10)] + big = _file_block("diff_diff/core.py", 60) + patch = b"".join(small) + big + stub_len = len(mod.stub(mod.split_blocks(small[0])[1][0], "a" * 40, "b" * 40)) + assert stub_len > len(small[0]) + budget = len(patch) - len(big) // 2 + out = mod.fit_to_budget(patch, budget, "a" * 40, "b" * 40) + assert len(out) <= budget + for blk in small: + assert blk in out + assert big not in out + assert b"[diff omitted for prompt budget] diff_diff/core.py (" in out + + def test_stub_without_shas_has_no_git_hint(self, mod): + patch = _file_block("tests/t.py", 100) + _file_block("diff_diff/a.py", 1) + out = mod.fit_to_budget(patch, len(patch) - 1) + assert b"[diff omitted for prompt budget] tests/t.py (" in out + assert b"git --no-pager" not in out + + def test_path_with_quote_is_shell_quoted_in_hint(self, mod): + patch = _file_block("tests/it's.py", 100) + _file_block("diff_diff/a.py", 1) + out = mod.fit_to_budget(patch, len(patch) - 1, "aaaaaaa", "bbbbbbb") + assert b"-- ':(literal)tests/it'\\''s.py'" in out + + def test_output_never_exceeds_budget_across_sizes(self, mod): + blocks = [_file_block(f"tests/t{i}.py", 5 + 37 * i) for i in range(12)] + blocks += [_file_block(f"diff_diff/m{i}.py", 50 + 11 * i) for i in range(6)] + blocks += [_file_block(f"docs/d{i}.md", 30 * i) for i in range(5)] + patch = b"".join(blocks) + for budget in list(range(0, 400)) + list(range(400, len(patch) + 1, 997)): + out = mod.fit_to_budget(patch, budget, "aaaaaaa", "bbbbbbb") + assert len(out) <= budget, budget + assert mod.fit_to_budget(patch, len(patch)) == patch + + +class TestCli: + def _run(self, args, cwd): + return subprocess.run( + [sys.executable, "-I", str(SCRIPT_PATH), *args], + cwd=cwd, + capture_output=True, + check=False, + ) + + def test_cli_passthrough_and_output_file(self, tmp_path): + patch = _file_block("diff_diff/a.py", 3) + src = tmp_path / "d.patch" + src.write_bytes(patch) + out = tmp_path / "o.patch" + r = self._run(["--diff", str(src), "--budget", "99999", "-o", str(out)], tmp_path) + assert r.returncode == 0, r.stderr + assert out.read_bytes() == patch + r2 = self._run(["--diff", str(src), "--budget", "99999"], tmp_path) + assert r2.returncode == 0 and r2.stdout == patch + + def test_cli_trims_with_shas(self, tmp_path): + patch = _file_block("tests/t.py", 200) + _file_block("diff_diff/a.py", 2) + src = tmp_path / "d.patch" + src.write_bytes(patch) + r = self._run( + [ + "--diff", + str(src), + "--budget", + str(len(patch) - 1), + "--base-sha", + "0" * 40, + "--head-sha", + "f" * 40, + ], + tmp_path, + ) + assert r.returncode == 0, r.stderr + assert len(r.stdout) <= len(patch) - 1 + assert b"[diff omitted for prompt budget] tests/t.py (" in r.stdout + assert b"diff --git a/diff_diff/a.py" in r.stdout + + @pytest.mark.parametrize("bad", ["main", "$(id)", "abc;rm", "ZZZZZZZ"]) + def test_cli_rejects_non_hex_sha(self, tmp_path, bad): + src = tmp_path / "d.patch" + src.write_bytes(_file_block("x.py", 1)) + r = self._run(["--diff", str(src), "--budget", "10", "--base-sha", bad], tmp_path) + assert r.returncode == 2 + assert b"not a hex object id" in r.stderr + + def test_cli_rejects_negative_budget(self, tmp_path): + src = tmp_path / "d.patch" + src.write_bytes(_file_block("x.py", 1)) + r = self._run(["--diff", str(src), "--budget", "-1"], tmp_path) + assert r.returncode == 2 + + +class TestWorkflowWiring: + """The workflow must stage the filter from BASE_SHA and append the + unified diff LAST under a budget derived from the compiled prompt size + (so notebook prose, PR text and the previous review are all counted). + The staging/exec ordering contract itself is enforced by + `tests/test_openai_review.py::TestWorkflowDoesNotExecutePRHeadCode`.""" + + WF = REPO_ROOT / ".github" / "workflows" / "ai_pr_review.yml" + + @pytest.fixture + def text(self): + if not self.WF.exists(): + pytest.skip("workflow not found") + return self.WF.read_text() + + def test_filter_is_staged_from_base_sha(self, text): + assert ( + 'git show "${BASE_SHA}":tools/ai_review_diff_budget.py > /tmp/ai_review_diff_budget.py' + in text + ) + + def test_budget_is_derived_from_the_compiled_prompt(self, text): + live = [ln for ln in text.splitlines() if not ln.lstrip().startswith("#")] + assert any(ln.strip() == "PROMPT_BUDGET=1000000" for ln in live), ( + "PROMPT_BUDGET must stay a whole-prompt byte target below the " + "1,048,576-character codex cap" + ) + assert any('DIFF_BUDGET=$(( PROMPT_BUDGET - $(wc -c < "$PROMPT") ))' in ln for ln in live) + assert any( + "python3 /tmp/ai_review_diff_budget.py --diff /tmp/ai-review-diff.patch" in ln + for ln in live + ) + + def test_unified_diff_is_appended_after_notebook_prose(self, text): + # The diff must come after the LAST prose wrapper so the budget sees it. + prose_end = text.rfind('echo ""') + diff_header = text.rfind('echo "Unified diff (context=5):"') + assert 0 < prose_end < diff_header + assert text.find("- name: Run Codex") > diff_header + + def test_bootstrap_fallback_appends_untrimmed(self, text): + assert 'cat /tmp/ai-review-diff.patch >> "$PROMPT"' in text diff --git a/tests/test_openai_review.py b/tests/test_openai_review.py index 9e17dee7..f3bc036a 100644 --- a/tests/test_openai_review.py +++ b/tests/test_openai_review.py @@ -3870,6 +3870,7 @@ class TestWorkflowDoesNotExecutePRHeadCode: # source so the staging-test asserts the FULL command line. ALLOWED_TMP_PYTHON_EXECUTIONS = { "/tmp/notebook_md_extract.py": "tools/notebook_md_extract.py", + "/tmp/ai_review_diff_budget.py": "tools/ai_review_diff_budget.py", } # Literal allowlist of single-line `python3 -c ` payloads. diff --git a/tools/ai_review_diff_budget.py b/tools/ai_review_diff_budget.py new file mode 100644 index 00000000..b836f182 --- /dev/null +++ b/tools/ai_review_diff_budget.py @@ -0,0 +1,326 @@ +"""Fit a unified diff into the AI reviewer's input budget. + +Used by `.github/workflows/ai_pr_review.yml`. Codex rejects a prompt longer +than 1,048,576 characters outright (`input_too_large`), and a large PR's +unified diff alone can exceed that (PR #861: 1.63 MB of diff after the +notebook exclusion). This filter reads the patch `git diff --unified=5` +produced, and when it is longer than the remaining budget replaces whole +per-file blocks with one-line stubs - largest files first, lowest-value +paths first - until the output fits. Every omitted file stays listed in the +prompt's `Changed files:` section, and each stub carries the exact +`git diff` command the (agentic, repo-checked-out) reviewer can run to see +that file's hunks on demand. + +Under budget the patch passes through byte-for-byte, so the prompt of every +PR that fitted before this filter existed is unchanged. + +Splitting is on `diff --git ` at column 0: in unified-diff output every body +line carries a one-character prefix (` `, `+`, `-`, `\\`), so a column-0 +`diff --git ` line is always a file header, whatever the file contains. + +Drop order (deterministic, greedy): + tier 0 - tests/, benchmarks/, tools/, .claude/ (largest first) + tier 1 - everything else (largest first) + tier 2 - diff_diff/, rust/ (largest first; the + reviewer's primary subject, dropped only when unavoidable) + +All sizes are BYTES: the cap counts characters, and for UTF-8 text the byte +count is an upper bound on both code points and UTF-16 units, so fitting in +bytes always fits in characters. stdlib only - the workflow has no +`pip install` step - and no shell is ever invoked (the patch is read from a +file; SHAs are validated hex and appear only as text in the stubs). +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path +from typing import List, NamedTuple, Optional, Sequence + +HEADER = b"diff --git " +_SHA_RE = re.compile(r"^[0-9a-f]{7,64}$") + +TIER_FIRST = ("tests/", "benchmarks/", "tools/", ".claude/") +TIER_LAST = ("diff_diff/", "rust/") + + +class FileBlock(NamedTuple): + path: str # display path (b-side, else a-side), decoded for the stub + data: bytes # the block's bytes, header line included + + +def split_blocks(patch: bytes) -> tuple[bytes, List[FileBlock]]: + """Split a unified diff into its preamble and per-file blocks. + + The preamble is anything before the first `diff --git ` header (normally + empty); it is kept verbatim. + """ + if not patch: + return b"", [] + starts: List[int] = [] + if patch.startswith(HEADER): + starts.append(0) + pos = 0 + needle = b"\n" + HEADER + while True: + i = patch.find(needle, pos) + if i < 0: + break + starts.append(i + 1) + pos = i + 1 + if not starts: + return patch, [] + preamble = patch[: starts[0]] + blocks: List[FileBlock] = [] + for k, s in enumerate(starts): + e = starts[k + 1] if k + 1 < len(starts) else len(patch) + data = patch[s:e] + blocks.append(FileBlock(path=_block_path(data), data=data)) + return preamble, blocks + + +_ESCAPES = { + b"a": b"\a", + b"b": b"\b", + b"t": b"\t", + b"n": b"\n", + b"v": b"\v", + b"f": b"\f", + b"r": b"\r", + b"\\": b"\\", + b'"': b'"', +} + + +def _c_unquote(raw: bytes) -> bytes: + """Undo git's C-style path quoting (`"a/b\303\251.py"`).""" + if not (len(raw) >= 2 and raw[:1] == b'"' and raw[-1:] == b'"'): + return raw + body = raw[1:-1] + out = bytearray() + i = 0 + while i < len(body): + c = body[i : i + 1] + if c != b"\\" or i + 1 >= len(body): + out += c + i += 1 + continue + nxt = body[i + 1 : i + 2] + if nxt in _ESCAPES: + out += _ESCAPES[nxt] + i += 2 + elif nxt.isdigit(): + digits = body[i + 1 : i + 4] + n = len(digits) + while n > 0 and not all(48 <= d <= 55 for d in digits[:n]): + n -= 1 + if n == 0: + out += nxt + i += 2 + else: + out.append(int(digits[:n], 8) & 0xFF) + i += 1 + n + else: + out += nxt + i += 2 + return bytes(out) + + +def _quoted_token_end(raw: bytes) -> Optional[int]: + """Index just past the closing quote of the C-quoted token at `raw[0]`.""" + if raw[:1] != b'"': + return None + i = 1 + while i < len(raw): + c = raw[i : i + 1] + if c == b"\\": + i += 2 + continue + if c == b'"': + return i + 1 + i += 1 + return None + + +def _strip_side(raw: bytes, prefix: bytes) -> Optional[bytes]: + """`b/path` (possibly C-quoted as `"b/path"`) -> `path`, else None.""" + val = _c_unquote(raw.rstrip(b"\t")) + if val.startswith(prefix): + return val[len(prefix) :] + return None + + +def _block_path(data: bytes) -> str: + """The block's canonical repo path, as git's own headers give it. + + Priority: the extended `rename to` / `copy to` header (renames and + copies, even without hunks), then `+++ b/` (else `--- a/` for a + deletion), then the `diff --git a/P b/P` line itself (binary, mode-only + and other hunk-less blocks). C-quoted paths are unquoted. Only ever + rendered as text in a stub and used for tiering. + """ + lines = data.split(b"\n", 60)[:60] + for line in lines[1:]: + for key in (b"rename to ", b"copy to "): + if line.startswith(key): + return _decode(_c_unquote(line[len(key) :])) + for line in lines[1:]: + if line.startswith(b"+++ "): + side = _strip_side(line[4:], b"b/") + if side is not None: + return _decode(side) + for line in lines[1:]: + if line.startswith(b"--- "): + side = _strip_side(line[4:], b"a/") + if side is not None: + return _decode(side) + rem = lines[0][len(HEADER) :] + if rem.startswith(b'"'): + # `"a/P" "b/P"`: scan the a-side token honouring backslash escapes + # (a path may itself contain `" `), then unquote the b-side. + end = _quoted_token_end(rem) + if end is not None and rem[end : end + 1] == b" ": + side = _strip_side(rem[end + 1 :], b"b/") + if side is not None: + return _decode(side) + elif rem.startswith(b"a/") and len(rem) >= 5 and (len(rem) - 5) % 2 == 0: + # Unquoted `a/P b/P`: the same path on both sides, so P is determined + # by the remainder's length even when P itself contains " b/". + n = (len(rem) - 5) // 2 + cand = rem[2 : 2 + n] + if rem == b"a/" + cand + b" b/" + cand: + return _decode(cand) + return _decode(rem) + + +def _decode(raw: bytes) -> str: + return raw.decode("utf-8", errors="replace") + + +def tier(path: str) -> int: + if path.startswith(TIER_LAST): + return 2 + if path.startswith(TIER_FIRST): + return 0 + return 1 + + +def _shell_quote(path: str) -> str: + return "'" + path.replace("'", "'\\''") + "'" + + +def stub(block: FileBlock, base_sha: Optional[str], head_sha: Optional[str]) -> bytes: + text = f"[diff omitted for prompt budget] {block.path} " f"({len(block.data):,} bytes)." + if base_sha and head_sha: + text += ( + " View it: git --no-pager diff --unified=5 " + f"{base_sha} {head_sha} -- {_shell_quote(':(literal)' + block.path)}" + ) + return text.encode("utf-8") + b"\n" + + +def _banner(n_omitted: int, n_total: int, budget: int, full_size: int) -> bytes: + text = ( + f"[prompt budget] {n_omitted} of {n_total} file diffs are omitted below so the " + f"prompt fits the reviewer input limit (budget {budget:,} bytes, full diff " + f'{full_size:,} bytes). Omitted files remain listed under "Changed files:" ' + f"above; each stub gives the git command that shows its hunks.\n\n" + ) + return text.encode("utf-8") + + +def _bounded(text: str, budget: int) -> bytes: + """A diagnostic line that never exceeds the budget (empty if it cannot fit).""" + raw = text.encode("utf-8") + if len(raw) > budget: + print(f"ai_review_diff_budget: {text.rstrip()}", file=sys.stderr) + return raw[:budget] + return raw + + +def fit_to_budget( + patch: bytes, + budget: int, + base_sha: Optional[str] = None, + head_sha: Optional[str] = None, +) -> bytes: + """Return `patch` unchanged if `len(patch) <= budget`, else a trimmed patch. + + Greedy: replace blocks by their stubs in (tier, -size) order until the + output fits, skipping blocks whose stub would not be smaller than the + block. The result is ALWAYS at most `budget` bytes: if even every + shrinking replacement cannot get there, the output is a single summary + line, itself truncated to the budget (empty at budget 0). + """ + if len(patch) <= budget: + return patch + preamble, blocks = split_blocks(patch) + if not blocks: + # No file headers - nothing to trim structurally; hard-truncate. + marker = b"\n[prompt budget] diff truncated here.\n" + if len(marker) > budget: + return _bounded("[prompt budget] diff truncated here.\n", budget) + return patch[: budget - len(marker)] + marker + + stubs = [stub(b, base_sha, head_sha) for b in blocks] + shrinking = [i for i in range(len(blocks)) if len(stubs[i]) < len(blocks[i].data)] + order = sorted(shrinking, key=lambda i: (tier(blocks[i].path), -len(blocks[i].data), i)) + omitted = [False] * len(blocks) + + def total(n_omitted: int) -> int: + body = sum(len(stubs[i]) if omitted[i] else len(blocks[i].data) for i in range(len(blocks))) + return len(preamble) + len(_banner(n_omitted, len(blocks), budget, len(patch))) + body + + n_omitted = 0 + for i in order: + if total(n_omitted) <= budget: + break + omitted[i] = True + n_omitted += 1 + + if total(n_omitted) > budget: + return _bounded( + f"[prompt budget] the unified diff ({len(patch):,} bytes across {len(blocks)} " + f"files) does not fit the remaining budget of {budget:,} bytes even as stubs; " + f'see "Changed files:" above and run git diff per file.\n', + budget, + ) + + out = [preamble, _banner(n_omitted, len(blocks), budget, len(patch))] + for i, b in enumerate(blocks): + out.append(stubs[i] if omitted[i] else b.data) + return b"".join(out) + + +def _sha(value: Optional[str]) -> Optional[str]: + if value is None or value == "": + return None + if not _SHA_RE.match(value): + raise argparse.ArgumentTypeError(f"not a hex object id: {value!r}") + return value + + +def main(argv: Optional[Sequence[str]] = None) -> int: + ap = argparse.ArgumentParser(description=(__doc__ or "").split("\n\n")[0]) + ap.add_argument("--diff", required=True, type=Path, help="unified diff file to fit") + ap.add_argument("--budget", required=True, type=int, help="max output size in bytes") + ap.add_argument("--base-sha", type=_sha, default=None, help="for the per-stub git command") + ap.add_argument("--head-sha", type=_sha, default=None, help="for the per-stub git command") + ap.add_argument("-o", "--output", type=Path, default=None, help="default: stdout") + args = ap.parse_args(argv) + if args.budget < 0: + ap.error("--budget must be >= 0") + patch = args.diff.read_bytes() + out = fit_to_budget(patch, args.budget, args.base_sha, args.head_sha) + if args.output is None: + sys.stdout.buffer.write(out) + sys.stdout.buffer.flush() + else: + args.output.write_bytes(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main())