fix(server): keep diff highlighting responsive and bounded

Tiny patches previously parsed both full source versions synchronously. Yield between parser work units and admit source by UTF-8 bytes, keeping large patches renderable without tokens. Cancel obsolete highlighting when a newer diff supersedes it.

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();
+  });
+});

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<void>; signal?: { readonly aborted: boolean } },
+): Promise<HighlightToken[][] | null> {
+  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;
+}

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";

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;
+});

packages/server/src/server/checkout-diff-cache.ts

@@ -4,6 +4,7 @@ import type { CheckoutDiffCompare, CheckoutDiffResult } from "../utils/checkout-
 interface ActiveRead {
   promise: Promise<CheckoutDiffResult>;
   invalidated: boolean;
+  controller: AbortController;
 }
 
 interface CachedDiff {
@@ -17,7 +18,7 @@ interface CheckoutDiffReadInput {
   compare: CheckoutDiffCompare;
   force?: boolean;
   reason?: string;
-  load: () => Promise<CheckoutDiffResult>;
+  load: (signal: AbortSignal) => Promise<CheckoutDiffResult>;
 }
 
 /** 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();
+      }
     }
   }
 }

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);

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 });
+  }
+});

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<string | null>;
   getNewFileContent?: (file: ParsedDiffFile) => Promise<string | null>;
 }
 
+/** Owns supersession for a whole diff, including sources that arrive after cancellation. */
+export class DiffHighlightRequests {
+  private readonly active = new Map<string, AbortController>();
+
+  async run<T>(
+    target: string,
+    operation: (signal: AbortSignal) => Promise<T>,
+    parentSignal?: AbortSignal,
+  ): Promise<T> {
+    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<number, HighlightToken[]> | null {
-  if (hasOversizedLine(fileContent)) return null;
+  signal?: AbortSignal,
+): Promise<Map<number, HighlightToken[]> | 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<number, HighlightToken[]>();
-  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<string | null | undefined> {
+  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<number, string>) => {
+    // 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<ParsedDiffFile> {
-  if (!isLanguageSupported(file.path) || hasOversizedDiffLine(file)) {
-    return file;
-  }
+  if (!isDiffHighlightEligible(file, options)) return file;
 
-  let newTokensByLine: Map<number, HighlightToken[]> | 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<Map<number, HighlightToken[]> | 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<number, HighlightToken[]>,
@@ -406,20 +510,24 @@ export async function parseAndHighlightDiff(
 ): Promise<ParsedDiffFile[]> {
   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;
 }
 

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,
+        ),
     });
   }
 

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<ReadonlyMap<
   return renames;
 }
 
+const diffHighlightRequests = new DiffHighlightRequests();
+
 /**
  * Fetches the unified diff of a single file as introduced by one commit and
  * parses it into the same {@link ParsedDiffFile} shape the diff subscription
@@ -2715,15 +2718,29 @@ async function readCommitRenames(cwd: string, sha: string): Promise<ReadonlyMap<
  * binary-only (no textual hunks). Throws on git failure (e.g. an unknown sha),
  * which the caller maps to a typed checkout error.
  */
-export async function getCommitFileDiff({
-  cwd,
-  sha,
-  path,
-}: {
+export function getCommitFileDiff(input: {
   cwd: string;
   sha: string;
   path: string;
 }): Promise<ParsedDiffFile | null> {
+  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<ParsedDiffFile | null> {
   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<string, FileStat>;
@@ -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<CheckoutDiffResult> {
+  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<CheckoutDiffResult> {
   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<CheckoutDiffResult> {
   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,