diff --git a/apps/web/src/components/ChatMarkdown.tsx b/apps/web/src/components/ChatMarkdown.tsx index 88d34b8d58bc..9be8b0a7b393 100644 --- a/apps/web/src/components/ChatMarkdown.tsx +++ b/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) { diff --git a/apps/web/src/lib/incrementalHighlighting.test.ts b/apps/web/src/lib/incrementalHighlighting.test.ts new file mode 100644 index 000000000000..b61cfc3b48a5 --- /dev/null +++ b/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 <\nconst x = 1;\n\n\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 = \n{value}\n;\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, + 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" }), + ); + } + }); +}); diff --git a/apps/web/src/lib/incrementalHighlighting.ts b/apps/web/src/lib/incrementalHighlighting.ts new file mode 100644 index 000000000000..a2f90b6935c2 --- /dev/null +++ b/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) { + 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; + children: ReturnType; + } + | 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] }), + }, + ], + }); + }; +}