fix(web): show math labels in Mermaid diagrams

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

@@ -0,0 +1,68 @@
+// @vitest-environment jsdom
+
+import { describe, expect, it } from "vite-plus/test";
+
+import { sanitizeMermaidSvg } from "./MermaidDiagram";
+
+// The label markup Mermaid's sequence diagrams emit for $$...$$ text.
+function noteSvg(label: string) {
+  return `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 100 100"><foreignObject height="34" width="35" x="57" y="85"><div style="width: fit-content;" xmlns="http://www.w3.org/1999/xhtml"><div style="display: flex; align-items: center; justify-content: center; white-space: nowrap;"><span class="katex">${label}</span></div></div></foreignObject></svg>`;
+}
+
+function sanitize(svg: string) {
+  const host = document.createElement("div");
+  host.innerHTML = sanitizeMermaidSvg(svg);
+  return host;
+}
+
+describe("sanitizeMermaidSvg", () => {
+  it("keeps the MathML Mermaid renders for math labels", () => {
+    const host = sanitize(
+      noteSvg(
+        '<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><mrow><mfrac><mi>d</mi><mrow><mi>d</mi><mi>x</mi></mrow></mfrac><mrow><mo fence="true">(</mo><mi mathvariant="bold">v</mi><mo fence="true">)</mo></mrow><msqrt><mn>2</mn></msqrt></mrow></math>',
+      ),
+    );
+
+    const math = host.querySelector(".katex > math");
+    expect(math?.namespaceURI).toBe("http://www.w3.org/1998/Math/MathML");
+    expect(math?.getAttribute("display")).toBe("block");
+    expect(math?.querySelector("mfrac")?.textContent).toBe("ddx");
+    expect(math?.querySelector("msqrt")?.textContent).toBe("2");
+    expect(math?.querySelector("mi[mathvariant='bold']")?.textContent).toBe("v");
+    expect(math?.querySelectorAll("mo[fence='true']")).toHaveLength(2);
+  });
+
+  it("keeps spaced math labels parseable as an SVG image", () => {
+    // KaTeX renders \text{total cost} with a no-break space.
+    const svg = sanitizeMermaidSvg(noteSvg("<math><mrow><mtext>total cost</mtext></mrow></math>"));
+
+    const image = new DOMParser().parseFromString(svg, "image/svg+xml");
+    expect(image.documentElement.localName).toBe("svg");
+    expect(image.querySelector("mtext")?.textContent).toBe("total cost");
+  });
+
+  it("strips links, scripts, and MathML outside what KaTeX emits", () => {
+    const host = sanitize(
+      noteSvg(
+        '<math href="javascript:alert(1)"><mi href="javascript:alert(1)" onclick="alert(1)">x</mi><maction actiontype="toggle"><mi>a</mi></maction><mglyph src="https://example.com/a.png"></mglyph><annotation-xml encoding="text/html"><img src="x" onerror="alert(1)"></annotation-xml><semantics><mi>y</mi><annotation encoding="application/x-tex">y</annotation></semantics></math>',
+      ),
+    );
+
+    const math = host.querySelector("math");
+    expect(math).not.toBeNull();
+    expect(math?.querySelector("mi")?.textContent).toBe("x");
+    expect(host.querySelector("[href], [onclick], [src], [onerror]")).toBeNull();
+    expect(host.querySelector("maction, mglyph, annotation-xml, annotation, semantics, img")).toBe(
+      null,
+    );
+  });
+
+  it("only accepts math as an HTML label and keeps its attributes on MathML", () => {
+    const host = sanitize(
+      '<svg xmlns="http://www.w3.org/2000/svg"><math><mi>x</mi></math><text mathvariant="bold" fence="true">label</text></svg>',
+    );
+
+    expect(host.querySelector("math, mi")).toBeNull();
+    expect(host.querySelector("text")?.getAttributeNames()).toEqual([]);
+  });
+});

apps/web/src/components/chat/MermaidDiagram.tsx

@@ -31,10 +31,70 @@ function loadMermaid(): Promise<Mermaid> {
 const REMOTE_CSS_URL = /url\(\s*(?!['"]?#)[^)]*\)/gi;
 let purifier: ReturnType<typeof DOMPurify> | null = null;
 
+// Mermaid renders $$...$$ labels with KaTeX's MathML output. These are KaTeX's
+// MathML element types, without annotation (Mermaid strips it), semantics (only
+// a wrapper, so its content is kept), and mglyph (a trusted-input image).
+const KATEX_MATHML_TAGS = new Set([
+  "math",
+  "menclose",
+  "mfrac",
+  "mi",
+  "mlabeledtr",
+  "mn",
+  "mo",
+  "mover",
+  "mpadded",
+  "mphantom",
+  "mroot",
+  "mrow",
+  "mspace",
+  "msqrt",
+  "mstyle",
+  "msub",
+  "msubsup",
+  "msup",
+  "mtable",
+  "mtd",
+  "mtext",
+  "mtr",
+  "munder",
+  "munderover",
+]);
+// Presentation attributes KaTeX sets on those elements. The existing profiles
+// already cover class, style, width, height, display, and xmlns.
+const KATEX_MATHML_ATTRIBUTES = new Set([
+  "accent",
+  "accentunder",
+  "columnalign",
+  "columnlines",
+  "columnspacing",
+  "depth",
+  "displaystyle",
+  "fence",
+  "largeop",
+  "linebreak",
+  "linethickness",
+  "lspace",
+  "mathbackground",
+  "mathcolor",
+  "mathsize",
+  "mathvariant",
+  "maxsize",
+  "minsize",
+  "notation",
+  "rowlines",
+  "rowspacing",
+  "rspace",
+  "scriptlevel",
+  "separator",
+  "stretchy",
+  "voffset",
+]);
+
 // Diagrams can come from untrusted PR descriptions, so strip anything that can
 // navigate, run script, or fetch remote content on top of Mermaid's own strict
 // sanitization. CSS keeps only local url(#id) references; label text is untouched.
-function sanitizeMermaidSvg(svg: string): string {
+export function sanitizeMermaidSvg(svg: string): string {
   if (!purifier) {
     purifier = DOMPurify(window);
     purifier.addHook("uponSanitizeElement", (node, data) => {
@@ -47,13 +107,20 @@ function sanitizeMermaidSvg(svg: string): string {
         data.attrValue = data.attrValue.replace(REMOTE_CSS_URL, "none");
     });
   }
-  return purifier.sanitize(svg, {
-    ADD_TAGS: ["foreignObject"],
+  const sanitized = purifier.sanitize(svg, {
+    // DOMPurify still only accepts <math> under an HTML parent, here the label
+    // <div> inside <foreignObject>, and checks every MathML child's namespace.
+    ADD_TAGS: ["foreignObject", ...KATEX_MATHML_TAGS],
+    ADD_ATTR: (attribute, tagName) =>
+      KATEX_MATHML_TAGS.has(tagName) && KATEX_MATHML_ATTRIBUTES.has(attribute),
     HTML_INTEGRATION_POINTS: { foreignobject: true },
     FORBID_ATTR: ["href", "xlink:href", "src", "srcset"],
     FORBID_TAGS: ["a", "img", "image", "script"],
     USE_PROFILES: { svg: true, svgFilters: true, html: true },
   });
+  // The expanded image parses this as XML, which has no &nbsp; entity. KaTeX
+  // writes spaces in math as U+00A0, which HTML serialization turns into &nbsp;.
+  return sanitized.replaceAll("&nbsp;", "&#160;");
 }
 
 // Mermaid also lazy-loads diagram chunks inside render(); losing the network
@@ -152,7 +219,11 @@ function mermaidImageUrl(svg: string): string {
   }
   element.setAttribute("xmlns", "http://www.w3.org/2000/svg");
   element.style.maxWidth = "none";
-  element.style.backgroundColor = getComputedStyle(document.body).backgroundColor;
+  // A standalone image inherits nothing from the page, including the text color
+  // that HTML labels such as math use.
+  const { backgroundColor, color } = getComputedStyle(document.body);
+  element.style.backgroundColor = backgroundColor;
+  element.style.color = color;
   if (expandedImageUrl) URL.revokeObjectURL(expandedImageUrl);
   expandedImageUrl = URL.createObjectURL(
     new Blob([new XMLSerializer().serializeToString(element)], { type: "image/svg+xml" }),