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