fix(frontend): improve chat math rendering (#3557)

* fix(frontend): improve chat math rendering

* fix(frontend): refine math rendering stability

* fix(frontend): address review feedback on math preprocessing

- Fix escaped-backslash blind spot: consume \\\\ as a unit so \\( and
  \\[ are not mis-interpreted as math delimiters when the backslash
  itself is escaped.
- Thread inInlineCode state across lines so multi-line backtick code
  spans are protected from delimiter conversion.
- Remove double math normalization: preprocessStreamdownMarkdown now
  only handles Mermaid; math normalization lives solely in
  ClipboardSafeStreamdown, preventing non-idempotent double passes.
- Wire parseIncompleteMarkdown to isLoading in MarkdownContent so
  Streamdown's streaming-incomplete-math handling is actually active
  during streaming.

* fix(frontend): address streamdown math review issues

* refactor(frontend): move streamdown preprocessing out of ai elements

frontend/src/app/layout.tsx

@@ -1,5 +1,5 @@
-import "@/styles/globals.css";
 import "katex/dist/katex.min.css";
+import "@/styles/globals.css";
 
 import { type Metadata } from "next";
 

frontend/src/components/ai-elements/streamdown.tsx

@@ -1,10 +1,9 @@
 "use client";
 
-import { Component, useMemo, type ComponentProps, type ReactNode } from "react";
+import { Component, type ComponentProps, type ReactNode } from "react";
 import { Streamdown } from "streamdown";
 
 import { installClipboardFallback } from "@/core/clipboard";
-import { capMarkdownNesting } from "@/core/streamdown/preprocess";
 
 export type ClipboardSafeStreamdownProps = ComponentProps<typeof Streamdown>;
 
@@ -58,19 +57,9 @@ export function ClipboardSafeStreamdown({
   children,
   ...props
 }: ClipboardSafeStreamdownProps) {
-  // Fast path for the dominant pathological inputs (deep ">" chains and deeply
-  // nested lists both blow up marked's recursive tokenizers) so the error
-  // boundary below rarely has to absorb a stack overflow — and never has to
-  // face the heap exhaustion the same lists cause on larger stacks, which it
-  // cannot catch.
-  const safeChildren = useMemo(
-    () =>
-      typeof children === "string" ? capMarkdownNesting(children) : children,
-    [children],
-  );
   return (
     <StreamdownFallbackBoundary raw={children}>
-      <Streamdown {...props}>{safeChildren}</Streamdown>
+      <Streamdown {...props}>{children}</Streamdown>
     </StreamdownFallbackBoundary>
   );
 }

frontend/src/components/workspace/artifacts/artifact-file-detail.tsx

@@ -19,7 +19,6 @@ import {
   ArtifactHeader,
   ArtifactTitle,
 } from "@/components/ai-elements/artifact";
-import { ClipboardSafeStreamdown } from "@/components/ai-elements/streamdown";
 import { Button } from "@/components/ui/button";
 import { Select, SelectItem } from "@/components/ui/select";
 import {
@@ -44,6 +43,7 @@ import { useI18n } from "@/core/i18n/hooks";
 import { findToolCallResult } from "@/core/messages/utils";
 import { installSkill } from "@/core/skills/api";
 import { streamdownPlugins } from "@/core/streamdown";
+import { SafeStreamdown } from "@/core/streamdown/components";
 import {
   canBrowserPreviewFile,
   checkCodeFile,
@@ -462,13 +462,13 @@ export function ArtifactFilePreview({
   if (language === "markdown") {
     return (
       <div className="size-full px-4">
-        <ClipboardSafeStreamdown
+        <SafeStreamdown
           className="size-full"
           {...streamdownPlugins}
           components={{ a: ArtifactLink }}
         >
           {content ?? ""}
-        </ClipboardSafeStreamdown>
+        </SafeStreamdown>
       </div>
     );
   }

frontend/src/components/workspace/messages/markdown-content.tsx

@@ -3,14 +3,12 @@
 import { useMemo } from "react";
 import type { AnchorHTMLAttributes } from "react";
 
-import {
-  MessageResponse,
-  type MessageResponseProps,
-} from "@/components/ai-elements/message";
+import { type ClipboardSafeStreamdownProps } from "@/components/ai-elements/streamdown";
 import {
   preprocessStreamdownMarkdown,
-  streamdownPlugins,
+  streamdownPluginsWithoutRawHtml,
 } from "@/core/streamdown";
+import { SafeMessageResponse } from "@/core/streamdown/components";
 import { cn } from "@/lib/utils";
 
 import { CitationLink } from "../citations/citation-link";
@@ -22,24 +20,30 @@ function isExternalUrl(href: string | undefined): boolean {
 export type MarkdownContentProps = {
   content: string;
   isLoading: boolean;
-  rehypePlugins: MessageResponseProps["rehypePlugins"];
+  rehypePlugins?: ClipboardSafeStreamdownProps["rehypePlugins"];
   className?: string;
-  remarkPlugins?: MessageResponseProps["remarkPlugins"];
-  components?: MessageResponseProps["components"];
+  remarkPlugins?: ClipboardSafeStreamdownProps["remarkPlugins"];
+  components?: ClipboardSafeStreamdownProps["components"];
 };
 
 /** Renders markdown content. */
 export function MarkdownContent({
   content,
+  isLoading,
   rehypePlugins,
   className,
-  remarkPlugins = streamdownPlugins.remarkPlugins,
+  remarkPlugins = streamdownPluginsWithoutRawHtml.remarkPlugins,
   components: componentsFromProps,
 }: MarkdownContentProps) {
   const normalizedContent = useMemo(
     () => preprocessStreamdownMarkdown(content),
     [content],
   );
+  const effectiveRehypePlugins = useMemo(() => {
+    const base = streamdownPluginsWithoutRawHtml.rehypePlugins ?? [];
+    const extra = rehypePlugins ?? [];
+    return [...base, ...extra] as ClipboardSafeStreamdownProps["rehypePlugins"];
+  }, [rehypePlugins]);
   const components = useMemo(() => {
     return {
       a: (props: AnchorHTMLAttributes<HTMLAnchorElement>) => {
@@ -71,13 +75,14 @@ export function MarkdownContent({
   if (!content) return null;
 
   return (
-    <MessageResponse
+    <SafeMessageResponse
       className={className}
       remarkPlugins={remarkPlugins}
-      rehypePlugins={rehypePlugins}
+      rehypePlugins={effectiveRehypePlugins}
       components={components}
+      parseIncompleteMarkdown={isLoading}
     >
       {normalizedContent}
-    </MessageResponse>
+    </SafeMessageResponse>
   );
 }

frontend/src/components/workspace/messages/message-list-item.tsx

@@ -14,7 +14,6 @@ import {
   type AnchorHTMLAttributes,
   type ImgHTMLAttributes,
 } from "react";
-import rehypeKatex from "rehype-katex";
 
 import { Loader } from "@/components/ai-elements/loader";
 import {
@@ -24,7 +23,6 @@ import {
 } from "@/components/ai-elements/message";
 import {
   Reasoning,
-  ReasoningContent,
   ReasoningTrigger,
 } from "@/components/ai-elements/reasoning";
 import { Task, TaskTrigger } from "@/components/ai-elements/task";
@@ -44,6 +42,7 @@ import {
   type FileInMessage,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
+import { SafeReasoningContent } from "@/core/streamdown/components";
 import { cn } from "@/lib/utils";
 
 import { CopyButton } from "../copy-button";
@@ -347,7 +346,7 @@ function MessageContent_({
           onTurnDurationChange={handleDurationChange}
         >
           <ReasoningTrigger />
-          <ReasoningContent>{reasoningContent}</ReasoningContent>
+          <SafeReasoningContent>{reasoningContent}</SafeReasoningContent>
         </Reasoning>
       </AIElementMessageContent>
     );
@@ -390,14 +389,14 @@ function MessageContent_({
           >
             <ReasoningTrigger hasContent={!!reasoningContent} />
             {reasoningContent && (
-              <ReasoningContent>{reasoningContent}</ReasoningContent>
+              <SafeReasoningContent>{reasoningContent}</SafeReasoningContent>
             )}
           </Reasoning>
         )}
       <MarkdownContent
         content={contentToDisplay}
         isLoading={isLoading}
-        rehypePlugins={[...rehypePlugins, [rehypeKatex, { output: "html" }]]}
+        rehypePlugins={rehypePlugins}
         className="my-3"
         components={components}
       />

frontend/src/components/workspace/messages/subtask-card.tsx

@@ -13,13 +13,13 @@ import {
   ChainOfThoughtStep,
 } from "@/components/ai-elements/chain-of-thought";
 import { Shimmer } from "@/components/ai-elements/shimmer";
-import { ClipboardSafeStreamdown } from "@/components/ai-elements/streamdown";
 import { Button } from "@/components/ui/button";
 import { ShineBorder } from "@/components/ui/shine-border";
 import { useI18n } from "@/core/i18n/hooks";
 import { hasToolCalls } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import { streamdownPluginsWithWordAnimation } from "@/core/streamdown";
+import { SafeStreamdown } from "@/core/streamdown/components";
 import { useSubtask } from "@/core/tasks/context";
 import { explainLastToolCall } from "@/core/tools/utils";
 import { cn } from "@/lib/utils";
@@ -126,12 +126,12 @@ export function SubtaskCard({
           {task.prompt && (
             <ChainOfThoughtStep
               label={
-                <ClipboardSafeStreamdown
+                <SafeStreamdown
                   {...streamdownPluginsWithWordAnimation}
                   components={{ a: CitationLink }}
                 >
                   {task.prompt}
-                </ClipboardSafeStreamdown>
+                </SafeStreamdown>
               }
             ></ChainOfThoughtStep>
           )}

frontend/src/components/workspace/settings/about-settings-page.tsx

@@ -1,9 +1,9 @@
 "use client";
 
-import { ClipboardSafeStreamdown } from "@/components/ai-elements/streamdown";
+import { SafeStreamdown } from "@/core/streamdown/components";
 
 import { aboutMarkdown } from "./about-content";
 
 export function AboutSettingsPage() {
-  return <ClipboardSafeStreamdown>{aboutMarkdown}</ClipboardSafeStreamdown>;
+  return <SafeStreamdown>{aboutMarkdown}</SafeStreamdown>;
 }

frontend/src/components/workspace/settings/memory-settings-page.tsx

@@ -11,7 +11,6 @@ import Link from "next/link";
 import { useDeferredValue, useId, useRef, useState } from "react";
 import { toast } from "sonner";
 
-import { ClipboardSafeStreamdown } from "@/components/ai-elements/streamdown";
 import { Button } from "@/components/ui/button";
 import {
   Dialog,
@@ -39,6 +38,7 @@ import type {
   MemoryFactPatchInput,
   UserMemory,
 } from "@/core/memory/types";
+import { SafeStreamdown } from "@/core/streamdown/components";
 import { streamdownPlugins } from "@/core/streamdown/plugins";
 import { pathOfThread } from "@/core/threads/utils";
 import { formatTimeAgo } from "@/core/utils/datetime";
@@ -639,12 +639,12 @@ export function MemorySettingsPage() {
                 <div className="text-muted-foreground mb-4 text-sm">
                   {summaryReadOnly}
                 </div>
-                <ClipboardSafeStreamdown
+                <SafeStreamdown
                   className="size-full min-w-0 [overflow-wrap:anywhere] [&>*:first-child]:mt-0 [&>*:last-child]:mb-0"
                   {...streamdownPlugins}
                 >
                   {summariesToMarkdown(memory, filteredSectionGroups, t)}
-                </ClipboardSafeStreamdown>
+                </SafeStreamdown>
               </div>
             ) : null}
 

frontend/src/core/streamdown/components.tsx

@@ -0,0 +1,48 @@
+"use client";
+
+import {
+  MessageResponse,
+  type MessageResponseProps,
+} from "@/components/ai-elements/message";
+import {
+  ReasoningContent,
+  type ReasoningContentProps,
+} from "@/components/ai-elements/reasoning";
+import {
+  ClipboardSafeStreamdown,
+  type ClipboardSafeStreamdownProps,
+} from "@/components/ai-elements/streamdown";
+
+import {
+  useSafeStreamdownChildren,
+  useSafeStreamdownMarkdown,
+} from "./safe-children";
+
+export function SafeStreamdown({
+  children,
+  ...props
+}: ClipboardSafeStreamdownProps) {
+  const safeChildren = useSafeStreamdownChildren(children);
+
+  return (
+    <ClipboardSafeStreamdown {...props}>{safeChildren}</ClipboardSafeStreamdown>
+  );
+}
+
+export function SafeMessageResponse({
+  children,
+  ...props
+}: MessageResponseProps) {
+  const safeChildren = useSafeStreamdownChildren(children);
+
+  return <MessageResponse {...props}>{safeChildren}</MessageResponse>;
+}
+
+export function SafeReasoningContent({
+  children,
+  ...props
+}: ReasoningContentProps) {
+  const safeChildren = useSafeStreamdownMarkdown(children);
+
+  return <ReasoningContent {...props}>{safeChildren}</ReasoningContent>;
+}

frontend/src/core/streamdown/index.ts

@@ -1,3 +1,5 @@
+export * from "./components";
 export * from "./mermaid";
 export * from "./preprocess";
 export * from "./plugins";
+export * from "./safe-children";

frontend/src/core/streamdown/plugins.ts

@@ -6,6 +6,12 @@ import type { StreamdownProps } from "streamdown";
 
 import { rehypeSplitWordsIntoSpans } from "../rehype";
 
+const katexOptions = {
+  output: "html",
+  throwOnError: false,
+  strict: false,
+} as const;
+
 export const streamdownPlugins = {
   remarkPlugins: [
     remarkGfm,
@@ -13,7 +19,7 @@ export const streamdownPlugins = {
   ] as StreamdownProps["remarkPlugins"],
   rehypePlugins: [
     rehypeRaw,
-    [rehypeKatex, { output: "html" }],
+    [rehypeKatex, katexOptions],
   ] as StreamdownProps["rehypePlugins"],
 };
 
@@ -23,16 +29,18 @@ export const streamdownPluginsWithWordAnimation = {
     [remarkMath, { singleDollarTextMath: true }],
   ] as StreamdownProps["remarkPlugins"],
   rehypePlugins: [
-    [rehypeKatex, { output: "html" }],
+    [rehypeKatex, katexOptions],
     rehypeSplitWordsIntoSpans,
   ] as StreamdownProps["rehypePlugins"],
 };
 
-// Plugins for reasoning/thinking content — derived from streamdownPlugins but without rehypeRaw,
-// to prevent LLM-hallucinated HTML tags (e.g. <simd>) from being rendered as DOM elements.
-export const reasoningPlugins = {
+export const streamdownPluginsWithoutRawHtml = {
   remarkPlugins: streamdownPlugins.remarkPlugins,
   rehypePlugins: streamdownPlugins.rehypePlugins?.filter(
     (p) => p !== rehypeRaw,
   ) as StreamdownProps["rehypePlugins"],
 };
+
+// Plugins for reasoning/thinking content — derived from streamdownPlugins but without rehypeRaw,
+// to prevent LLM-hallucinated HTML tags (e.g. <simd>) from being rendered as DOM elements.
+export const reasoningPlugins = streamdownPluginsWithoutRawHtml;

frontend/src/core/streamdown/preprocess.ts

@@ -100,6 +100,223 @@ export function capMarkdownNesting(markdown: string): string {
   return capListNesting(capBlockquoteNesting(markdown));
 }
 
+type MathDelimiter = {
+  close: "\\)" | "\\]";
+  replacement: "$" | "$$";
+};
+
+type DelimiterState = {
+  openBlock: MathDelimiter | null;
+  inlineCodeDelimiterLength: number | null;
+};
+
+function consumeBacktickRun(line: string, index: number): number {
+  let runLength = 0;
+  while (line[index + runLength] === "`") {
+    runLength += 1;
+  }
+  return runLength;
+}
+
+function convertLatexDelimitersInLine(
+  line: string,
+  state: DelimiterState,
+): { line: string; state: DelimiterState } {
+  let result = "";
+  let i = 0;
+  let inlineCodeDelimiterLength = state.inlineCodeDelimiterLength;
+  let currentBlock = state.openBlock;
+
+  while (i < line.length) {
+    if (line[i] === "`") {
+      const runLength = consumeBacktickRun(line, i);
+      result += line.slice(i, i + runLength);
+      if (!currentBlock) {
+        if (inlineCodeDelimiterLength === null) {
+          inlineCodeDelimiterLength = runLength;
+        } else if (runLength === inlineCodeDelimiterLength) {
+          inlineCodeDelimiterLength = null;
+        }
+      }
+      i += runLength;
+      continue;
+    }
+
+    const two = line.slice(i, i + 2);
+    const inInlineCode = inlineCodeDelimiterLength !== null;
+
+    // Consume escaped backslash as a unit — `\\` is never part of a math
+    // delimiter, so skip past both characters to avoid the second `\` being
+    // mis-paired with a following `(` or `[`.
+    if (two === "\\\\" && !inInlineCode) {
+      result += two;
+      i += 2;
+      continue;
+    }
+
+    // Close an open math block
+    if (!inInlineCode && currentBlock?.close === two) {
+      result += currentBlock.replacement;
+      currentBlock = null;
+      i += 2;
+      continue;
+    }
+
+    // Open a new math block
+    if (!inInlineCode && !currentBlock && (two === "\\(" || two === "\\[")) {
+      const isDisplay = two === "\\[";
+      currentBlock = {
+        close: isDisplay ? "\\]" : "\\)",
+        replacement: isDisplay ? "$$" : "$",
+      };
+      result += currentBlock.replacement;
+      i += 2;
+      continue;
+    }
+
+    result += line[i];
+    i += 1;
+  }
+
+  return {
+    line: result,
+    state: { openBlock: currentBlock, inlineCodeDelimiterLength },
+  };
+}
+
+/**
+ * Normalize common LLM LaTeX delimiters for remark-math.
+ *
+ * remark-math recognizes `$...$` and `$$...$$`, but many models output
+ * `\(...\)` and `\[...\]`. Convert those delimiters outside fenced/indented
+ * code so KaTeX can render equations without corrupting code blocks. The
+ * conversion is stateful across lines, because display math normally spans
+ * several lines:
+ *
+ *   \[
+ *   ...
+ *   \]
+ */
+export function normalizeLatexMathDelimiters(markdown: string): string {
+  if (!/[\\][([\])]/.test(markdown)) {
+    return markdown;
+  }
+
+  let insideFence = false;
+  let mathState: DelimiterState = {
+    openBlock: null,
+    inlineCodeDelimiterLength: null,
+  };
+
+  return markdown
+    .split("\n")
+    .map((line) => {
+      if (CODE_FENCE_RE.test(line) && !mathState.openBlock) {
+        insideFence = !insideFence;
+        return line;
+      }
+      if (
+        insideFence ||
+        (INDENTED_CODE_RE.test(line) && !mathState.openBlock)
+      ) {
+        return line;
+      }
+      const converted = convertLatexDelimitersInLine(line, mathState);
+      mathState = converted.state;
+      return converted.line;
+    })
+    .join("\n");
+}
+
+function hasUnescapedTexComment(line: string): boolean {
+  for (let i = 0; i < line.length; i++) {
+    if (line[i] !== "%") {
+      continue;
+    }
+
+    let backslashCount = 0;
+    for (let j = i - 1; j >= 0 && line[j] === "\\"; j--) {
+      backslashCount += 1;
+    }
+
+    if (backslashCount % 2 === 0) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
+function flattenDisplayMathBody(lines: string[]): string[] {
+  if (lines.some(hasUnescapedTexComment)) {
+    return lines;
+  }
+
+  return [lines.map((line) => line.trim()).join(" ")];
+}
+
+/**
+ * Keep complete display-math blocks atomic for Streamdown.
+ *
+ * Streamdown first splits Markdown into render blocks with marked, then runs
+ * react-markdown on each block. Multi-line `$$ ... $$` can be split before
+ * remark-math sees the matching delimiters, especially in long numbered
+ * responses. Compacting the content between the opening and closing `$$`
+ * preserves the LaTeX semantics (visual line breaks still come from `\\`,
+ * `aligned`, `matrix`, `cases`, etc.) while keeping the display-math block
+ * atomic for Streamdown's splitter.
+ */
+export function compactDisplayMathBlocks(markdown: string): string {
+  if (!markdown.includes("$$")) {
+    return markdown;
+  }
+
+  const output: string[] = [];
+  let insideFence = false;
+  let mathLines: string[] | null = null;
+
+  for (const line of markdown.split("\n")) {
+    if (CODE_FENCE_RE.test(line) && mathLines === null) {
+      insideFence = !insideFence;
+      output.push(line);
+      continue;
+    }
+
+    if (insideFence || (INDENTED_CODE_RE.test(line) && mathLines === null)) {
+      output.push(line);
+      continue;
+    }
+
+    if (line.trim() === "$$") {
+      if (mathLines === null) {
+        mathLines = [];
+      } else {
+        const flattenedMathLines = flattenDisplayMathBody(mathLines);
+        output.push("$$", ...flattenedMathLines, "$$");
+        mathLines = null;
+      }
+      continue;
+    }
+
+    if (mathLines !== null) {
+      mathLines.push(line);
+      continue;
+    }
+
+    output.push(line);
+  }
+
+  if (mathLines !== null) {
+    output.push("$$", ...mathLines);
+  }
+
+  return output.join("\n");
+}
+
+export function normalizeStreamdownMathMarkdown(markdown: string): string {
+  return compactDisplayMathBlocks(normalizeLatexMathDelimiters(markdown));
+}
+
 export function preprocessStreamdownMarkdown(markdown: string): string {
   if (!MERMAID_BLOCK_HINT_RE.test(markdown) || !markdown.includes("-.->")) {
     return markdown;

frontend/src/core/streamdown/safe-children.ts

@@ -0,0 +1,34 @@
+import { useMemo } from "react";
+import type { ComponentProps } from "react";
+import { type Streamdown } from "streamdown";
+
+import {
+  capMarkdownNesting,
+  normalizeStreamdownMathMarkdown,
+} from "./preprocess";
+
+type StreamdownChildren = ComponentProps<typeof Streamdown>["children"];
+
+export function getSafeStreamdownMarkdown(markdown: string): string {
+  return normalizeStreamdownMathMarkdown(capMarkdownNesting(markdown));
+}
+
+export function getSafeStreamdownChildren(
+  children: StreamdownChildren,
+): StreamdownChildren {
+  if (typeof children !== "string") {
+    return children;
+  }
+
+  return getSafeStreamdownMarkdown(children);
+}
+
+export function useSafeStreamdownChildren(
+  children: StreamdownChildren,
+): StreamdownChildren {
+  return useMemo(() => getSafeStreamdownChildren(children), [children]);
+}
+
+export function useSafeStreamdownMarkdown(markdown: string): string {
+  return useMemo(() => getSafeStreamdownMarkdown(markdown), [markdown]);
+}

frontend/src/styles/globals.css

@@ -382,6 +382,41 @@
     -webkit-text-fill-color: transparent;
     text-fill-color: transparent;
   }
+
+  /* KaTeX rendering polish for chat messages.
+     Keep this narrowly scoped to KaTeX classes so normal Markdown text is not
+     affected. Display equations can be wider than the chat column; horizontal
+     scrolling is safer than wrapping/corrupting matrices, aligned equations, or
+     long PGD/DDPM expressions. */
+  .katex {
+    font-size: 1.02em;
+    line-height: 1.35;
+    text-rendering: auto;
+  }
+
+  .katex-display {
+    max-width: 100%;
+    overflow-x: auto;
+    padding: 0.35rem 0.125rem 0.45rem;
+    margin: 0.85rem 0;
+    text-align: center;
+  }
+
+  .katex-display > .katex {
+    display: inline-block;
+    max-width: 100%;
+    white-space: nowrap;
+    text-align: initial;
+  }
+
+  .katex-display::-webkit-scrollbar {
+    height: 6px;
+  }
+
+  .katex-display::-webkit-scrollbar-thumb {
+    background: color-mix(in oklch, var(--muted-foreground) 35%, transparent);
+    border-radius: 999px;
+  }
 }
 
 :root {

frontend/tests/unit/core/streamdown/preprocess.test.ts

@@ -4,6 +4,8 @@ import {
   capBlockquoteNesting,
   capListNesting,
   capMarkdownNesting,
+  compactDisplayMathBlocks,
+  normalizeStreamdownMathMarkdown,
   preprocessStreamdownMarkdown,
 } from "@/core/streamdown/preprocess";
 
@@ -100,7 +102,148 @@ test("capMarkdownNesting caps both blockquote and list nesting", () => {
   expect(/^[ \t]*/.exec(lines[1]!)![0].length).toBe(200);
 });
 
-test("preprocessStreamdownMarkdown leaves non-mermaid content unchanged", () => {
-  const input = "just some text";
-  expect(preprocessStreamdownMarkdown(input)).toBe(input);
+test("normalizeStreamdownMathMarkdown converts inline math delimiters", () => {
+  expect(
+    normalizeStreamdownMathMarkdown("Given \\(x\\), compute \\(x^2\\)."),
+  ).toBe("Given $x$, compute $x^2$.");
+});
+
+test("normalizeStreamdownMathMarkdown converts multiline display math delimiters", () => {
+  const input = [
+    "Before",
+    "\\[",
+    "\\begin{aligned}",
+    "x_t &= \\sqrt{\\bar{\\alpha}_t}x_0 + \\sqrt{1-\\bar{\\alpha}_t}\\epsilon, \\\\",
+    "\\hat{x}_0 &= x_t",
+    "\\end{aligned}",
+    "\\]",
+    "After",
+  ].join("\n");
+  const expected = [
+    "Before",
+    "$$",
+    "\\begin{aligned} x_t &= \\sqrt{\\bar{\\alpha}_t}x_0 + \\sqrt{1-\\bar{\\alpha}_t}\\epsilon, \\\\ \\hat{x}_0 &= x_t \\end{aligned}",
+    "$$",
+    "After",
+  ].join("\n");
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(expected);
+});
+
+test("normalizeStreamdownMathMarkdown leaves fenced and indented code untouched", () => {
+  const input = [
+    "Text \\(x\\)",
+    "```tex",
+    "\\[",
+    "x^2",
+    "\\]",
+    "```",
+    "    \\(literal\\)",
+  ].join("\n");
+  const expected = [
+    "Text $x$",
+    "```tex",
+    "\\[",
+    "x^2",
+    "\\]",
+    "```",
+    "    \\(literal\\)",
+  ].join("\n");
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(expected);
+});
+
+test("compactDisplayMathBlocks keeps display math as display math", () => {
+  const input = ["Before", "$$", "x", "=", "y", "$$", "After"].join("\n");
+  const expected = ["Before", "$$", "x = y", "$$", "After"].join("\n");
+  expect(compactDisplayMathBlocks(input)).toBe(expected);
+});
+
+test("compactDisplayMathBlocks preserves TeX comments in display math", () => {
+  const input = ["Before", "$$", "a % step 1", "+ b", "$$", "After"].join("\n");
+  expect(compactDisplayMathBlocks(input)).toBe(input);
+});
+
+test("compactDisplayMathBlocks compacts escaped percent in display math", () => {
+  const input = ["Before", "$$", "a \\% step 1", "+ b", "$$", "After"].join(
+    "\n",
+  );
+  const expected = ["Before", "$$", "a \\% step 1 + b", "$$", "After"].join(
+    "\n",
+  );
+  expect(compactDisplayMathBlocks(input)).toBe(expected);
+});
+
+test("compactDisplayMathBlocks leaves fenced code content untouched", () => {
+  const input = [
+    "```md",
+    "$$",
+    "x = y",
+    "$$",
+    "```",
+    "$$",
+    "a",
+    "=",
+    "b",
+    "$$",
+  ].join("\n");
+  const expected = [
+    "```md",
+    "$$",
+    "x = y",
+    "$$",
+    "```",
+    "$$",
+    "a = b",
+    "$$",
+  ].join("\n");
+  expect(compactDisplayMathBlocks(input)).toBe(expected);
+});
+
+test("preprocessStreamdownMarkdown applies only Mermaid fixes (not math)", () => {
+  const input = [
+    "Before \\(x\\)",
+    "```mermaid",
+    "graph TD",
+    "  A -.-> B",
+    "```",
+  ].join("\n");
+  const expected = [
+    "Before \\(x\\)",
+    "```mermaid",
+    "graph TD",
+    "  A -.-> B",
+    "```",
+  ].join("\n");
+  expect(preprocessStreamdownMarkdown(input)).toBe(expected);
+});
+
+test("normalizeStreamdownMathMarkdown preserves escaped backslash before parens", () => {
+  // When the backslash itself is escaped (\\), the following ( is not a math open
+  const input = "Use \\\\( to start inline math.";
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(
+    "Use \\\\( to start inline math.",
+  );
+});
+
+test("normalizeStreamdownMathMarkdown preserves escaped backslash before brackets", () => {
+  const input = "Escape: \\\\[ is not math.";
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(
+    "Escape: \\\\[ is not math.",
+  );
+});
+
+test("normalizeStreamdownMathMarkdown preserves delimiters inside multi-line code spans", () => {
+  // A backtick code span opened on line 1 should protect line 2 content
+  const input = ["`code span", "with \\(x\\) inside`"].join("\n");
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(input);
+});
+
+test("normalizeStreamdownMathMarkdown preserves delimiters inside multi-backtick code spans", () => {
+  const input = "Use ``\\(literal\\)`` here";
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(input);
+});
+
+test("normalizeStreamdownMathMarkdown requires matching backtick run to close code spans", () => {
+  const input = "Use ``\\(literal\\)` and still code`` then \\(x\\)";
+  const expected = "Use ``\\(literal\\)` and still code`` then $x$";
+  expect(normalizeStreamdownMathMarkdown(input)).toBe(expected);
 });