feat(web): render LaTeX equations in chat

apps/web/package.json

@@ -53,17 +53,20 @@
     "jose": "catalog:",
     "jsonc-parser": "3.3.1",
     "jszip": "3.10.1",
+    "katex": "^0.16.47",
     "lucide": "^0.564.0",
     "lucide-react": "^0.564.0",
     "mermaid": "^11.17.2",
     "morphicons": "^1.7.1",
     "react": "19.2.6",
     "react-dom": "19.2.6",
     "react-markdown": "^10.1.0",
+    "rehype-katex": "^7.0.1",
     "rehype-raw": "^7.0.0",
     "rehype-sanitize": "^6.0.0",
     "remark-breaks": "^4.0.0",
     "remark-gfm": "^4.0.1",
+    "remark-math": "^6.0.0",
     "tailwind-merge": "^3.4.0",
     "unbash": "^5.0.0",
     "zustand": "^5.0.11"

apps/web/src/components/ChatMarkdown.test.tsx

@@ -3,6 +3,7 @@
 import { EnvironmentId, type AuthEnvironmentScope } from "@t3tools/contracts";
 import { act, type ComponentProps, type ReactNode } from "react";
 import { renderToStaticMarkup } from "react-dom/server";
+import { createRoot } from "react-dom/client";
 import { create, type ReactTestRenderer } from "react-test-renderer";
 import { describe, expect, it, vi } from "vite-plus/test";
 
@@ -84,6 +85,150 @@ function codeButton(renderer: ReactTestRenderer, label: string) {
   return button.props as ComponentProps<typeof Button>;
 }
 
+describe("ChatMarkdown equations", () => {
+  async function withMarkdown(
+    check: (
+      container: HTMLDivElement,
+      render: (props: ComponentProps<typeof ChatMarkdown>) => Promise<void>,
+    ) => Promise<void>,
+  ) {
+    vi.stubGlobal("IS_REACT_ACT_ENVIRONMENT", true);
+    const container = document.createElement("div");
+    document.body.append(container);
+    const root = createRoot(container);
+    const render = async (props: ComponentProps<typeof ChatMarkdown>) => {
+      await act(async () => root.render(<ChatMarkdown {...props} />));
+    };
+    try {
+      await check(container, render);
+    } finally {
+      await act(async () => root.unmount());
+      container.remove();
+      vi.unstubAllGlobals();
+    }
+  }
+
+  it.each([true, false])(
+    "renders inline and both display delimiters with parseRawHtml=%s",
+    async (parseRawHtml) => {
+      await withMarkdown(async (container, render) => {
+        await render({
+          cwd: undefined,
+          parseRawHtml,
+          text: String.raw`The objectives are $f_1$ and $f_2$.
+
+$$f_1(x) \le f_1(y), \quad f_2(x) \le f_2(y)$$
+
+$$
+\frac{P}{W}\bigg|_{\text{hover}} = \frac{\kappa_c}{\mathrm{FM}}\sqrt{\frac{\mathrm{DL}}{2\rho}}
+$$`,
+        });
+        expect(container.querySelectorAll(".katex")).toHaveLength(4);
+        expect(container.querySelectorAll(".katex-display")).toHaveLength(2);
+        expect(container.querySelector("math msub")?.textContent).toBe("f1");
+        expect(container.querySelector("math mfrac")).not.toBeNull();
+        expect(container.querySelector("math msqrt")).not.toBeNull();
+        expect(container.querySelectorAll(".katex-error")).toHaveLength(0);
+      });
+    },
+  );
+
+  it("leaves escaped dollars and code literal while rendering equations in lists", async () => {
+    await withMarkdown(async (container, render) => {
+      await render({
+        cwd: undefined,
+        lineBreaks: true,
+        text: String.raw`Costs are \$5 and \$10. Code: \`$f_1$\`.
+
+- $\rho$: air density
+- $\mathrm{DL} = W/A$: disk loading`.replaceAll("\\`", "`"),
+      });
+      expect(container.querySelector("p")?.textContent).toContain("Costs are $5 and $10.");
+      expect(container.querySelector("code")?.textContent).toBe("$f_1$");
+      expect(container.querySelectorAll("li .katex")).toHaveLength(2);
+    });
+  });
+
+  it("preserves prices and multiple skill mentions beside real equations", async () => {
+    await withMarkdown(async (container, render) => {
+      await render({
+        cwd: undefined,
+        skills: [
+          { name: "alpha", displayName: "Alpha" },
+          { name: "beta", displayName: "Beta" },
+        ],
+        text: String.raw`Use $alpha and $beta with $\rho$. Prices are $5, $10 and $20. Then $x^2$.`,
+      });
+      expect(
+        [...container.querySelectorAll("[data-markdown-copy]")].map((node) =>
+          node.getAttribute("data-markdown-copy"),
+        ),
+      ).toEqual(["$alpha", "$beta"]);
+      expect(container.textContent).toContain("Prices are $5, $10 and $20.");
+      expect(container.querySelectorAll(".katex")).toHaveLength(2);
+      expect(container.querySelector("math mi")?.textContent).toBe("ρ");
+    });
+  });
+
+  it("recovers from incomplete streamed equations and invalid LaTeX", async () => {
+    await withMarkdown(async (container, render) => {
+      for (const text of [
+        "Start $",
+        String.raw`Start $\frac{P}{`,
+        String.raw`Start $\frac{P}{W}$`,
+      ]) {
+        await render({ cwd: undefined, text, isStreaming: true });
+        expect(container.textContent).toContain("Start");
+      }
+      expect(container.querySelector("math mfrac")).not.toBeNull();
+      await render({ cwd: undefined, text: String.raw`$\frac{$` });
+      expect(container.querySelector(".katex-error")?.textContent).toBe(String.raw`\frac{`);
+      await render({ cwd: undefined, text: String.raw`$\rho$` });
+      expect(container.querySelector("math mi")?.textContent).toBe("ρ");
+    });
+  });
+
+  it("keeps raw HTML sanitization before math rendering", async () => {
+    await withMarkdown(async (container, render) => {
+      await render({
+        cwd: undefined,
+        text: String.raw`<img src="x" onerror="alert(1)"><script>alert(1)</script>
+
+$\href{javascript:alert(1)}{x}$`,
+      });
+      expect(container.querySelector("script, [onerror], a[href^='javascript:']")).toBeNull();
+    });
+  });
+
+  it.each([true, false])(
+    "rejects math commands that create links, images or HTML with parseRawHtml=%s",
+    async (parseRawHtml) => {
+      await withMarkdown(async (container, render) => {
+        await render({
+          cwd: undefined,
+          parseRawHtml,
+          text: [
+            "$x^2$",
+            String.raw`$$\href{https://example.com/unsafe}{click}$$`,
+            String.raw`$$\href{javascript:alert(1)}{click}$$`,
+            String.raw`$$\includegraphics{https://example.com/unsafe.png}$$`,
+            String.raw`$$\htmlClass{math-injection}{x}$$`,
+            String.raw`$$\htmlId{math-injection}{x}$$`,
+            String.raw`$$\htmlStyle{background-image:url(https://example.com/unsafe)}{x}$$`,
+            String.raw`$$\htmlData{math-injection=value}{x}$$`,
+          ].join("\n\n"),
+        });
+        expect(container.querySelector("math msup")?.textContent).toBe("x2");
+        expect(
+          container.querySelector(
+            "a, img, .math-injection, #math-injection, [data-math-injection], [style*='example.com']",
+          ),
+        ).toBeNull();
+      });
+    },
+  );
+});
+
 describe("ChatMarkdown bare anchor placeholders", () => {
   it.each(["<A>", "<a>", "<a >", "<a/>", "<A/>", "<a />"])(
     "preserves unmatched %s without linking later blocks",

apps/web/src/components/ChatMarkdown.tsx

@@ -84,7 +84,9 @@ import { toHtml } from "hast-util-to-html";
 import { createIncrementalMarkdownPlugin } from "../markdown-incremental";
 import { defaultUrlTransform } from "react-markdown";
 import rehypeRaw from "rehype-raw";
+import rehypeKatex from "rehype-katex";
 import rehypeSanitize, { defaultSchema } from "rehype-sanitize";
+import "katex/dist/katex.min.css";
 import remarkBreaks from "remark-breaks";
 import { parseAssistantCitationHref } from "@t3tools/shared/assistantCitations";
 import { parseComposerContextHref } from "@t3tools/shared/composerContextReferences";
@@ -154,6 +156,7 @@ import {
   serializeTableElementToMarkdown,
 } from "../markdown-clipboard";
 import { remarkNormalizeListItemIndentation } from "../markdown-list-indentation";
+import { remarkChatMath } from "../markdown-math";
 import {
   extractMarkdownLinkHrefs,
   isWindowsDrivePathHref,
@@ -527,7 +530,11 @@ const CHAT_MARKDOWN_SANITIZE_SCHEMA = {
   attributes: {
     ...defaultSchema.attributes,
     "*": (defaultSchema.attributes?.["*"] ?? []).filter((attribute) => attribute !== "title"),
-    code: [...(defaultSchema.attributes?.code ?? []), "dataCodeMeta", "dataInlineCode"],
+    code: [
+      ["className", /^language-./, "math-inline", "math-display"],
+      "dataCodeMeta",
+      "dataInlineCode",
+    ],
     blockquote: [...(defaultSchema.attributes?.blockquote ?? []), "dataAlert"],
     div: [...(defaultSchema.attributes?.div ?? []), ...CODEX_ARTIFACT_TEMPLATE_HAST_PROPERTIES],
     a: [...(defaultSchema.attributes?.a ?? []), "dataPullRequestAutolink"],
@@ -553,6 +560,7 @@ const CHAT_MARKDOWN_SANITIZE_SCHEMA = {
 
 const CHAT_MARKDOWN_REMARK_PLUGINS = [
   remarkGfm,
+  remarkChatMath,
   remarkKeepWindowsPathDestinations,
   remarkGithubAlerts,
   remarkNormalizeListItemIndentation,
@@ -563,6 +571,7 @@ const CHAT_MARKDOWN_REMARK_PLUGINS = [
 
 const CHAT_MARKDOWN_REMARK_PLUGINS_WITH_BREAKS = [
   remarkGfm,
+  remarkChatMath,
   remarkKeepWindowsPathDestinations,
   remarkGithubAlerts,
   remarkNormalizeListItemIndentation,
@@ -577,6 +586,12 @@ const CHAT_MARKDOWN_REHYPE_PLUGINS = [
   rehypeRaw,
   rehypePreserveImageSourceMeta,
   [rehypeSanitize, CHAT_MARKDOWN_SANITIZE_SCHEMA],
+  // Sanitize authored HTML before KaTeX generates its own HTML and MathML.
+  [rehypeKatex, { trust: false }],
+] satisfies NonNullable<ReactMarkdownOptions["rehypePlugins"]>;
+
+const CHAT_MARKDOWN_LITERAL_REHYPE_PLUGINS = [
+  [rehypeKatex, { trust: false }],
 ] satisfies NonNullable<ReactMarkdownOptions["rehypePlugins"]>;
 
 /** GitHub's own five alert kinds, in its colors: the glyph names the urgency, the title says it. */
@@ -3673,7 +3688,9 @@ function ChatMarkdown({
       <ChatMarkdownRendererContext value={componentState}>
         <ReactMarkdown
           remarkPlugins={remarkPlugins}
-          rehypePlugins={parseRawHtml ? CHAT_MARKDOWN_REHYPE_PLUGINS : undefined}
+          rehypePlugins={
+            parseRawHtml ? CHAT_MARKDOWN_REHYPE_PLUGINS : CHAT_MARKDOWN_LITERAL_REHYPE_PLUGINS
+          }
           skipHtml={false}
           components={CHAT_MARKDOWN_COMPONENTS}
           urlTransform={markdownUrlTransform}

apps/web/src/index.css

@@ -1750,6 +1750,18 @@ code {
   margin-bottom: 0;
 }
 
+.chat-markdown .katex-display {
+  max-width: 100%;
+  overflow-x: auto;
+  overflow-y: hidden;
+  padding-block: 0.15em;
+}
+
+.chat-markdown .katex {
+  overflow-wrap: normal;
+  word-break: normal;
+}
+
 .chat-markdown p,
 .chat-markdown ul,
 .chat-markdown ol,

apps/web/src/markdown-clipboard.ts

@@ -188,6 +188,17 @@ function serializeChildren(node: Node): string {
   return out;
 }
 
+/** KaTeX keeps the TeX source in MathML; partial selections may only retain its visible HTML. */
+function serializeMathElement(element: Element): string | null {
+  const display = element.classList.contains("katex-display");
+  if (!display && !element.classList.contains("katex")) return null;
+  const source = element.querySelector('annotation[encoding="application/x-tex"]')?.textContent;
+  if (source !== undefined && source !== null) {
+    return display ? `\n\n$$\n${source}\n$$\n\n` : `$${source}$`;
+  }
+  return element.querySelector(".katex-html")?.textContent ?? null;
+}
+
 function serializeNode(node: Node): string {
   if (node.nodeType === Node.TEXT_NODE) {
     const text = node.textContent ?? "";
@@ -203,6 +214,8 @@ function serializeNode(node: Node): string {
   }
   const markdownCopy = element.getAttribute("data-markdown-copy");
   if (markdownCopy !== null) return markdownCopy;
+  const math = serializeMathElement(element);
+  if (math !== null) return math;
   if (isSkippedElement(element)) return "";
 
   const headingLevel = /^H([1-6])$/.exec(element.tagName)?.[1];
@@ -288,6 +301,11 @@ function scanForSoleCodeBlock(node: Node, scan: SoleCodeBlockScan): void {
       if (markdownCopy.trim().length > 0) scan.other = true;
       continue;
     }
+    const math = serializeMathElement(element);
+    if (math !== null) {
+      if (math.trim().length > 0) scan.other = true;
+      continue;
+    }
     if (isSkippedElement(element)) continue;
     if (element.tagName === "PRE") {
       if (scan.pre) scan.other = true;
@@ -390,6 +408,23 @@ export function chatMarkdownClipboardPayload(
     if (range.collapsed) continue;
     const container = document.createElement("div");
     container.appendChild(range.cloneContents());
+    const endElement =
+      range.endContainer.nodeType === Node.ELEMENT_NODE
+        ? (range.endContainer as Element)
+        : range.endContainer.parentElement;
+    const endMath = endElement?.closest(".katex");
+    if (endMath) {
+      const mathRange = document.createRange();
+      mathRange.selectNodeContents(endMath);
+      // Hidden MathML precedes the visible equation, so a partial drag can
+      // include its complete annotation before reaching the final symbol.
+      if (range.compareBoundaryPoints(Range.END_TO_END, mathRange) < 0) {
+        const mathFragments = container.querySelectorAll(".katex");
+        mathFragments[mathFragments.length - 1]
+          ?.querySelector('annotation[encoding="application/x-tex"]')
+          ?.remove();
+      }
+    }
     const ancestor = range.commonAncestorContainer;
     const ancestorElement =
       ancestor.nodeType === Node.ELEMENT_NODE ? (ancestor as Element) : ancestor.parentElement;

apps/web/src/markdown-incremental.test.tsx

@@ -11,6 +11,7 @@ import { remarkCodexDirectives } from "@t3tools/client-runtime/codex-markdown-di
 import { remarkGithubAlerts } from "./markdown-github-alerts";
 import { createIncrementalMarkdownPlugin } from "./markdown-incremental";
 import { remarkNormalizeListItemIndentation } from "./markdown-list-indentation";
+import { remarkChatMath } from "./markdown-math";
 
 function render(source: string, incremental?: Plugin<[], Root>, parsedSources?: string[]) {
   let tree: Root | undefined;
@@ -32,6 +33,7 @@ function render(source: string, incremental?: Plugin<[], Root>, parsedSources?:
         observeParsing,
         capture,
         remarkGfm,
+        remarkChatMath,
         remarkGithubAlerts,
         remarkNormalizeListItemIndentation,
         remarkCodexDirectives,
@@ -73,6 +75,8 @@ describe("incremental Markdown parsing", () => {
     "```\na\n```\n\nnext\n\n~~~\nb\n~~~\n\nmore",
     "\n\n\tcode\n\nmore",
     "text <https://example.com> *bold*",
+    String.raw`Inline $f_1(x)$ and $$\frac{P}{W}$$`,
+    "$$f_1(x) \\le f_1(y)$$\n\n$$\n\\frac{P}{W}\n$$\n",
     "> [!NOTE]\n> alert\n\n- [ ] task",
     "\uFEFFtext after a byte-order mark",
   ])("preserves the parse tree, positions, and HTML while streaming %j", (tail) => {