authorgravatar for julius0216@outlook.comJulius Marminge <julius0216@outlook.com> 2026-09-10 23:50:32-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-09-10 23:50:32-07:00
log8078c532ceeee5cb951325031f216171507f3d5e
tree4ca76678909eae59d9a63dd05c19759629d34286
parentd7d7f8f3eb7bf593925281eb62348e15a91252b8
signature Signed by PGP key B5690EEEBB952194

perf(web): preserve completed code-line DOM while streaming (#11198)

8 files changed, 164 insertions(+), 31 deletions(-)

apps/web/package.json+2
......@@ -34,6 +34,8 @@
3434 "class-variance-authority": "^0.7.1",
3535 "culori": "^4.0.2",
3636 "effect": "catalog:",
37 "hast-util-to-html": "^9.0.5",
38 "hast-util-to-jsx-runtime": "^2.3.6",
3739 "heic-to": "^1.5.2",
3840 "jose": "catalog:",
3941 "jsonc-parser": "3.3.1",
apps/web/src/components/ChatMarkdown.test.tsx+27-4
......@@ -110,13 +110,36 @@ describe("ChatMarkdown favicon privacy", () => {
110110});
111111
112112describe("ChatMarkdown streaming", () => {
113 it("does not retokenize completed lines when streaming finishes", async () => {
114 const highlighter = await getSyntaxHighlighterPromise("typescript");
115 const highlight = vi.spyOn(highlighter, "codeToHast");
116 vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
117 let renderer: ReactTestRenderer | undefined;
118 const text = "```typescript\nconst completed = 1;\nconst current = 2;";
119 try {
120 await act(async () => {
121 renderer = create(<ChatMarkdown cwd="/tmp/project" text={text} isStreaming />);
122 });
123 expect(highlight).toHaveBeenCalled();
124 highlight.mockClear();
125 await act(async () => {
126 renderer!.update(<ChatMarkdown cwd="/tmp/project" text={text + "\n```"} />);
127 });
128 expect(highlight.mock.calls.every(([code]) => !code.includes("const completed"))).toBe(true);
129 } finally {
130 await act(async () => renderer?.unmount());
131 vi.unstubAllGlobals();
132 vi.restoreAllMocks();
133 }
134 });
135
113136 it("recovers highlighting after a failed fence changes without resetting its controls", async () => {
114137 const highlighter = await getSyntaxHighlighterPromise("text");
115 const codeToHtml = highlighter.codeToHtml.bind(highlighter);
138 const codeToHast = highlighter.codeToHast.bind(highlighter);
116139 let fail = true;
117 vi.spyOn(highlighter, "codeToHtml").mockImplementation((...args) => {
140 vi.spyOn(highlighter, "codeToHast").mockImplementation((...args) => {
118141 if (fail) throw new Error("Temporary highlighter failure");
119 return codeToHtml(...args);
142 return codeToHast(...args);
120143 });
121144 vi.spyOn(console, "error").mockImplementation(() => {});
122145 vi.spyOn(console, "warn").mockImplementation(() => {});
......@@ -156,7 +179,7 @@ describe("ChatMarkdown streaming", () => {
156179
157180 it("preserves code controls and details without highlighting an unchanged fence again", async () => {
158181 const highlighter = await getSyntaxHighlighterPromise("text");
159 const highlight = vi.spyOn(highlighter, "codeToHtml");
182 const highlight = vi.spyOn(highlighter, "codeToHast");
160183 const writeText = vi.fn(async (_text: string) => {});
161184 vi.stubGlobal("navigator", { clipboard: { writeText } });
162185 vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
apps/web/src/components/ChatMarkdown.tsx+31-12
......@@ -69,6 +69,7 @@ import React, {
6969} from "react";
7070import type { Components, Options as ReactMarkdownOptions } from "react-markdown";
7171import ReactMarkdown from "react-markdown";
72import { toHtml } from "hast-util-to-html";
7273import { createIncrementalMarkdownPlugin } from "../markdown-incremental";
7374import { defaultUrlTransform } from "react-markdown";
7475import rehypeRaw from "rehype-raw";
......@@ -123,7 +124,8 @@ import { fnv1a32 } from "../lib/diffRendering";
123124import { LRUCache } from "../lib/lruCache";
124125import { getSyntaxHighlighterPromise } from "../lib/syntaxHighlighting";
125126import { GitHubIcon } from "./Icons";
126import { createIncrementalHighlighter } from "../lib/incrementalHighlighting";
127import { createIncrementalHighlightedDocument } from "../lib/incrementalHighlighting";
128import { HighlightedCodeLines } from "./chat/HighlightedCodeLines";
127129import { RenderErrorBoundary } from "./RenderErrorBoundary";
128130import { useTheme } from "../hooks/useTheme";
129131import { getClientSettings, useClientSettings } from "../hooks/useSettings";
......@@ -1011,9 +1013,14 @@ function SuspenseShikiCodeBlock({
10111013 themeName,
10121014 isStreaming,
10131015}: SuspenseShikiCodeBlockProps) {
1016 const [hasStreamed, setHasStreamed] = useState(isStreaming);
1017 if (isStreaming && !hasStreamed) setHasStreamed(true);
10141018 const language = extractFenceLanguage(className);
10151019 const cacheKey = createHighlightCacheKey(code, language, themeName);
1016 const cachedHighlightedHtml = !isStreaming ? highlightedCodeCache.get(cacheKey) : null;
1020 // Once lines are mounted individually, keep that renderer when streaming
1021 // finishes so switching to cached HTML cannot clear an existing selection.
1022 const cachedHighlightedHtml =
1023 !isStreaming && !hasStreamed ? highlightedCodeCache.get(cacheKey) : null;
10171024
10181025 if (cachedHighlightedHtml != null) {
10191026 return (
......@@ -1031,6 +1038,7 @@ function SuspenseShikiCodeBlock({
10311038 themeName={themeName}
10321039 cacheKey={cacheKey}
10331040 isStreaming={isStreaming}
1041 preserveLines={isStreaming || hasStreamed}
10341042 />
10351043 );
10361044}
......@@ -1041,6 +1049,7 @@ interface UncachedShikiCodeBlockProps {
10411049 themeName: DiffThemeName;
10421050 cacheKey: string;
10431051 isStreaming: boolean;
1052 preserveLines: boolean;
10441053}
10451054
10461055function UncachedShikiCodeBlock({
......@@ -1049,16 +1058,19 @@ function UncachedShikiCodeBlock({
10491058 themeName,
10501059 cacheKey,
10511060 isStreaming,
1061 preserveLines,
10521062}: UncachedShikiCodeBlockProps) {
10531063 const highlighter = use(getSyntaxHighlighterPromise(language));
10541064 const incrementalHighlight = useMemo(
1055 () => (isStreaming ? createIncrementalHighlighter(highlighter, language, themeName) : null),
1056 [highlighter, isStreaming, language, themeName],
1065 () =>
1066 preserveLines ? createIncrementalHighlightedDocument(highlighter, language, themeName) : null,
1067 [highlighter, preserveLines, language, themeName],
10571068 );
1058 const highlightedHtml = useMemo(() => {
1069 const highlighted = useMemo(() => {
10591070 try {
1060 return incrementalHighlight
1061 ? incrementalHighlight(code)
1071 if (incrementalHighlight) return incrementalHighlight(code);
1072 return preserveLines
1073 ? highlighter.codeToHast(code, { lang: language, theme: themeName })
10621074 : highlighter.codeToHtml(code, { lang: language, theme: themeName });
10631075 } catch (error) {
10641076 // Log highlighting failures for debugging while falling back to plain text
......@@ -1067,22 +1079,29 @@ function UncachedShikiCodeBlock({
10671079 error instanceof Error ? error.message : error,
10681080 );
10691081 // If highlighting fails for this language, render as plain text
1070 return highlighter.codeToHtml(code, { lang: "text", theme: themeName });
1082 return preserveLines
1083 ? highlighter.codeToHast(code, { lang: "text", theme: themeName })
1084 : highlighter.codeToHtml(code, { lang: "text", theme: themeName });
10711085 }
1072 }, [code, highlighter, incrementalHighlight, language, themeName]);
1086 }, [code, highlighter, incrementalHighlight, language, preserveLines, themeName]);
10731087
10741088 useEffect(() => {
10751089 if (!isStreaming) {
1090 const highlightedHtml = typeof highlighted === "string" ? highlighted : toHtml(highlighted);
10761091 highlightedCodeCache.set(
10771092 cacheKey,
10781093 highlightedHtml,
10791094 estimateHighlightedSize(highlightedHtml, code),
10801095 );
10811096 }
1082 }, [cacheKey, code, highlightedHtml, isStreaming]);
1097 }, [cacheKey, code, highlighted, isStreaming]);
10831098
1084 return (
1085 <div className="chat-markdown-shiki" dangerouslySetInnerHTML={{ __html: highlightedHtml }} />
1099 return typeof highlighted === "string" ? (
1100 <div className="chat-markdown-shiki" dangerouslySetInnerHTML={{ __html: highlighted }} />
1101 ) : (
1102 <div className="chat-markdown-shiki">
1103 <HighlightedCodeLines root={highlighted} />
1104 </div>
10861105 );
10871106}
10881107
apps/web/src/components/chat/HighlightedCodeLines.test.tsx created+26
......@@ -0,0 +1,26 @@
1import { getSharedHighlighter } from "@pierre/diffs";
2import { toHtml } from "hast-util-to-html";
3import { renderToStaticMarkup } from "react-dom/server";
4import { describe, expect, it } from "vite-plus/test";
5
6import { createIncrementalHighlightedDocument } from "../../lib/incrementalHighlighting";
7import { HighlightedCodeLines } from "./HighlightedCodeLines";
8
9describe("highlighted code lines", () => {
10 it("preserves Shiki HTML, including colors, escaping, whitespace, and blank lines", async () => {
11 const highlighter = await getSharedHighlighter({
12 langs: ["typescript"],
13 themes: ["pierre-dark", "pierre-light"],
14 preferredHighlighter: "shiki-wasm",
15 });
16 for (const theme of ["pierre-dark", "pierre-light"] as const) {
17 const highlight = createIncrementalHighlightedDocument(highlighter, "typescript", theme);
18 const code =
19 'const html = "<img src=x onerror=alert(1)>";\n\n/* multi\nline */\n\tconst x = 1;\n';
20 for (let end = 0; end <= code.length; end++) {
21 const root = highlight(code.slice(0, end));
22 expect(renderToStaticMarkup(<HighlightedCodeLines root={root} />)).toBe(toHtml(root));
23 }
24 }
25 });
26});
apps/web/src/components/chat/HighlightedCodeLines.tsx created+48
......@@ -0,0 +1,48 @@
1import type { DiffsHighlighter } from "@pierre/diffs";
2import { toHtml } from "hast-util-to-html";
3import { toJsxRuntime } from "hast-util-to-jsx-runtime";
4import { cloneElement, isValidElement, memo, type DOMAttributes } from "react";
5import { Fragment, jsx, jsxs } from "react/jsx-runtime";
6
7type HighlightedRoot = ReturnType<DiffsHighlighter["codeToHast"]>;
8type HighlightedNode = HighlightedRoot["children"][number];
9const runtime = { Fragment, jsx, jsxs };
10
11function elementShell(node: Extract<HighlightedNode, { type: "element" }>) {
12 const element = toJsxRuntime({ ...node, children: [] }, runtime);
13 if (!isValidElement<DOMAttributes<HTMLElement>>(element)) {
14 throw new Error("Expected a highlighted code element");
15 }
16 return element;
17}
18
19const HighlightedLine = memo(function HighlightedLine({ node }: { node: HighlightedNode }) {
20 if (node.type !== "element") return toJsxRuntime(node, runtime);
21 return cloneElement(elementShell(node), {
22 dangerouslySetInnerHTML: { __html: toHtml({ type: "root", children: node.children }) },
23 });
24});
25
26/** Completed line nodes retain their identity in the incremental highlighter.
27 * Keep their DOM mounted too: replacing the entire pre makes the browser parse
28 * and resolve styles for thousands of unchanged token spans on each update.
29 */
30export function HighlightedCodeLines({ root }: { root: HighlightedRoot }) {
31 const pre = root.children[0];
32 if (pre?.type !== "element" || pre.tagName !== "pre") return toJsxRuntime(root, runtime);
33 const code = pre.children[0];
34 if (code?.type !== "element" || code.tagName !== "code") return toJsxRuntime(root, runtime);
35 return cloneElement(
36 elementShell(pre),
37 undefined,
38 cloneElement(
39 elementShell(code),
40 undefined,
41 code.children.map((node, index) => (
42 // A line's position is stable as tokens and new lines are appended.
43 // oxlint-disable-next-line react/no-array-index-key
44 <HighlightedLine key={index} node={node} />
45 )),
46 ),
47 );
48}
apps/web/src/lib/incrementalHighlighting.test.ts+18-9
......@@ -1,7 +1,8 @@
1import { toHtml } from "hast-util-to-html";
12import { getSharedHighlighter } from "@pierre/diffs";
23import { describe, expect, it } from "vite-plus/test";
34
4import { createIncrementalHighlighter } from "./incrementalHighlighting";
5import { createIncrementalHighlightedDocument } from "./incrementalHighlighting";
56
67const samples = {
78 typescript: "/* multi\nline comment */\nconst x = `template\n${1 + 2}`;\nconst re = /abc/;\n",
......@@ -29,10 +30,10 @@ describe("incremental code highlighting", () => {
2930 async (language, code) => {
3031 const highlighter = await highlighterPromise;
3132 for (const theme of ["pierre-dark", "pierre-light"] as const) {
32 const highlight = createIncrementalHighlighter(highlighter, language, theme);
33 const highlight = createIncrementalHighlightedDocument(highlighter, language, theme);
3334 for (let end = 0; end <= code.length; end++) {
3435 const text = code.slice(0, end);
35 expect(highlight(text), `${theme}, prefix ${end}`).toBe(
36 expect(toHtml(highlight(text)), `${theme}, prefix ${end}`).toBe(
3637 highlighter.codeToHtml(text, { lang: language, theme }),
3738 );
3839 }
......@@ -42,7 +43,11 @@ describe("incremental code highlighting", () => {
4243
4344 it("resets after edits and truncation, including edits to a completed line", async () => {
4445 const highlighter = await highlighterPromise;
45 const highlight = createIncrementalHighlighter(highlighter, "typescript", "pierre-dark");
46 const highlight = createIncrementalHighlightedDocument(
47 highlighter,
48 "typescript",
49 "pierre-dark",
50 );
4651 const inputs = [
4752 "/* open\ncomment\n",
4853 "/* open\ncomment\n*/\nconst x = 1;",
......@@ -53,7 +58,7 @@ describe("incremental code highlighting", () => {
5358 "\n\n\nconst fresh = true;\n",
5459 ];
5560 for (const text of inputs) {
56 expect(highlight(text)).toBe(
61 expect(toHtml(highlight(text))).toBe(
5762 highlighter.codeToHtml(text, { lang: "typescript", theme: "pierre-dark" }),
5863 );
5964 }
......@@ -63,9 +68,9 @@ describe("incremental code highlighting", () => {
6368 "preserves %s without requesting grammar state",
6469 async (language) => {
6570 const highlighter = await highlighterPromise;
66 const highlight = createIncrementalHighlighter(highlighter, language, "pierre-dark");
71 const highlight = createIncrementalHighlightedDocument(highlighter, language, "pierre-dark");
6772 for (const text of ["plain\ntext", "\u001b[31mred\ncontinued", "\n"]) {
68 expect(highlight(text)).toBe(
73 expect(toHtml(highlight(text))).toBe(
6974 highlighter.codeToHtml(text, { lang: language, theme: "pierre-dark" }),
7075 );
7176 }
......@@ -74,11 +79,15 @@ describe("incremental code highlighting", () => {
7479
7580 it("preserves partial CRLF and CR line endings", async () => {
7681 const highlighter = await highlighterPromise;
77 const highlight = createIncrementalHighlighter(highlighter, "typescript", "pierre-dark");
82 const highlight = createIncrementalHighlightedDocument(
83 highlighter,
84 "typescript",
85 "pierre-dark",
86 );
7887 const code = "/* multi\r\nline */\r\nconst x = 1;\r\n";
7988 for (let end = 0; end <= code.length; end++) {
8089 const text = code.slice(0, end);
81 expect(highlight(text)).toBe(
90 expect(toHtml(highlight(text))).toBe(
8291 highlighter.codeToHtml(text, { lang: "typescript", theme: "pierre-dark" }),
8392 );
8493 }
apps/web/src/lib/incrementalHighlighting.ts+6-6
......@@ -14,7 +14,7 @@ function codeChildren(root: ReturnType<DiffsHighlighter["codeToHast"]>) {
1414 * multiline strings, comments, and embedded languages continue to highlight as
1515 * they do in a full pass. The current line is always highlighted again.
1616 */
17export function createIncrementalHighlighter(
17export function createIncrementalHighlightedDocument(
1818 highlighter: DiffsHighlighter,
1919 language: string,
2020 theme: DiffThemeName,
......@@ -29,7 +29,7 @@ export function createIncrementalHighlighter(
2929 }
3030 | undefined;
3131
32 return (code: string): string => {
32 return (code: string) => {
3333 // Plain text and ANSI do not have a TextMate grammar state. A CR at the end
3434 // of a chunk can still become a CRLF, so keep that input on the full path.
3535 if (
......@@ -37,7 +37,7 @@ export function createIncrementalHighlighter(
3737 ["text", "plaintext", "plain", "txt", "ansi"].includes(language) ||
3838 code.includes("\r")
3939 ) {
40 return highlighter.codeToHtml(code, options);
40 return highlighter.codeToHast(code, options);
4141 }
4242 if (cached && !code.startsWith(cached.prefix)) cached = undefined;
4343 const end = code.lastIndexOf("\n") + 1;
......@@ -51,7 +51,7 @@ export function createIncrementalHighlighter(
5151 const state = highlighter.getLastGrammarState(root);
5252 if (!state) {
5353 cached = undefined;
54 return highlighter.codeToHtml(code, options);
54 return highlighter.codeToHast(code, options);
5555 }
5656 cached = {
5757 prefix: code.slice(0, end),
......@@ -60,8 +60,8 @@ export function createIncrementalHighlighter(
6060 };
6161 }
6262 const prefix = cached;
63 if (!prefix) return highlighter.codeToHtml(code, options);
64 return highlighter.codeToHtml(code.slice(prefix.prefix.length), {
63 if (!prefix) return highlighter.codeToHast(code, options);
64 return highlighter.codeToHast(code.slice(prefix.prefix.length), {
6565 ...options,
6666 grammarState: prefix.state,
6767 transformers: [
pnpm-lock.yaml+6
......@@ -637,6 +637,12 @@ importers:
637637 effect:
638638 specifier: 4.0.0-rc.112
639639 version: 4.0.0-rc.112(patch_hash=8bef799f35729cf3465eb428196f617b434a7c9f658d56acb2874d74b21b98b2)
640 hast-util-to-html:
641 specifier: ^9.0.5
642 version: 9.0.5
643 hast-util-to-jsx-runtime:
644 specifier: ^2.3.6
645 version: 2.3.6
640646 heic-to:
641647 specifier: ^1.5.2
642648 version: 1.5.2