perf(web): resume syntax highlighting from completed lines

apps/web/src/components/ChatMarkdown.tsx

@@ -123,6 +123,7 @@ import { fnv1a32 } from "../lib/diffRendering";
 import { LRUCache } from "../lib/lruCache";
 import { getSyntaxHighlighterPromise } from "../lib/syntaxHighlighting";
 import { GitHubIcon } from "./Icons";
+import { createIncrementalHighlighter } from "../lib/incrementalHighlighting";
 import { RenderErrorBoundary } from "./RenderErrorBoundary";
 import { useTheme } from "../hooks/useTheme";
 import { getClientSettings, useClientSettings } from "../hooks/useSettings";
@@ -1050,9 +1051,15 @@ function UncachedShikiCodeBlock({
   isStreaming,
 }: UncachedShikiCodeBlockProps) {
   const highlighter = use(getSyntaxHighlighterPromise(language));
+  const incrementalHighlight = useMemo(
+    () => (isStreaming ? createIncrementalHighlighter(highlighter, language, themeName) : null),
+    [highlighter, isStreaming, language, themeName],
+  );
   const highlightedHtml = useMemo(() => {
     try {
-      return highlighter.codeToHtml(code, { lang: language, theme: themeName });
+      return incrementalHighlight
+        ? incrementalHighlight(code)
+        : highlighter.codeToHtml(code, { lang: language, theme: themeName });
     } catch (error) {
       // Log highlighting failures for debugging while falling back to plain text
       console.warn(
@@ -1062,7 +1069,7 @@ function UncachedShikiCodeBlock({
       // If highlighting fails for this language, render as plain text
       return highlighter.codeToHtml(code, { lang: "text", theme: themeName });
     }
-  }, [code, highlighter, language, themeName]);
+  }, [code, highlighter, incrementalHighlight, language, themeName]);
 
   useEffect(() => {
     if (!isStreaming) {

apps/web/src/lib/incrementalHighlighting.test.ts

@@ -0,0 +1,86 @@
+import { getSharedHighlighter } from "@pierre/diffs";
+import { describe, expect, it } from "vite-plus/test";
+
+import { createIncrementalHighlighter } from "./incrementalHighlighting";
+
+const samples = {
+  typescript: "/* multi\nline comment */\nconst x = `template\n${1 + 2}`;\nconst re = /abc/;\n",
+  python: '#!/usr/bin/python\nx = """multi\nline"""\nprint(x)\n',
+  bash: "#!/bin/bash\ncat <<EOF\nhello\nEOF\necho hi\n",
+  html: "<script>\nconst x = 1;\n</script>\n<style>\np {color:red;}\n</style>\n",
+  markdown: "# heading\n\n```ts\nconst a = 1;\n```\n\ntext\n",
+  rust: 'fn main() {\n let x = r#"multi\nline"#;\n}\n',
+  tsx: 'const element = <div\n className="x">\n{value}\n</div>;\n',
+  json: '{\n "value": [1,\n 2, 3]\n}\n',
+  yaml: "key: |\n  multiline\n  value\nnext: true\n",
+  css: '/* comment\n continued */\np::before {\n content: "text";\n}\n',
+  sql: "SELECT 'multi\nline'\nFROM table_name;\n",
+} as const;
+
+const highlighterPromise = getSharedHighlighter({
+  langs: Object.keys(samples) as Array<keyof typeof samples>,
+  themes: ["pierre-dark", "pierre-light"],
+  preferredHighlighter: "shiki-wasm",
+});
+
+describe("incremental code highlighting", () => {
+  it.each(Object.entries(samples))(
+    "matches full HTML at every streaming prefix in %s",
+    async (language, code) => {
+      const highlighter = await highlighterPromise;
+      for (const theme of ["pierre-dark", "pierre-light"] as const) {
+        const highlight = createIncrementalHighlighter(highlighter, language, theme);
+        for (let end = 0; end <= code.length; end++) {
+          const text = code.slice(0, end);
+          expect(highlight(text), `${theme}, prefix ${end}`).toBe(
+            highlighter.codeToHtml(text, { lang: language, theme }),
+          );
+        }
+      }
+    },
+  );
+
+  it("resets after edits and truncation, including edits to a completed line", async () => {
+    const highlighter = await highlighterPromise;
+    const highlight = createIncrementalHighlighter(highlighter, "typescript", "pierre-dark");
+    const inputs = [
+      "/* open\ncomment\n",
+      "/* open\ncomment\n*/\nconst x = 1;",
+      "const edited = 2;\nconst x = 1;",
+      "const edited = 2;\nconst x = 10;",
+      "const edited = 2;\n",
+      "",
+      "\n\n\nconst fresh = true;\n",
+    ];
+    for (const text of inputs) {
+      expect(highlight(text)).toBe(
+        highlighter.codeToHtml(text, { lang: "typescript", theme: "pierre-dark" }),
+      );
+    }
+  });
+
+  it.each(["text", "plaintext", "plain", "txt", "ansi"])(
+    "preserves %s without requesting grammar state",
+    async (language) => {
+      const highlighter = await highlighterPromise;
+      const highlight = createIncrementalHighlighter(highlighter, language, "pierre-dark");
+      for (const text of ["plain\ntext", "\u001b[31mred\ncontinued", "\n"]) {
+        expect(highlight(text)).toBe(
+          highlighter.codeToHtml(text, { lang: language, theme: "pierre-dark" }),
+        );
+      }
+    },
+  );
+
+  it("preserves partial CRLF and CR line endings", async () => {
+    const highlighter = await highlighterPromise;
+    const highlight = createIncrementalHighlighter(highlighter, "typescript", "pierre-dark");
+    const code = "/* multi\r\nline */\r\nconst x = 1;\r\n";
+    for (let end = 0; end <= code.length; end++) {
+      const text = code.slice(0, end);
+      expect(highlight(text)).toBe(
+        highlighter.codeToHtml(text, { lang: "typescript", theme: "pierre-dark" }),
+      );
+    }
+  });
+});

apps/web/src/lib/incrementalHighlighting.ts

@@ -0,0 +1,74 @@
+import type { DiffsHighlighter } from "@pierre/diffs";
+
+import type { DiffThemeName } from "./diffRendering";
+
+function codeChildren(root: ReturnType<DiffsHighlighter["codeToHast"]>) {
+  const pre = root.children.find((node) => node.type === "element" && node.tagName === "pre");
+  if (pre?.type !== "element") throw new Error("Missing highlighted pre element");
+  const code = pre.children.find((node) => node.type === "element" && node.tagName === "code");
+  if (code?.type !== "element") throw new Error("Missing highlighted code element");
+  return code.children;
+}
+
+/** Resume tokenization after the last completed line. Keep its grammar state so
+ * multiline strings, comments, and embedded languages continue to highlight as
+ * they do in a full pass. The current line is always highlighted again.
+ */
+export function createIncrementalHighlighter(
+  highlighter: DiffsHighlighter,
+  language: string,
+  theme: DiffThemeName,
+) {
+  const options = { lang: language, theme };
+  const newline = { type: "text" as const, value: "\n" };
+  let cached:
+    | {
+        prefix: string;
+        state: ReturnType<DiffsHighlighter["getLastGrammarState"]>;
+        children: ReturnType<typeof codeChildren>;
+      }
+    | undefined;
+
+  return (code: string): string => {
+    // Plain text and ANSI do not have a TextMate grammar state. A CR at the end
+    // of a chunk can still become a CRLF, so keep that input on the full path.
+    if (
+      !language ||
+      ["text", "plaintext", "plain", "txt", "ansi"].includes(language) ||
+      code.includes("\r")
+    ) {
+      return highlighter.codeToHtml(code, options);
+    }
+    if (cached && !code.startsWith(cached.prefix)) cached = undefined;
+    const end = code.lastIndexOf("\n") + 1;
+    if (end > (cached?.prefix.length ?? 0)) {
+      // Omit the final newline: Shiki would tokenize an extra empty line and
+      // advance the grammar state twice before we process the following line.
+      const root = highlighter.codeToHast(code.slice(cached?.prefix.length ?? 0, end - 1), {
+        ...options,
+        ...(cached ? { grammarState: cached.state } : {}),
+      });
+      const state = highlighter.getLastGrammarState(root);
+      if (!state) {
+        cached = undefined;
+        return highlighter.codeToHtml(code, options);
+      }
+      cached = {
+        prefix: code.slice(0, end),
+        state,
+        children: [...(cached ? [...cached.children, newline] : []), ...codeChildren(root)],
+      };
+    }
+    const prefix = cached;
+    if (!prefix) return highlighter.codeToHtml(code, options);
+    return highlighter.codeToHtml(code.slice(prefix.prefix.length), {
+      ...options,
+      grammarState: prefix.state,
+      transformers: [
+        {
+          code: (node) => ({ ...node, children: [...prefix.children, newline, ...node.children] }),
+        },
+      ],
+    });
+  };
+}