diff --git a/packages/highlight/src/__tests__/highlighter.test.ts b/packages/highlight/src/__tests__/highlighter.test.ts index fde3f9d740..095f9689a8 100644 --- a/packages/highlight/src/__tests__/highlighter.test.ts +++ b/packages/highlight/src/__tests__/highlighter.test.ts @@ -1,4 +1,6 @@ import { describe, it, expect } from "vitest"; +import { setImmediate } from "node:timers/promises"; +import { highlightCodeWithBudget, HIGHLIGHT_MAX_SOURCE_BYTES } from "../budgeted-highlighter.js"; import type { HighlightToken } from "../types.js"; import { highlightCode, highlightCodeInSteps, highlightLine } from "../highlighter.js"; @@ -420,3 +422,55 @@ describe("highlightLine", () => { expect(tokens).toEqual([{ text: "hello", style: null }]); }); }); + +describe("highlightCodeWithBudget", () => { + it("yields to IO while producing byte-identical output below the source budget", async () => { + const code = "const value = Math.max(42, 7);\n".repeat(3_000); + const expected = JSON.stringify(highlightCode(code, "file.ts")); + let ticks = 0; + let longestGap = 0; + let lastTick = performance.now(); + const timer = setInterval(() => { + const now = performance.now(); + longestGap = Math.max(longestGap, now - lastTick); + lastTick = now; + ticks++; + }, 1); + try { + const actual = await highlightCodeWithBudget(code, "file.ts", { + yieldToEventLoop: setImmediate, + }); + await new Promise((resolve) => setTimeout(resolve, 5)); + expect(ticks).toBeGreaterThan(1); + expect(longestGap).toBeLessThan(50); + expect(JSON.stringify(actual)).toBe(expected); + } finally { + clearInterval(timer); + } + }); + + it("budgets UTF-8 bytes before parsing, including multibyte text", async () => { + expect( + await highlightCodeWithBudget("é".repeat(HIGHLIGHT_MAX_SOURCE_BYTES / 2 + 1), "file.ts", { + yieldToEventLoop: setImmediate, + }), + ).toBeNull(); + expect( + await highlightCodeWithBudget("a".repeat(HIGHLIGHT_MAX_SOURCE_BYTES + 1), "file.ts", { + yieldToEventLoop: setImmediate, + }), + ).toBeNull(); + }); + + it("stops a parser that is cancelled between real event-loop turns", async () => { + const controller = new AbortController(); + const pending = highlightCodeWithBudget( + "const value = Math.max(42, 7);\n".repeat(3_000), + "file.ts", + { yieldToEventLoop: setImmediate, signal: controller.signal }, + ); + await setImmediate(); + controller.abort(); + expect(await pending).toBeNull(); + }); +}); diff --git a/packages/highlight/src/budgeted-highlighter.ts b/packages/highlight/src/budgeted-highlighter.ts new file mode 100644 index 0000000000..e949a619ff --- /dev/null +++ b/packages/highlight/src/budgeted-highlighter.ts @@ -0,0 +1,49 @@ +import { highlightCodeInSteps } from "./highlighter.js"; +import type { HighlightToken } from "./types.js"; + +// Full source, rather than patch size, bounds parser and token allocation work. +// Each token consumes a source code unit or an empty line: at most the byte cap + 1 records. +export const HIGHLIGHT_MAX_SOURCE_BYTES = 768 * 1024; +const SLICE_MS = 4; + +/** + * Returns null when source exceeds the budget or the request is cancelled. + * The host supplies a real event-loop yield (setImmediate on the daemon). + * Existing synchronous highlighting remains available for its existing consumers. + */ +export async function highlightCodeWithBudget( + code: string, + filename: string, + options: { yieldToEventLoop: () => Promise; signal?: { readonly aborted: boolean } }, +): Promise { + if (options.signal?.aborted || exceedsSourceBudget(code)) return null; + const steps = highlightCodeInSteps(code, filename); + let sliceStarted = Date.now(); + try { + for (;;) { + if (options.signal?.aborted) return null; + const step = steps.next(); + if (step.done) return step.value; + if (Date.now() - sliceStarted >= SLICE_MS) { + await options.yieldToEventLoop(); + sliceStarted = Date.now(); + } + } + } finally { + steps.return([]); + } +} + +function exceedsSourceBudget(code: string): boolean { + if (code.length > HIGHLIGHT_MAX_SOURCE_BYTES) return true; + let bytes = 0; + for (const char of code) { + const point = char.codePointAt(0)!; + if (point < 0x80) bytes++; + else if (point < 0x800) bytes += 2; + else if (point < 0x10000) bytes += 3; + else bytes += 4; + if (bytes > HIGHLIGHT_MAX_SOURCE_BYTES) return true; + } + return false; +} diff --git a/packages/highlight/src/index.ts b/packages/highlight/src/index.ts index a69820ee28..da468e2f87 100644 --- a/packages/highlight/src/index.ts +++ b/packages/highlight/src/index.ts @@ -15,3 +15,4 @@ export { isSyntaxThemeId, resolveSyntaxColors, } from "./themes.js"; +export { highlightCodeWithBudget, HIGHLIGHT_MAX_SOURCE_BYTES } from "./budgeted-highlighter.js"; diff --git a/packages/server/src/server/checkout-diff-cache.test.ts b/packages/server/src/server/checkout-diff-cache.test.ts index 37f9283320..a587f89894 100644 --- a/packages/server/src/server/checkout-diff-cache.test.ts +++ b/packages/server/src/server/checkout-diff-cache.test.ts @@ -132,3 +132,22 @@ test.each(["base", "uncommitted"] as const)( expect(await read("next", "next-after")).toEqual({ diff: "next-after" }); }, ); + +test("superseding a diff cancels its highlighting while identical readers share work", async () => { + const cache = new CheckoutDiffCache(() => 0); + const release = Promise.withResolvers<{ diff: string }>(); + let activeSignal: AbortSignal | undefined; + const load = (signal: AbortSignal) => { + activeSignal = signal; + return release.promise; + }; + const first = cache.read({ cwd: "repo", compare: { mode: "base" }, load }); + await Promise.resolve(); + expect(activeSignal).toBeDefined(); + expect(activeSignal!.aborted).toBe(false); + expect(cache.read({ cwd: "repo", compare: { mode: "base" }, load })).toBe(first); + cache.invalidate("repo", "base"); + expect(activeSignal!.aborted).toBe(true); + release.resolve({ diff: "old snapshot" }); + await first; +}); diff --git a/packages/server/src/server/checkout-diff-cache.ts b/packages/server/src/server/checkout-diff-cache.ts index 4547cdf2e5..71d7d95497 100644 --- a/packages/server/src/server/checkout-diff-cache.ts +++ b/packages/server/src/server/checkout-diff-cache.ts @@ -4,6 +4,7 @@ import type { CheckoutDiffCompare, CheckoutDiffResult } from "../utils/checkout- interface ActiveRead { promise: Promise; invalidated: boolean; + controller: AbortController; } interface CachedDiff { @@ -17,7 +18,7 @@ interface CheckoutDiffReadInput { compare: CheckoutDiffCompare; force?: boolean; reason?: string; - load: () => Promise; + load: (signal: AbortSignal) => Promise; } /** Completed payloads are bounded; active work survives invalidation and eviction. */ @@ -51,6 +52,7 @@ export class CheckoutDiffCache { if (pending) { if (force) { pending.invalidated = true; + pending.controller.abort(); this.values.delete(key); } if (!pending.invalidated) return pending.promise; @@ -66,8 +68,9 @@ export class CheckoutDiffCache { // Defer load until ownership is installed, including synchronous invalidations. const read: ActiveRead = { invalidated: false, + controller: new AbortController(), promise: Promise.resolve() - .then(load) + .then(() => load(read.controller.signal)) .then((value) => { // Never repopulate the cache with a snapshot invalidated during its build. if (!read.invalidated) @@ -86,7 +89,10 @@ export class CheckoutDiffCache { if (cachedCwd !== cwd || (cachedMode !== mode && cachedMode !== "all")) continue; this.values.delete(key); const pending = this.active.get(key); - if (pending) pending.invalidated = true; + if (pending) { + pending.invalidated = true; + pending.controller.abort(); + } } } } diff --git a/packages/server/src/server/daemon-e2e/checkout-diff-content.e2e.test.ts b/packages/server/src/server/daemon-e2e/checkout-diff-content.e2e.test.ts index 09da7506c9..fad67ae7b9 100644 --- a/packages/server/src/server/daemon-e2e/checkout-diff-content.e2e.test.ts +++ b/packages/server/src/server/daemon-e2e/checkout-diff-content.e2e.test.ts @@ -1,6 +1,7 @@ -import { execFileSync } from "node:child_process"; +import { execFile, execFileSync } from "node:child_process"; import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; -import { rm } from "node:fs/promises"; +import { rm, writeFile } from "node:fs/promises"; +import { promisify } from "node:util"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { expect, test } from "vitest"; @@ -223,3 +224,73 @@ test("All changes subscriptions refresh on file edits and commits", async () => await rm(cwd, { recursive: true, force: true, maxRetries: 10, retryDelay: 100 }); } }, 60_000); + +test("large source diffs stay renderable for reads, subscriptions, and commit files", async () => { + const ctx = await createDaemonTestContext(); + const cwd = mkdtempSync(join(tmpdir(), "checkout-diff-budget-")); + const runGit = promisify(execFile); + try { + git(cwd, "init", "-b", "main"); + git(cwd, "config", "user.name", "Test"); + git(cwd, "config", "user.email", "test@example.com"); + git(cwd, "config", "core.autocrlf", "false"); + const original = "const value = Math.max(42, 7);\n".repeat(Math.ceil((5 * 1024 * 1024) / 30)); + writeFileSync(join(cwd, "generated.ts"), original); + git(cwd, "add", "."); + git(cwd, "-c", "commit.gpgsign=false", "commit", "-m", "initial"); + await writeFile(join(cwd, "generated.ts"), original.replace("42", "43")); + const pending = ctx.client.getCheckoutDiff(cwd, { mode: "uncommitted" }); + const rpcStarted = performance.now(); + await ctx.client.fetchAgents({}); + const rpcMs = performance.now() - rpcStarted; + const initial = await pending; + expect(initial.error).toBeNull(); + const check = (file: (typeof initial.files)[number] | null | undefined, value: string) => { + expect(file).toMatchObject({ path: "generated.ts", additions: 1, deletions: 1 }); + const lines = file!.hunks.flatMap((hunk) => hunk.lines); + expect(lines.filter((line) => line.type === "add").map((line) => line.content)).toEqual([ + `const value = Math.max(${value}, 7);`, + ]); + expect(lines.every((line) => line.tokens === undefined)).toBe(true); + }; + check(initial.files[0], "43"); + const subscription = ctx.client.observeCheckoutDiff(cwd, { mode: "uncommitted" }); + try { + let latest = await subscription.ready; + check(latest.files[0], "43"); + subscription.subscribe({ + snapshot: (snapshot) => { + latest = snapshot; + }, + update: (message) => { + if (message.type === "checkout_diff_update") latest = message.payload; + }, + }); + await writeFile(join(cwd, "generated.ts"), original.replace("42", "44")); + await expect + .poll( + () => + latest.files[0]?.hunks.flatMap((hunk) => hunk.lines).find((line) => line.type === "add") + ?.content, + { timeout: 15_000 }, + ) + .toBe("const value = Math.max(44, 7);"); + check(latest.files[0], "44"); + } finally { + await subscription.release(); + } + await runGit("git", ["add", "."], { cwd }); + await runGit("git", ["-c", "commit.gpgsign=false", "commit", "-m", "edit"], { cwd }); + const { stdout } = await runGit("git", ["rev-parse", "HEAD"], { cwd }); + const commit = await ctx.client.getCommitFileDiff(cwd, stdout.trim(), "generated.ts"); + check(commit.file, "44"); + console.info({ + sourceBytes: Buffer.byteLength(original), + cheapRpcMs: rpcMs, + isolatedDaemonPort: ctx.daemon.port, + }); + } finally { + await ctx.cleanup(); + await rm(cwd, { recursive: true, force: true, maxRetries: 10, retryDelay: 100 }); + } +}, 60_000); diff --git a/packages/server/src/server/utils/diff-highlighter.test.ts b/packages/server/src/server/utils/diff-highlighter.test.ts index fda98d6cbf..7599ca68e7 100644 --- a/packages/server/src/server/utils/diff-highlighter.test.ts +++ b/packages/server/src/server/utils/diff-highlighter.test.ts @@ -1,8 +1,12 @@ -import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { mkdtempSync, rmSync, writeFileSync, readFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { describe, expect, it } from "vitest"; +import { setImmediate } from "node:timers/promises"; +import { performance } from "node:perf_hooks"; +import { highlightCode } from "@getpaseo/highlight"; import { + DiffHighlightRequests, parseDiff, reconstructNewFile, reconstructOldFile, @@ -819,3 +823,118 @@ index 1111111..2222222 100644 } }); }); + +describe("diff highlighting budget", () => { + it.each(["supplied", "local"] as const)( + "keeps a one-line diff on a 5 MiB %s source responsive and renders it without tokens", + async (source) => { + const cwd = mkdtempSync(join(tmpdir(), "diff-highlight-budget-")); + const oldFileContent = "const value = Math.max(42, 7);\n".repeat( + Math.ceil((5 * 1024 * 1024) / 30), + ); + const newFileContent = oldFileContent.replace("42", "43"); + writeFileSync(join(cwd, "example.ts"), newFileContent); + const file = parseDiff(`diff --git a/example.ts b/example.ts +--- a/example.ts ++++ b/example.ts +@@ -1 +1 @@ +-const value = Math.max(42, 7); ++const value = Math.max(43, 7);`)[0]; + let lastTick = performance.now(); + let longestGap = 0; + const timer = setInterval(() => { + const now = performance.now(); + longestGap = Math.max(longestGap, now - lastTick); + lastTick = now; + }, 1); + try { + await setImmediate(); + const highlighted = await highlightDiffWithFileContent( + file, + cwd, + source === "supplied" ? { oldFileContent } : {}, + ); + await new Promise((resolve) => setTimeout(resolve, 5)); + console.info({ sourceBytes: Buffer.byteLength(oldFileContent), longestGap }); + expect(longestGap).toBeLessThan(50); + expect(highlighted).toEqual(file); + } finally { + clearInterval(timer); + rmSync(cwd, { recursive: true, force: true }); + } + }, + 30_000, + ); + + it("preserves byte-identical tokens below the budget, including multiline syntax", async () => { + const oldFileContent = "/*\nold comment\n*/\nconst value = 42;"; + const newFileContent = oldFileContent.replace("old", "new"); + const file = parseDiff(`diff --git a/example.ts b/example.ts +--- a/example.ts ++++ b/example.ts +@@ -2 +2 @@ +-old comment ++new comment`)[0]; + const highlighted = await highlightDiffWithFileContent(file, ".", { + oldFileContent, + newFileContent, + }); + expect(JSON.stringify(highlighted.hunks[0].lines[1].tokens)).toBe( + JSON.stringify(highlightCode(oldFileContent, file.path)[1]), + ); + expect(JSON.stringify(highlighted.hunks[0].lines[2].tokens)).toBe( + JSON.stringify(highlightCode(newFileContent, file.path)[1]), + ); + }); +}); + +it("a newer request cancels the whole previous diff even before its sources arrive", async () => { + const requests = new DiffHighlightRequests(); + const file = parseDiff(SIMPLE_DIFF)[0]; + const source = "const value = Math.max(42, 7);\n".repeat(3_000); + const old = requests.run("checkout-target", async (signal) => { + await setImmediate(); + return highlightDiffWithFileContent(file, ".", { + oldFileContent: source, + newFileContent: source, + signal, + }); + }); + const current = requests.run("checkout-target", (signal) => + highlightDiffWithFileContent(file, ".", { + oldFileContent: "const bar = 2;", + newFileContent: "const bar = 3;", + signal, + }), + ); + expect(await old).toEqual(file); + expect((await current).hunks[0].lines.some((line) => line.tokens)).toBe(true); +}); + +it("keeps real session.ts diffs highlighted and byte-identical under the source budget", async () => { + const cwd = mkdtempSync(join(tmpdir(), "diff-highlight-real-source-")); + try { + const oldFileContent = readFileSync(new URL("../session.ts", import.meta.url), "utf8"); + const newFileContent = oldFileContent.replace( + "WorkspaceRuntimeStore", + "WorkspaceRuntimeSnapshot", + ); + writeFileSync(join(cwd, "session.ts"), newFileContent); + const oldLine = oldFileContent.split("\n")[0]; + const newLine = newFileContent.split("\n")[0]; + const [file] = parseDiff(`diff --git a/session.ts b/session.ts +--- a/session.ts ++++ b/session.ts +@@ -1 +1 @@ +-${oldLine} ++${newLine}`); + const highlighted = await highlightDiffWithFileContent(file, cwd, { oldFileContent }); + const removed = highlighted.hunks[0].lines.find((line) => line.type === "remove"); + const added = highlighted.hunks[0].lines.find((line) => line.type === "add"); + expect(removed?.tokens).toEqual(highlightCode(oldFileContent, "session.ts")[0]); + expect(added?.tokens).toEqual(highlightCode(newFileContent, "session.ts")[0]); + expect(added?.tokens?.some((token) => token.style !== null)).toBe(true); + } finally { + rmSync(cwd, { recursive: true, force: true }); + } +}); diff --git a/packages/server/src/server/utils/diff-highlighter.ts b/packages/server/src/server/utils/diff-highlighter.ts index 3c2df279e2..d290155899 100644 --- a/packages/server/src/server/utils/diff-highlighter.ts +++ b/packages/server/src/server/utils/diff-highlighter.ts @@ -1,6 +1,14 @@ -import { readFile } from "node:fs/promises"; +import { open } from "node:fs/promises"; +import { constants } from "node:fs"; +import { setImmediate } from "node:timers/promises"; import { resolve } from "node:path"; -import { highlightCode, isLanguageSupported, type HighlightToken } from "@getpaseo/highlight"; +import { + highlightCode, + highlightCodeWithBudget, + HIGHLIGHT_MAX_SOURCE_BYTES, + isLanguageSupported, + type HighlightToken, +} from "@getpaseo/highlight"; import { unquoteGitPath } from "../../utils/git-quoted-path.js"; const MAX_DIFF_HIGHLIGHT_LINE_CHARS = 10_000; @@ -30,15 +38,40 @@ export interface ParsedDiffFile { } interface HighlightDiffWithFileContentOptions { + signal?: AbortSignal; oldFileContent?: string | null; newFileContent?: string | null; } interface ParseAndHighlightDiffOptions { + signal?: AbortSignal; getOldFileContent?: (file: ParsedDiffFile) => Promise; getNewFileContent?: (file: ParsedDiffFile) => Promise; } +/** Owns supersession for a whole diff, including sources that arrive after cancellation. */ +export class DiffHighlightRequests { + private readonly active = new Map(); + + async run( + target: string, + operation: (signal: AbortSignal) => Promise, + parentSignal?: AbortSignal, + ): Promise { + this.active.get(target)?.abort(); + const controller = new AbortController(); + this.active.set(target, controller); + const signal = parentSignal + ? AbortSignal.any([parentSignal, controller.signal]) + : controller.signal; + try { + return await operation(signal); + } finally { + if (this.active.get(target) === controller) this.active.delete(target); + } + } +} + /** * Parse a unified diff into structured data */ @@ -269,20 +302,71 @@ function buildTokenLookup( return lookup; } -function buildFullFileTokenLookup( +async function buildFullFileTokenLookup( fileContent: string, path: string, -): Map | null { - if (hasOversizedLine(fileContent)) return null; + signal?: AbortSignal, +): Promise | null> { + const highlighted = await highlightCodeWithBudget(fileContent, path, { + yieldToEventLoop: setImmediate, + signal, + }); + if (!highlighted) return null; + return new Map(highlighted.map((tokens, index) => [index + 1, tokens])); +} - const lookup = new Map(); - const highlighted = highlightCode(fileContent, path); - - for (let i = 0; i < highlighted.length; i++) { - lookup.set(i + 1, highlighted[i]); +// null means over budget; undefined means unavailable and permits reconstruction. +async function readHighlightSource(path: string): Promise { + let handle; + try { + handle = await open(path, constants.O_RDONLY | constants.O_NONBLOCK); + const stat = await handle.stat(); + if (!stat.isFile() || stat.size > HIGHLIGHT_MAX_SOURCE_BYTES) return null; + const buffer = Buffer.alloc(HIGHLIGHT_MAX_SOURCE_BYTES + 1); + let offset = 0; + while (offset < buffer.length) { + const { bytesRead } = await handle.read(buffer, offset, buffer.length - offset, null); + if (bytesRead === 0) break; + offset += bytesRead; + } + return offset > HIGHLIGHT_MAX_SOURCE_BYTES ? null : buffer.toString("utf8", 0, offset); + } catch { + return undefined; + } finally { + await handle?.close(); } +} - return lookup; +function exceedsSourceBudget(content: string): boolean { + return ( + content.length > HIGHLIGHT_MAX_SOURCE_BYTES || + Buffer.byteLength(content) > HIGHLIGHT_MAX_SOURCE_BYTES + ); +} + +async function buildBudgetedReconstructedTokens(file: ParsedDiffFile, signal?: AbortSignal) { + const newLines = reconstructNewFile(file.hunks); + const oldLines = reconstructOldFile(file.hunks); + const build = async (lines: Map) => { + // Sparse hunks can be far apart. Admit the span before filling its gaps. + const numbers = Array.from(lines.keys()); + if (numbers.length && numbers[numbers.length - 1] - numbers[0] > HIGHLIGHT_MAX_SOURCE_BYTES) + return null; + let bytes = 0; + for (const line of lines.values()) { + bytes += Buffer.byteLength(line) + 1; + if (bytes > HIGHLIGHT_MAX_SOURCE_BYTES) return null; + } + const highlighted = await highlightCodeWithBudget(buildFileContent(lines), file.path, { + yieldToEventLoop: setImmediate, + signal, + }); + return highlighted ? buildTokenLookup(lines, highlighted) : null; + }; + const newTokensByLine = await build(newLines); + if (!newTokensByLine) return null; + const oldTokensByLine = await build(oldLines); + return oldTokensByLine ? { newTokensByLine, oldTokensByLine } : null; } function buildReconstructedTokenLookups(file: ParsedDiffFile): { @@ -328,36 +412,56 @@ export async function highlightDiffWithFileContent( cwd: string, options: HighlightDiffWithFileContentOptions = {}, ): Promise { - if (!isLanguageSupported(file.path) || hasOversizedDiffLine(file)) { - return file; - } + if (!isDiffHighlightEligible(file, options)) return file; - let newTokensByLine: Map | null = null; - const oldTokensByLine = - typeof options.oldFileContent === "string" - ? buildFullFileTokenLookup(options.oldFileContent, file.path) - : null; - if (typeof options.newFileContent === "string") { - newTokensByLine = buildFullFileTokenLookup(options.newFileContent, file.path); - } else { - try { - const fileContent = await readFile(resolve(cwd, file.path), "utf-8"); - newTokensByLine = buildFullFileTokenLookup(fileContent, file.path); - } catch { - // Deleted or unavailable files use reconstructed tokens below. - } - } + const newContent = + typeof options.newFileContent === "string" + ? options.newFileContent + : await readHighlightSource(resolve(cwd, file.path)); + if (newContent === null || options.signal?.aborted) return file; - if (newTokensByLine && oldTokensByLine) - return applyTokensToHunks(file, newTokensByLine, oldTokensByLine); - const reconstructedTokens = buildReconstructedTokenLookups(file); + const oldTokens = await highlightAvailableSource( + options.oldFileContent, + file.path, + options.signal, + ); + if (oldTokens === null || options.signal?.aborted) return file; + const newTokens = await highlightAvailableSource(newContent, file.path, options.signal); + if (newTokens === null || options.signal?.aborted) return file; + if (oldTokens && newTokens) return applyTokensToHunks(file, newTokens, oldTokens); + const reconstructed = await buildBudgetedReconstructedTokens(file, options.signal); + if (!reconstructed || options.signal?.aborted) return file; return applyTokensToHunks( file, - newTokensByLine ?? reconstructedTokens.newTokensByLine, - oldTokensByLine ?? reconstructedTokens.oldTokensByLine, + newTokens ?? reconstructed.newTokensByLine, + oldTokens ?? reconstructed.oldTokensByLine, ); } +function isDiffHighlightEligible( + file: ParsedDiffFile, + options: HighlightDiffWithFileContentOptions, +): boolean { + return ( + !options.signal?.aborted && + isLanguageSupported(file.path) && + !hasOversizedDiffLine(file) && + ![options.oldFileContent, options.newFileContent].some( + (content) => typeof content === "string" && exceedsSourceBudget(content), + ) + ); +} + +async function highlightAvailableSource( + content: string | null | undefined, + path: string, + signal?: AbortSignal, +): Promise | null | undefined> { + // Preserve the existing hunk-context fallback for oversized individual lines. + if (typeof content !== "string" || hasOversizedLine(content)) return undefined; + return buildFullFileTokenLookup(content, path, signal); +} + function applyTokensToHunks( file: ParsedDiffFile, newTokensByLine: Map, @@ -406,20 +510,24 @@ export async function parseAndHighlightDiff( ): Promise { const files = parseDiff(diffText); - const highlightedFiles = await Promise.all( - files.map(async (file) => { - const [oldFileContent, newFileContent] = await Promise.all([ - options.getOldFileContent?.(file), - options.getNewFileContent?.(file), - ]); - - return highlightDiffWithFileContent(file, cwd, { + const highlightedFiles: ParsedDiffFile[] = []; + for (const file of files) { + if (options.signal?.aborted) { + highlightedFiles.push(file); + continue; + } + const [oldFileContent, newFileContent] = await Promise.all([ + options.getOldFileContent?.(file), + options.getNewFileContent?.(file), + ]); + highlightedFiles.push( + await highlightDiffWithFileContent(file, cwd, { oldFileContent: oldFileContent ?? undefined, newFileContent: newFileContent ?? undefined, - }); - }), - ); - + signal: options.signal, + }), + ); + } return highlightedFiles; } diff --git a/packages/server/src/server/workspace-git-service.ts b/packages/server/src/server/workspace-git-service.ts index dcd2ead214..493daa1ddd 100644 --- a/packages/server/src/server/workspace-git-service.ts +++ b/packages/server/src/server/workspace-git-service.ts @@ -815,11 +815,16 @@ export class WorkspaceGitServiceImpl implements WorkspaceGitService { cwd: normalizedCwd, compare: normalizedOptions, ...readOptions, - load: () => - this.deps.getCheckoutDiff(normalizedCwd, normalizedOptions, { - paseoHome: this.paseoHome, - worktreesRoot: this.worktreesRoot, - }), + load: (signal) => + this.deps.getCheckoutDiff( + normalizedCwd, + normalizedOptions, + { + paseoHome: this.paseoHome, + worktreesRoot: this.worktreesRoot, + }, + signal, + ), }); } diff --git a/packages/server/src/utils/checkout-git.ts b/packages/server/src/utils/checkout-git.ts index fbbb494b4a..293a5dcc4e 100644 --- a/packages/server/src/utils/checkout-git.ts +++ b/packages/server/src/utils/checkout-git.ts @@ -11,6 +11,7 @@ import { maxSecuredPlaintextByteLength } from "@getpaseo/relay"; import type { Logger } from "pino"; import type { ParsedDiffFile } from "../server/utils/diff-highlighter.js"; import { + DiffHighlightRequests, highlightDiffWithFileContent, parseAndHighlightDiff, parseDiff, @@ -2703,6 +2704,8 @@ async function readCommitRenames(cwd: string, sha: string): Promise { + return diffHighlightRequests.run( + JSON.stringify(["commit", resolve(input.cwd), input.sha, input.path]), + (signal) => readCommitFileDiff(input, signal), + ); +} + +async function readCommitFileDiff( + { + cwd, + sha, + path, + }: { + cwd: string; + sha: string; + path: string; + }, + signal: AbortSignal, +): Promise { const worktreeRoot = await requireGitWorktreeRoot(cwd); const renames = await getCommitRenames(worktreeRoot, sha); const oldPath = renames.get(path); @@ -2750,6 +2767,7 @@ export async function getCommitFileDiff({ } const parsedFiles = await parseAndHighlightDiff(stdout, worktreeRoot, { + signal, getOldFileContent: (file) => readGitFileContentAtRef(worktreeRoot, `${sha}^`, file.path), getNewFileContent: (file) => readGitFileContentAtRef(worktreeRoot, sha, file.path), }); @@ -3132,6 +3150,7 @@ export function warmCheckoutShortstatInBackground( } interface AppendStructuredTrackedDiffsInput { + signal: AbortSignal; cwd: string; trackedChanges: CheckoutFileChange[]; trackedNumstatByPath: Map; @@ -3147,6 +3166,7 @@ interface AppendStructuredTrackedDiffsInput { } async function buildHighlightedTrackedDiffFile(input: { + signal: AbortSignal; cwd: string; change: CheckoutFileChange; parsedFile: ParsedDiffFile; @@ -3169,6 +3189,7 @@ async function buildHighlightedTrackedDiffFile(input: { : null, ]); const highlightedFile = await highlightDiffWithFileContent(parsedFile, cwd, { + signal: input.signal, oldFileContent, newFileContent, }); @@ -3254,6 +3275,7 @@ async function appendStructuredTrackedDiffs( const parsedFile = parsedTrackedByPath.get(change.path); if (parsedFile) { const file = await buildHighlightedTrackedDiffFile({ + signal: input.signal, cwd, change, parsedFile, @@ -3295,6 +3317,7 @@ async function appendStructuredTrackedDiffs( } interface ProcessUntrackedChangeInput { + signal: AbortSignal; cwd: string; change: CheckoutFileChange; ignoreWhitespace: boolean; @@ -3345,7 +3368,7 @@ async function processUntrackedChange(input: ProcessUntrackedChangeInput): Promi } appendDiff(text); - const parsed = await parseAndHighlightDiff(text, cwd); + const parsed = await parseAndHighlightDiff(text, cwd, { signal: input.signal }); const parsedFile = parsed[0] ?? ({ @@ -3537,10 +3560,31 @@ async function prepareAllChangesIndex(cwd: string, baseRef: string) { } } -export async function getCheckoutDiff( +export function getCheckoutDiff( cwd: string, compare: CheckoutDiffCompare, context?: CheckoutContext, + signal?: AbortSignal, +): Promise { + return diffHighlightRequests.run( + JSON.stringify([ + "checkout", + resolve(cwd), + compare.mode, + compare.mode !== "uncommitted" ? (compare.baseRef ?? null) : null, + compare.ignoreWhitespace === true, + compare.includeStructured === true, + ]), + (highlightSignal) => loadCheckoutDiff(cwd, compare, context, highlightSignal), + signal, + ); +} + +async function loadCheckoutDiff( + cwd: string, + compare: CheckoutDiffCompare, + context: CheckoutContext | undefined, + signal: AbortSignal, ): Promise { cwd = await requireGitWorktreeRoot(cwd); @@ -3552,10 +3596,15 @@ export async function getCheckoutDiff( const index = compare.mode === "all" ? await prepareAllChangesIndex(cwd, refsForDiff.baseRef) : null; try { - return await readCheckoutDiff(cwd, compare, { - ...refsForDiff, - indexEnv: index ? { GIT_INDEX_FILE: index.indexFile } : undefined, - }); + return await readCheckoutDiff( + cwd, + compare, + { + ...refsForDiff, + indexEnv: index ? { GIT_INDEX_FILE: index.indexFile } : undefined, + }, + signal, + ); } finally { await index?.cleanup(); } @@ -3565,6 +3614,7 @@ async function readCheckoutDiff( cwd: string, compare: CheckoutDiffCompare, refsForDiff: CheckoutDiffRefs, + signal: AbortSignal, ): Promise { const ignoreWhitespace = compare.ignoreWhitespace === true; let effectiveRefsForDiff = refsForDiff; @@ -3625,6 +3675,7 @@ async function readCheckoutDiff( if (compare.includeStructured) { const didAppendTrackedDiffs = await appendStructuredTrackedDiffs({ + signal, cwd, trackedChanges, trackedNumstatByPath: trackedDiff.trackedNumstatByPath, @@ -3652,6 +3703,7 @@ async function readCheckoutDiff( break; } const didAppendUntrackedDiff = await processUntrackedChange({ + signal, cwd, change, ignoreWhitespace,