diff --git a/infra/flue-review/.flue/lib/diff-budget.ts b/infra/flue-review/.flue/lib/diff-budget.ts new file mode 100644 index 0000000000..526b43b172 --- /dev/null +++ b/infra/flue-review/.flue/lib/diff-budget.ts @@ -0,0 +1,74 @@ +// Size budget for the staged PR diff. The review agent starts by reading the +// whole diff file, so an oversized diff (generated types, lockfiles, large +// catalogs) lands in the model context verbatim and kills the model call. +// Oversized per-file sections are elided down to their headers with a note; +// the agent reads those files from the checkout instead. + +const DEFAULT_PER_FILE_BYTES = 48 * 1024; +const DEFAULT_TOTAL_BYTES = 384 * 1024; + +export interface DiffBudget { + readonly perFileBytes?: number; + readonly totalBytes?: number; +} + +interface Section { + text: string; + elided: boolean; +} + +export function elideLargeDiffSections(diff: string, budget: DiffBudget = {}): string { + const perFileBytes = budget.perFileBytes ?? DEFAULT_PER_FILE_BYTES; + const totalBytes = budget.totalBytes ?? DEFAULT_TOTAL_BYTES; + if (diff.length <= Math.min(perFileBytes, totalBytes)) return diff; + + const sections = splitSections(diff); + for (const section of sections) { + if (!section.elided && section.text.length > perFileBytes) elide(section); + } + // Still over the total budget: elide the largest remaining sections until + // under it (or nothing left to elide). + let total = sections.reduce((n, s) => n + s.text.length, 0); + while (total > totalBytes) { + const next = sections + .filter((s) => !s.elided) + .toSorted((a, b) => b.text.length - a.text.length)[0]; + if (!next) break; + total -= next.text.length; + elide(next); + total += next.text.length; + } + return sections.map((s) => s.text).join(""); +} + +function splitSections(diff: string): Section[] { + const starts: number[] = []; + const re = /^diff --git /gm; + for (let m = re.exec(diff); m; m = re.exec(diff)) starts.push(m.index); + if (starts.length === 0) return [{ text: diff, elided: false }]; + const sections: Section[] = []; + if (starts[0] !== 0) sections.push({ text: diff.slice(0, starts[0]), elided: true }); + for (let i = 0; i < starts.length; i++) { + const end = i + 1 < starts.length ? starts[i + 1] : diff.length; + sections.push({ text: diff.slice(starts[i], end), elided: false }); + } + return sections; +} + +function elide(section: Section): void { + // Mark unconditionally: a section this function cannot reduce must still + // leave the total-budget loop's candidate pool, or the loop never shrinks. + section.elided = true; + const lines = section.text.split("\n"); + // Keep the file header: everything up to and including the `+++` line (or + // the whole header for binary/rename-only sections with no hunks). + let headerEnd = lines.findIndex((line) => line.startsWith("+++ ")); + if (headerEnd === -1) headerEnd = lines.findIndex((line) => line.startsWith("@@ ")) - 1; + if (headerEnd < 0) return; + const body = lines.length - (headerEnd + 1); + section.text = [ + ...lines.slice(0, headerEnd + 1), + `(diff content elided: ${body} lines over the size budget -- read this file from the checkout instead)`, + "", + ].join("\n"); +} diff --git a/infra/flue-review/.flue/workflows/review.ts b/infra/flue-review/.flue/workflows/review.ts index 651f4beec1..dd4f631acc 100644 --- a/infra/flue-review/.flue/workflows/review.ts +++ b/infra/flue-review/.flue/workflows/review.ts @@ -27,6 +27,7 @@ import { getCloudflareContext, getDurableObjectIdentity } from "@flue/runtime/cl import * as v from "valibot"; import { withCapacityRetry } from "../lib/capacity.js"; +import { elideLargeDiffSections } from "../lib/diff-budget.js"; import { readAppCreds, mintInstallationToken, @@ -358,7 +359,7 @@ async function run(context: ActionContext): Promise< payload.baseSha, payload.headSha, ); - await context.harness.fs.writeFile(DIFF_PATH, diff); + await context.harness.fs.writeFile(DIFF_PATH, elideLargeDiffSections(diff)); stage = "model_review"; if ( diff --git a/infra/flue-review/test/diff-budget.test.ts b/infra/flue-review/test/diff-budget.test.ts new file mode 100644 index 0000000000..511e65805f --- /dev/null +++ b/infra/flue-review/test/diff-budget.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it } from "vitest"; + +import { elideLargeDiffSections } from "../.flue/lib/diff-budget.js"; + +function fileSection(path: string, lines: number, line = "+const x = 1;"): string { + return [ + `diff --git a/${path} b/${path}`, + "index 0000000..1111111 100644", + `--- a/${path}`, + `+++ b/${path}`, + "@@ -0,0 +1 @@", + ...Array.from({ length: lines }).fill(line).map(String), + "", + ].join("\n"); +} + +describe("elideLargeDiffSections", () => { + it("returns a small diff unchanged", () => { + const diff = fileSection("src/a.ts", 10) + fileSection("src/b.ts", 20); + expect(elideLargeDiffSections(diff)).toBe(diff); + }); + + it("elides a section over the per-file budget, keeping its header", () => { + const big = fileSection("types.d.ts", 2_000); + const small = fileSection("src/a.ts", 5); + const out = elideLargeDiffSections(small + big, { perFileBytes: 1_000 }); + expect(out).toContain(small); + expect(out).toContain("diff --git a/types.d.ts b/types.d.ts"); + expect(out).toContain("+++ b/types.d.ts"); + expect(out).toMatch(/diff content elided: \d+ lines/); + expect(out).not.toContain("+const x = 1;\n+const x = 1;\n".repeat(50)); + }); + + it("elides largest sections first until under the total budget", () => { + const a = fileSection("a.ts", 30); + const b = fileSection("b.ts", 60); + const c = fileSection("c.ts", 10); + const out = elideLargeDiffSections(a + b + c, { + perFileBytes: 10_000, + totalBytes: a.length + c.length + 400, + }); + expect(out).toContain("diff content elided"); + expect(out).toContain(a); + expect(out).toContain(c); + expect(out).not.toContain(b); + }); + + it("skips an unreducible headerless section instead of looping on it", () => { + const headerless = `diff --git a/blob.bin b/blob.bin\nBinary files differ\n${"x\n".repeat(1_000)}`; + const small = fileSection("src/a.ts", 5); + const out = elideLargeDiffSections(small + headerless, { + perFileBytes: 500, + totalBytes: 600, + }); + expect(out).toContain("Binary files differ"); + expect(out).toContain("+++ b/src/a.ts"); + }); + + it("leaves a header-only section (no hunks) alone", () => { + const rename = [ + "diff --git a/old.ts b/new.ts", + "similarity index 100%", + "rename from old.ts", + "rename to new.ts", + "", + ].join("\n"); + const filler = fileSection("big.ts", 2_000); + const out = elideLargeDiffSections(rename + filler, { perFileBytes: 1_000, totalBytes: 1_500 }); + expect(out).toContain("rename from old.ts"); + }); +});