From 8078c532ceeee5cb951325031f216171507f3d5e Mon Sep 17 00:00:00 2001 From: Julius Marminge Date: Thu, 10 Sep 2026 23:50:32 -0700 Subject: [PATCH] perf(web): preserve completed code-line DOM while streaming (#11198) --- apps/web/package.json | 2 + apps/web/src/components/ChatMarkdown.test.tsx | 31 ++++++++++-- apps/web/src/components/ChatMarkdown.tsx | 43 ++++++++++++----- .../chat/HighlightedCodeLines.test.tsx | 26 ++++++++++ .../components/chat/HighlightedCodeLines.tsx | 48 +++++++++++++++++++ .../src/lib/incrementalHighlighting.test.ts | 27 +++++++---- apps/web/src/lib/incrementalHighlighting.ts | 12 ++--- pnpm-lock.yaml | 6 +++ 8 files changed, 164 insertions(+), 31 deletions(-) create mode 100644 apps/web/src/components/chat/HighlightedCodeLines.test.tsx create mode 100644 apps/web/src/components/chat/HighlightedCodeLines.tsx diff --git a/apps/web/package.json b/apps/web/package.json index 069a222ca1227e8332eb548dcf941d1d11feed96..16579e730f1bf5b2f362eb5abaf4fe50c1d562ac 100644 --- a/apps/web/package.json +++ b/apps/web/package.json @@ -34,6 +34,8 @@ "class-variance-authority": "^0.7.1", "culori": "^4.0.2", "effect": "catalog:", + "hast-util-to-html": "^9.0.5", + "hast-util-to-jsx-runtime": "^2.3.6", "heic-to": "^1.5.2", "jose": "catalog:", "jsonc-parser": "3.3.1", diff --git a/apps/web/src/components/ChatMarkdown.test.tsx b/apps/web/src/components/ChatMarkdown.test.tsx index d21960a28e341a6d03b14e075f029adc9fa7c494..adc3ddb77444b3ca5a8903188ac12e5cb2b16df1 100644 --- a/apps/web/src/components/ChatMarkdown.test.tsx +++ b/apps/web/src/components/ChatMarkdown.test.tsx @@ -110,13 +110,36 @@ describe("ChatMarkdown favicon privacy", () => { }); describe("ChatMarkdown streaming", () => { + it("does not retokenize completed lines when streaming finishes", async () => { + const highlighter = await getSyntaxHighlighterPromise("typescript"); + const highlight = vi.spyOn(highlighter, "codeToHast"); + vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); + let renderer: ReactTestRenderer | undefined; + const text = "```typescript\nconst completed = 1;\nconst current = 2;"; + try { + await act(async () => { + renderer = create(); + }); + expect(highlight).toHaveBeenCalled(); + highlight.mockClear(); + await act(async () => { + renderer!.update(); + }); + expect(highlight.mock.calls.every(([code]) => !code.includes("const completed"))).toBe(true); + } finally { + await act(async () => renderer?.unmount()); + vi.unstubAllGlobals(); + vi.restoreAllMocks(); + } + }); + it("recovers highlighting after a failed fence changes without resetting its controls", async () => { const highlighter = await getSyntaxHighlighterPromise("text"); - const codeToHtml = highlighter.codeToHtml.bind(highlighter); + const codeToHast = highlighter.codeToHast.bind(highlighter); let fail = true; - vi.spyOn(highlighter, "codeToHtml").mockImplementation((...args) => { + vi.spyOn(highlighter, "codeToHast").mockImplementation((...args) => { if (fail) throw new Error("Temporary highlighter failure"); - return codeToHtml(...args); + return codeToHast(...args); }); vi.spyOn(console, "error").mockImplementation(() => {}); vi.spyOn(console, "warn").mockImplementation(() => {}); @@ -156,7 +179,7 @@ describe("ChatMarkdown streaming", () => { it("preserves code controls and details without highlighting an unchanged fence again", async () => { const highlighter = await getSyntaxHighlighterPromise("text"); - const highlight = vi.spyOn(highlighter, "codeToHtml"); + const highlight = vi.spyOn(highlighter, "codeToHast"); const writeText = vi.fn(async (_text: string) => {}); vi.stubGlobal("navigator", { clipboard: { writeText } }); vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true); diff --git a/apps/web/src/components/ChatMarkdown.tsx b/apps/web/src/components/ChatMarkdown.tsx index 9be8b0a7b393a37c731457b05f039ee1d331b4fe..0dd85e9147dca07584dc0d5bb678ce4bf91672ae 100644 --- a/apps/web/src/components/ChatMarkdown.tsx +++ b/apps/web/src/components/ChatMarkdown.tsx @@ -69,6 +69,7 @@ import React, { } from "react"; import type { Components, Options as ReactMarkdownOptions } from "react-markdown"; import ReactMarkdown from "react-markdown"; +import { toHtml } from "hast-util-to-html"; import { createIncrementalMarkdownPlugin } from "../markdown-incremental"; import { defaultUrlTransform } from "react-markdown"; import rehypeRaw from "rehype-raw"; @@ -123,7 +124,8 @@ 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 { createIncrementalHighlightedDocument } from "../lib/incrementalHighlighting"; +import { HighlightedCodeLines } from "./chat/HighlightedCodeLines"; import { RenderErrorBoundary } from "./RenderErrorBoundary"; import { useTheme } from "../hooks/useTheme"; import { getClientSettings, useClientSettings } from "../hooks/useSettings"; @@ -1011,9 +1013,14 @@ function SuspenseShikiCodeBlock({ themeName, isStreaming, }: SuspenseShikiCodeBlockProps) { + const [hasStreamed, setHasStreamed] = useState(isStreaming); + if (isStreaming && !hasStreamed) setHasStreamed(true); const language = extractFenceLanguage(className); const cacheKey = createHighlightCacheKey(code, language, themeName); - const cachedHighlightedHtml = !isStreaming ? highlightedCodeCache.get(cacheKey) : null; + // Once lines are mounted individually, keep that renderer when streaming + // finishes so switching to cached HTML cannot clear an existing selection. + const cachedHighlightedHtml = + !isStreaming && !hasStreamed ? highlightedCodeCache.get(cacheKey) : null; if (cachedHighlightedHtml != null) { return ( @@ -1031,6 +1038,7 @@ function SuspenseShikiCodeBlock({ themeName={themeName} cacheKey={cacheKey} isStreaming={isStreaming} + preserveLines={isStreaming || hasStreamed} /> ); } @@ -1041,6 +1049,7 @@ interface UncachedShikiCodeBlockProps { themeName: DiffThemeName; cacheKey: string; isStreaming: boolean; + preserveLines: boolean; } function UncachedShikiCodeBlock({ @@ -1049,16 +1058,19 @@ function UncachedShikiCodeBlock({ themeName, cacheKey, isStreaming, + preserveLines, }: UncachedShikiCodeBlockProps) { const highlighter = use(getSyntaxHighlighterPromise(language)); const incrementalHighlight = useMemo( - () => (isStreaming ? createIncrementalHighlighter(highlighter, language, themeName) : null), - [highlighter, isStreaming, language, themeName], + () => + preserveLines ? createIncrementalHighlightedDocument(highlighter, language, themeName) : null, + [highlighter, preserveLines, language, themeName], ); - const highlightedHtml = useMemo(() => { + const highlighted = useMemo(() => { try { - return incrementalHighlight - ? incrementalHighlight(code) + if (incrementalHighlight) return incrementalHighlight(code); + return preserveLines + ? highlighter.codeToHast(code, { lang: language, theme: themeName }) : highlighter.codeToHtml(code, { lang: language, theme: themeName }); } catch (error) { // Log highlighting failures for debugging while falling back to plain text @@ -1067,22 +1079,29 @@ function UncachedShikiCodeBlock({ error instanceof Error ? error.message : error, ); // If highlighting fails for this language, render as plain text - return highlighter.codeToHtml(code, { lang: "text", theme: themeName }); + return preserveLines + ? highlighter.codeToHast(code, { lang: "text", theme: themeName }) + : highlighter.codeToHtml(code, { lang: "text", theme: themeName }); } - }, [code, highlighter, incrementalHighlight, language, themeName]); + }, [code, highlighter, incrementalHighlight, language, preserveLines, themeName]); useEffect(() => { if (!isStreaming) { + const highlightedHtml = typeof highlighted === "string" ? highlighted : toHtml(highlighted); highlightedCodeCache.set( cacheKey, highlightedHtml, estimateHighlightedSize(highlightedHtml, code), ); } - }, [cacheKey, code, highlightedHtml, isStreaming]); + }, [cacheKey, code, highlighted, isStreaming]); - return ( -
+ return typeof highlighted === "string" ? ( +
+ ) : ( +
+ +
); } diff --git a/apps/web/src/components/chat/HighlightedCodeLines.test.tsx b/apps/web/src/components/chat/HighlightedCodeLines.test.tsx new file mode 100644 index 0000000000000000000000000000000000000000..66af6dcc5ed0bf3774ebbc4f3a01d00d27f1adcf --- /dev/null +++ b/apps/web/src/components/chat/HighlightedCodeLines.test.tsx @@ -0,0 +1,26 @@ +import { getSharedHighlighter } from "@pierre/diffs"; +import { toHtml } from "hast-util-to-html"; +import { renderToStaticMarkup } from "react-dom/server"; +import { describe, expect, it } from "vite-plus/test"; + +import { createIncrementalHighlightedDocument } from "../../lib/incrementalHighlighting"; +import { HighlightedCodeLines } from "./HighlightedCodeLines"; + +describe("highlighted code lines", () => { + it("preserves Shiki HTML, including colors, escaping, whitespace, and blank lines", async () => { + const highlighter = await getSharedHighlighter({ + langs: ["typescript"], + themes: ["pierre-dark", "pierre-light"], + preferredHighlighter: "shiki-wasm", + }); + for (const theme of ["pierre-dark", "pierre-light"] as const) { + const highlight = createIncrementalHighlightedDocument(highlighter, "typescript", theme); + const code = + 'const html = "";\n\n/* multi\nline */\n\tconst x = 1;\n'; + for (let end = 0; end <= code.length; end++) { + const root = highlight(code.slice(0, end)); + expect(renderToStaticMarkup()).toBe(toHtml(root)); + } + } + }); +}); diff --git a/apps/web/src/components/chat/HighlightedCodeLines.tsx b/apps/web/src/components/chat/HighlightedCodeLines.tsx new file mode 100644 index 0000000000000000000000000000000000000000..f0f971681c109c3c6d0c810e6dc656c882772e5e --- /dev/null +++ b/apps/web/src/components/chat/HighlightedCodeLines.tsx @@ -0,0 +1,48 @@ +import type { DiffsHighlighter } from "@pierre/diffs"; +import { toHtml } from "hast-util-to-html"; +import { toJsxRuntime } from "hast-util-to-jsx-runtime"; +import { cloneElement, isValidElement, memo, type DOMAttributes } from "react"; +import { Fragment, jsx, jsxs } from "react/jsx-runtime"; + +type HighlightedRoot = ReturnType; +type HighlightedNode = HighlightedRoot["children"][number]; +const runtime = { Fragment, jsx, jsxs }; + +function elementShell(node: Extract) { + const element = toJsxRuntime({ ...node, children: [] }, runtime); + if (!isValidElement>(element)) { + throw new Error("Expected a highlighted code element"); + } + return element; +} + +const HighlightedLine = memo(function HighlightedLine({ node }: { node: HighlightedNode }) { + if (node.type !== "element") return toJsxRuntime(node, runtime); + return cloneElement(elementShell(node), { + dangerouslySetInnerHTML: { __html: toHtml({ type: "root", children: node.children }) }, + }); +}); + +/** Completed line nodes retain their identity in the incremental highlighter. + * Keep their DOM mounted too: replacing the entire pre makes the browser parse + * and resolve styles for thousands of unchanged token spans on each update. + */ +export function HighlightedCodeLines({ root }: { root: HighlightedRoot }) { + const pre = root.children[0]; + if (pre?.type !== "element" || pre.tagName !== "pre") return toJsxRuntime(root, runtime); + const code = pre.children[0]; + if (code?.type !== "element" || code.tagName !== "code") return toJsxRuntime(root, runtime); + return cloneElement( + elementShell(pre), + undefined, + cloneElement( + elementShell(code), + undefined, + code.children.map((node, index) => ( + // A line's position is stable as tokens and new lines are appended. + // oxlint-disable-next-line react/no-array-index-key + + )), + ), + ); +} diff --git a/apps/web/src/lib/incrementalHighlighting.test.ts b/apps/web/src/lib/incrementalHighlighting.test.ts index b61cfc3b48a5ee8e53e9de1ae890835c9c9b6f13..36c27cea14a221adfa2ab79afe973272f22e17d9 100644 --- a/apps/web/src/lib/incrementalHighlighting.test.ts +++ b/apps/web/src/lib/incrementalHighlighting.test.ts @@ -1,7 +1,8 @@ +import { toHtml } from "hast-util-to-html"; import { getSharedHighlighter } from "@pierre/diffs"; import { describe, expect, it } from "vite-plus/test"; -import { createIncrementalHighlighter } from "./incrementalHighlighting"; +import { createIncrementalHighlightedDocument } from "./incrementalHighlighting"; const samples = { typescript: "/* multi\nline comment */\nconst x = `template\n${1 + 2}`;\nconst re = /abc/;\n", @@ -29,10 +30,10 @@ describe("incremental code highlighting", () => { async (language, code) => { const highlighter = await highlighterPromise; for (const theme of ["pierre-dark", "pierre-light"] as const) { - const highlight = createIncrementalHighlighter(highlighter, language, theme); + const highlight = createIncrementalHighlightedDocument(highlighter, language, theme); for (let end = 0; end <= code.length; end++) { const text = code.slice(0, end); - expect(highlight(text), `${theme}, prefix ${end}`).toBe( + expect(toHtml(highlight(text)), `${theme}, prefix ${end}`).toBe( highlighter.codeToHtml(text, { lang: language, theme }), ); } @@ -42,7 +43,11 @@ describe("incremental code highlighting", () => { 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 highlight = createIncrementalHighlightedDocument( + highlighter, + "typescript", + "pierre-dark", + ); const inputs = [ "/* open\ncomment\n", "/* open\ncomment\n*/\nconst x = 1;", @@ -53,7 +58,7 @@ describe("incremental code highlighting", () => { "\n\n\nconst fresh = true;\n", ]; for (const text of inputs) { - expect(highlight(text)).toBe( + expect(toHtml(highlight(text))).toBe( highlighter.codeToHtml(text, { lang: "typescript", theme: "pierre-dark" }), ); } @@ -63,9 +68,9 @@ describe("incremental code highlighting", () => { "preserves %s without requesting grammar state", async (language) => { const highlighter = await highlighterPromise; - const highlight = createIncrementalHighlighter(highlighter, language, "pierre-dark"); + const highlight = createIncrementalHighlightedDocument(highlighter, language, "pierre-dark"); for (const text of ["plain\ntext", "\u001b[31mred\ncontinued", "\n"]) { - expect(highlight(text)).toBe( + expect(toHtml(highlight(text))).toBe( highlighter.codeToHtml(text, { lang: language, theme: "pierre-dark" }), ); } @@ -74,11 +79,15 @@ describe("incremental code highlighting", () => { it("preserves partial CRLF and CR line endings", async () => { const highlighter = await highlighterPromise; - const highlight = createIncrementalHighlighter(highlighter, "typescript", "pierre-dark"); + const highlight = createIncrementalHighlightedDocument( + 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( + expect(toHtml(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 index a2f90b6935c2a0482dc75903a909d86c6f4bf0e1..fef3598a21520c9386b093ac512029876c74e38c 100644 --- a/apps/web/src/lib/incrementalHighlighting.ts +++ b/apps/web/src/lib/incrementalHighlighting.ts @@ -14,7 +14,7 @@ function codeChildren(root: ReturnType) { * 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( +export function createIncrementalHighlightedDocument( highlighter: DiffsHighlighter, language: string, theme: DiffThemeName, @@ -29,7 +29,7 @@ export function createIncrementalHighlighter( } | undefined; - return (code: string): string => { + return (code: 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 ( @@ -37,7 +37,7 @@ export function createIncrementalHighlighter( ["text", "plaintext", "plain", "txt", "ansi"].includes(language) || code.includes("\r") ) { - return highlighter.codeToHtml(code, options); + return highlighter.codeToHast(code, options); } if (cached && !code.startsWith(cached.prefix)) cached = undefined; const end = code.lastIndexOf("\n") + 1; @@ -51,7 +51,7 @@ export function createIncrementalHighlighter( const state = highlighter.getLastGrammarState(root); if (!state) { cached = undefined; - return highlighter.codeToHtml(code, options); + return highlighter.codeToHast(code, options); } cached = { prefix: code.slice(0, end), @@ -60,8 +60,8 @@ export function createIncrementalHighlighter( }; } const prefix = cached; - if (!prefix) return highlighter.codeToHtml(code, options); - return highlighter.codeToHtml(code.slice(prefix.prefix.length), { + if (!prefix) return highlighter.codeToHast(code, options); + return highlighter.codeToHast(code.slice(prefix.prefix.length), { ...options, grammarState: prefix.state, transformers: [ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 704437983bacc58584b6b6c653dedcf6b207d08f..eb2ff380f6ce0f084c4fadf321fc7edd1daf0a42 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -637,6 +637,12 @@ importers: effect: specifier: 4.0.0-rc.112 version: 4.0.0-rc.112(patch_hash=8bef799f35729cf3465eb428196f617b434a7c9f658d56acb2874d74b21b98b2) + hast-util-to-html: + specifier: ^9.0.5 + version: 9.0.5 + hast-util-to-jsx-runtime: + specifier: ^2.3.6 + version: 2.3.6 heic-to: specifier: ^1.5.2 version: 1.5.2 -- 2.54.0