perf(web): revisiting a thread reuses its rendered Markdown instead of parsing again

apps/web/package.json

@@ -61,6 +61,7 @@
     "hast-util-to-html": "^9.0.5",
     "hast-util-to-jsx-runtime": "^2.3.6",
     "heic-to": "^1.5.2",
+    "html-url-attributes": "^3.0.1",
     "jose": "catalog:",
     "jsonc-parser": "3.3.1",
     "jszip": "3.10.1",
@@ -76,13 +77,15 @@
     "remark-gfm": "catalog:",
     "remark-math": "6.0.0",
     "remark-parse": "catalog:",
+    "remark-rehype": "catalog:",
     "saxes": "6.0.0",
     "sql.js": "1.15.0",
     "tailwind-merge": "^3.4.0",
     "tiff": "7.1.3",
     "unbash": "^5.0.0",
     "unified": "catalog:",
     "use-sync-external-store": "^1.6.0",
+    "vfile": "^6.0.3",
     "zustand": "^5.0.11"
   },
   "devDependencies": {
@@ -92,6 +95,7 @@
     "@tanstack/router-plugin": "^1.161.0",
     "@types/compression": "^1.8.1",
     "@types/culori": "^4.0.1",
+    "@types/hast": "^3.0.4",
     "@types/mdast": "^4.0.4",
     "@types/react": "catalog:",
     "@types/react-dom": "catalog:",

apps/web/src/components/ChatMarkdown.tsx

@@ -87,8 +87,8 @@ import type {
   ExtraProps as ReactMarkdownExtraProps,
   Options as ReactMarkdownOptions,
 } from "react-markdown";
-import ReactMarkdown from "react-markdown";
 import { createIncrementalMarkdownPlugin } from "../markdown-incremental";
+import type { Pluggable, PluggableList } from "unified";
 import { defaultUrlTransform } from "react-markdown";
 import { parseAssistantCitationHref } from "@t3tools/shared/assistantCitations";
 import { parseComposerContextHref } from "@t3tools/shared/composerContextReferences";
@@ -189,6 +189,7 @@ import {
 import { resolveLinkTarget } from "../browser/browserLinkTarget";
 import { PullRequestLinkPreview } from "./pullRequest/PullRequestLinkPreview";
 import { MarkdownCodeBlock, SuspenseShikiCodeBlock } from "./chat/MarkdownCodeBlock";
+import { createMarkdownRenderer } from "./chat/markdownRenderer";
 import { splitMarkdownBlocks } from "./chat/markdownStreamBlocks";
 import { MarkdownMermaidCodeBlock } from "./chat/MarkdownMermaidCodeBlock";
 import { MarkdownTable } from "./chat/MarkdownTable";
@@ -397,6 +398,14 @@ const CHAT_MARKDOWN_LITERAL_HTML_REHYPE_PLUGINS = [
   rehypeTextDirection,
 ] satisfies NonNullable<ReactMarkdownOptions["rehypePlugins"]>;
 
+function markdownUrlTransform(href: string) {
+  if (parseAssistantCitationHref(href)) return href;
+  if (parseComposerContextHref(href)) return href;
+  if (parseThreadLinkHref(href)) return href;
+  if (isWindowsDrivePathHref(href)) return href;
+  return rewriteMarkdownFileUriHref(href) ?? defaultUrlTransform(href);
+}
+
 const NO_MARKDOWN_EDITORS: ReadonlyArray<EditorId> = [];
 
 type MarkdownServerConfigFields = ReturnType<typeof markdownServerConfigFields>;
@@ -644,13 +653,6 @@ function useChatMarkdownState({
     });
   };
   const diffThemeName = resolveDiffThemeName(resolvedTheme);
-  const markdownUrlTransform = (href: string) => {
-    if (parseAssistantCitationHref(href)) return href;
-    if (parseComposerContextHref(href)) return href;
-    if (parseThreadLinkHref(href)) return href;
-    if (isWindowsDrivePathHref(href)) return href;
-    return rewriteMarkdownFileUriHref(href) ?? defaultUrlTransform(href);
-  };
   const handleCopy = (event: ReactClipboardEvent<HTMLElement>) => {
     const selection = window.getSelection();
     if (!selection || selection.isCollapsed || !event.clipboardData) return;
@@ -925,7 +927,6 @@ function useChatMarkdownState({
     componentState,
     handleCopy,
     markdownRef,
-    markdownUrlTransform,
     localMediaPreview,
     closeLocalMediaPreview,
   };
@@ -1633,29 +1634,62 @@ const CHAT_MARKDOWN_COMPONENTS = {
   },
 } satisfies Components;
 
+const REMARK_MATH: Pluggable = [remarkMath, { singleDollarTextMath: true }];
+// Shared lists, so every message with the same options shares a processor and its cache.
+const CHAT_MARKDOWN_REMARK_PLUGIN_SETS = {
+  plain: CHAT_MARKDOWN_REMARK_PLUGINS,
+  breaks: CHAT_MARKDOWN_REMARK_PLUGINS_WITH_BREAKS,
+  math: [REMARK_MATH, ...CHAT_MARKDOWN_REMARK_PLUGINS],
+  mathBreaks: [REMARK_MATH, ...CHAT_MARKDOWN_REMARK_PLUGINS_WITH_BREAKS],
+} satisfies Record<string, PluggableList>;
+
+const markdownRenderers = new WeakMap<
+  PluggableList,
+  WeakMap<PluggableList, ReturnType<typeof createMarkdownRenderer>>
+>();
+
+function markdownRenderer(remarkPlugins: PluggableList, rehypePlugins: PluggableList) {
+  let byRehype = markdownRenderers.get(remarkPlugins);
+  if (!byRehype) markdownRenderers.set(remarkPlugins, (byRehype = new WeakMap()));
+  let renderer = byRehype.get(rehypePlugins);
+  if (!renderer) {
+    renderer = createMarkdownRenderer({
+      remarkPlugins,
+      rehypePlugins,
+      components: CHAT_MARKDOWN_COMPONENTS,
+      urlTransform: markdownUrlTransform,
+    });
+    byRehype.set(rehypePlugins, renderer);
+  }
+  return renderer;
+}
+
 const MarkdownSource = memo(function MarkdownSource({
   text,
   remarkPlugins,
   rehypePlugins,
-  urlTransform,
+  shared,
+  cache,
 }: {
   text: string;
-  remarkPlugins: NonNullable<ReactMarkdownOptions["remarkPlugins"]>;
-  rehypePlugins: NonNullable<ReactMarkdownOptions["rehypePlugins"]>;
-  urlTransform: ReactMarkdownOptions["urlTransform"];
+  remarkPlugins: PluggableList;
+  rehypePlugins: PluggableList;
+  /** The plugins are one of the shared sets. */
+  shared: boolean;
+  /** The text will not change again; streamed text would only fill the cache with prefixes. */
+  cache: boolean;
 }) {
+  // Per-instance plugins, such as the incremental parser, set up their state on every render.
+  const render = shared
+    ? markdownRenderer(remarkPlugins, rehypePlugins)
+    : createMarkdownRenderer({
+        remarkPlugins,
+        rehypePlugins,
+        components: CHAT_MARKDOWN_COMPONENTS,
+        urlTransform: markdownUrlTransform,
+      });
   return (
-    <ChatMarkdownTextContext value={text}>
-      <ReactMarkdown
-        remarkPlugins={remarkPlugins}
-        rehypePlugins={rehypePlugins}
-        skipHtml={false}
-        components={CHAT_MARKDOWN_COMPONENTS}
-        urlTransform={urlTransform}
-      >
-        {text}
-      </ReactMarkdown>
-    </ChatMarkdownTextContext>
+    <ChatMarkdownTextContext value={text}>{render(text, shared && cache)}</ChatMarkdownTextContext>
   );
 });
 
@@ -1668,14 +1702,8 @@ function ChatMarkdown({
   extraRemarkPlugins = EMPTY_REMARK_PLUGINS,
   ...props
 }: ChatMarkdownProps) {
-  const {
-    componentState,
-    handleCopy,
-    markdownRef,
-    markdownUrlTransform,
-    localMediaPreview,
-    closeLocalMediaPreview,
-  } = useChatMarkdownState({ text, inline, ...props });
+  const { componentState, handleCopy, markdownRef, localMediaPreview, closeLocalMediaPreview } =
+    useChatMarkdownState({ text, inline, ...props });
   // Streaming re-renders only the last block; offsets for task edits need the whole text.
   const streamingBlocks =
     props.isStreaming === true &&
@@ -1693,13 +1721,26 @@ function ChatMarkdown({
   const bindMarkdownRef = (node: HTMLElement | null) => {
     markdownRef.current = node;
   };
-  const remarkPlugins = [
-    ...(lineBreaks ? CHAT_MARKDOWN_REMARK_PLUGINS_WITH_BREAKS : CHAT_MARKDOWN_REMARK_PLUGINS),
-    ...extraRemarkPlugins,
-    ...(incrementalParsing ? [createIncrementalMarkdownPlugin()] : []),
-  ];
-  if (componentState.mathEnabled)
-    remarkPlugins.unshift([remarkMath, { singleDollarTextMath: true }]);
+  const sharedRemarkPlugins =
+    CHAT_MARKDOWN_REMARK_PLUGIN_SETS[
+      componentState.mathEnabled
+        ? lineBreaks
+          ? "mathBreaks"
+          : "math"
+        : lineBreaks
+          ? "breaks"
+          : "plain"
+    ];
+  // Extra and incremental plugins are per instance, so their output is not shared.
+  const uncachedPlugins = extraRemarkPlugins.length > 0 || incrementalParsing;
+  const remarkPlugins = uncachedPlugins
+    ? [
+        ...sharedRemarkPlugins,
+        ...extraRemarkPlugins,
+        ...(incrementalParsing ? [createIncrementalMarkdownPlugin()] : []),
+      ]
+    : sharedRemarkPlugins;
+  const blocks = streamingBlocks ?? [text];
 
   return (
     <Root
@@ -1714,7 +1755,7 @@ function ChatMarkdown({
       onCopy={handleCopy}
     >
       <ChatMarkdownRendererContext value={componentState}>
-        {(streamingBlocks ?? [text]).map((source, index) => (
+        {blocks.map((source, index) => (
           <MarkdownSource
             // oxlint-disable-next-line react/no-array-index-key -- streamed blocks only change at the end
             key={index}
@@ -1725,7 +1766,8 @@ function ChatMarkdown({
                 ? CHAT_MARKDOWN_RENDER_REHYPE_PLUGINS
                 : CHAT_MARKDOWN_LITERAL_HTML_REHYPE_PLUGINS
             }
-            urlTransform={markdownUrlTransform}
+            shared={!uncachedPlugins}
+            cache={!props.isStreaming}
           />
         ))}
       </ChatMarkdownRendererContext>

apps/web/src/components/chat/markdownRenderer.test.tsx

@@ -0,0 +1,49 @@
+import {
+  CHAT_MARKDOWN_REHYPE_PLUGINS,
+  CHAT_MARKDOWN_REMARK_PLUGINS,
+} from "@t3tools/shared/markdownPipeline";
+import { renderToStaticMarkup } from "react-dom/server";
+import ReactMarkdown from "react-markdown";
+import { describe, expect, it } from "vite-plus/test";
+
+import { createMarkdownRenderer } from "./markdownRenderer";
+
+// Visible in the markup, so every URL attribute must pass through it.
+const urlTransform = (url: string) => `${url}#seen`;
+
+describe("createMarkdownRenderer", () => {
+  const render = createMarkdownRenderer({
+    remarkPlugins: CHAT_MARKDOWN_REMARK_PLUGINS,
+    rehypePlugins: CHAT_MARKDOWN_REHYPE_PLUGINS,
+    components: {},
+    urlTransform,
+  });
+
+  it.each([
+    '# Title\n\nText with `code`, **bold**, [a link](https://example.com "t") and https://bare.example.com.',
+    "![shot](./shot.png) and [file](file:///C:/repo/a.ts) and [bad](javascript:alert(1))",
+    "<details><summary>More</summary>\n\nHidden <b>bold</b> <script>x()</script>\n\n</details>",
+    "| a | b |\n| - | - |\n| <br> | [x](#frag) |\n\n- [ ] task\n- [x] done",
+    '```ts title="a.ts"\nconst a = 1;\n```\n\n> [!WARNING]\n> careful <a>bare</a>',
+    "Footnote[^1] and ~~strike~~\n\n[^1]: the note",
+  ])("renders what react-markdown renders: %j", (source) => {
+    const expected = renderToStaticMarkup(
+      <ReactMarkdown
+        remarkPlugins={CHAT_MARKDOWN_REMARK_PLUGINS}
+        rehypePlugins={CHAT_MARKDOWN_REHYPE_PLUGINS}
+        urlTransform={urlTransform}
+      >
+        {source}
+      </ReactMarkdown>,
+    );
+    const element = render(source, true);
+    expect(renderToStaticMarkup(element)).toBe(expected);
+    expect(render(source, true)).toBe(element);
+  });
+
+  it("keeps cached messages when one far larger than the cache is rendered", () => {
+    const element = render("Kept *message*", true);
+    render("x".repeat(600_000), true);
+    expect(render("Kept *message*", true)).toBe(element);
+  });
+});

apps/web/src/components/chat/markdownRenderer.ts

@@ -0,0 +1,105 @@
+import type { Element, Nodes, Root } from "hast";
+import { toJsxRuntime, type Components } from "hast-util-to-jsx-runtime";
+import { urlAttributes } from "html-url-attributes";
+import type { ReactElement } from "react";
+import { Fragment, jsx, jsxs } from "react/jsx-runtime";
+import remarkParse from "remark-parse";
+import remarkRehype from "remark-rehype";
+import { type PluggableList, unified } from "unified";
+import { VFile } from "vfile";
+
+export type MarkdownUrlTransform = (url: string, key: string, node: Element) => string;
+
+interface MarkdownRendererOptions {
+  readonly remarkPlugins: PluggableList;
+  readonly rehypePlugins: PluggableList;
+  readonly components: Partial<Components>;
+  readonly urlTransform: MarkdownUrlTransform;
+}
+
+// Elements are immutable, so a remount (a thread switch, a row scrolled back into view)
+// reuses one instead of parsing again. Parsing is most of the cost of rendering markdown.
+// An element and the hast it keeps hold about 20 bytes per source character, so ~10 MB.
+const MAX_CACHED_CHARS = 500_000;
+const rendered = new Map<string, ReactElement>();
+let cachedChars = 0;
+let nextRendererId = 0;
+
+// A short block still holds a whole element tree.
+const ENTRY_OVERHEAD_CHARS = 200;
+const entrySize = (key: string) => key.length + ENTRY_OVERHEAD_CHARS;
+
+function remember(key: string, element: ReactElement) {
+  // One huge message would otherwise evict everything else.
+  if (entrySize(key) > MAX_CACHED_CHARS / 4) return;
+  rendered.set(key, element);
+  cachedChars += entrySize(key);
+  for (const oldest of rendered.keys()) {
+    if (cachedChars <= MAX_CACHED_CHARS) break;
+    rendered.delete(oldest);
+    cachedChars -= entrySize(oldest);
+  }
+}
+
+/** Renders markdown as react-markdown does, with one processor per plugin set. */
+export function createMarkdownRenderer({
+  remarkPlugins,
+  rehypePlugins,
+  components,
+  urlTransform,
+}: MarkdownRendererOptions) {
+  const processor = unified()
+    .use(remarkParse)
+    .use(remarkPlugins)
+    .use(remarkRehype, { allowDangerousHtml: true })
+    .use(rehypePlugins)
+    .freeze();
+  const keyPrefix = `${nextRendererId++}\u0000`;
+
+  // Literal raw HTML that survived the rehype plugins, and URL attributes, as react-markdown treats them.
+  const prepare = (node: Nodes) => {
+    if (!("children" in node)) return;
+    node.children.forEach((child, index) => {
+      if (child.type === "raw") {
+        node.children[index] = { type: "text", value: child.value };
+        return;
+      }
+      if (child.type === "element") {
+        for (const [key, tags] of Object.entries(urlAttributes)) {
+          if (
+            Object.hasOwn(child.properties, key) &&
+            (tags === null || tags.includes(child.tagName))
+          ) {
+            child.properties[key] = urlTransform(String(child.properties[key] || ""), key, child);
+          }
+        }
+      }
+      prepare(child);
+    });
+  };
+
+  return (text: string, cache: boolean): ReactElement => {
+    const key = keyPrefix + text;
+    const cached = rendered.get(key);
+    if (cached !== undefined) {
+      // Refresh its place in the eviction order.
+      rendered.delete(key);
+      rendered.set(key, cached);
+      return cached;
+    }
+    const file = new VFile(text);
+    const tree = processor.runSync(processor.parse(file), file) as Root;
+    prepare(tree);
+    const element = toJsxRuntime(tree, {
+      Fragment,
+      components,
+      ignoreInvalidStyle: true,
+      jsx,
+      jsxs,
+      passKeys: true,
+      passNode: true,
+    });
+    if (cache) remember(key, element);
+    return element;
+  };
+}