From 8ce4f07f923067d81c4efb0bd82faccfa8a3b1b6 Mon Sep 17 00:00:00 2001
From: ZiyaZhang <199893125+ZiyaZhang@users.noreply.github.com>
Date: Sat, 1 Aug 2026 00:58:50 -0700
Subject: [PATCH] fix(web): defer math rendering while streaming

---
 .../ui-primitives/src/markdown/MarkdownText.tsx   | 11 ++++++-----
 .../client/ui-primitives/tests/markdown.spec.tsx  | 15 +++++++++++++++
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/packages/client/ui-primitives/src/markdown/MarkdownText.tsx b/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
index 35d29ca720..a14f2c8a69 100644
--- a/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
+++ b/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
@@ -8,8 +8,9 @@ import { CodeBlock } from './CodeBlock.tsx'
 import 'katex/dist/katex.min.css'
 import css from './MarkdownText.module.css'
 
-const remarkPlugins = [remarkGfm, remarkMath]
-const rehypePlugins = [rehypeKatex]
+const streamingRemarkPlugins = [remarkGfm]
+const settledRemarkPlugins = [remarkGfm, remarkMath]
+const settledRehypePlugins = [rehypeKatex]
 
 function sanitizeUrl(url: string): string {
   try {
@@ -92,7 +93,7 @@ const streamingComponents = buildComponents(true)
 /**
  * Render untrusted assistant-authored Markdown as semantic React elements.
  * @param props - Markdown source text preserved by the session projection;
- * `streaming` renders fences plain (highlighting lands on the finalize swap);
+ * `streaming` renders fences and TeX plain (highlighting and KaTeX land on the finalize swap);
  * `codeLabels` forwards localized copy-button labels to fence CodeBlocks —
  * pass a reference-stable object (memoized per locale revision), because the
  * component table memoizes on its identity and a fresh literal per render
@@ -113,8 +114,8 @@ export function MarkdownText({ text, streaming = false, codeLabels }: {
   return (
     <div className={css.markdown}>
       <ReactMarkdown
-        remarkPlugins={remarkPlugins}
-        rehypePlugins={rehypePlugins}
+        remarkPlugins={streaming ? streamingRemarkPlugins : settledRemarkPlugins}
+        rehypePlugins={streaming ? undefined : settledRehypePlugins}
         components={components}
         urlTransform={safeUrl}
       >
diff --git a/packages/client/ui-primitives/tests/markdown.spec.tsx b/packages/client/ui-primitives/tests/markdown.spec.tsx
index 1431e2ae5b..030856a261 100644
--- a/packages/client/ui-primitives/tests/markdown.spec.tsx
+++ b/packages/client/ui-primitives/tests/markdown.spec.tsx
@@ -144,6 +144,21 @@ describe('MarkdownText', () => {
     expect(container.querySelector('.katex-display annotation')?.textContent).toContain('\\frac{\\partial \\mathbf{u}}')
     expect(container.querySelector('a')).toBeNull()
   })
+
+  it('defers TeX rendering while streaming so incomplete formulas never flash KaTeX errors', () => {
+    const partial = '$$\n\\frac{\\partial \\mathbf{u}}{\\partial'
+    const complete = '$$\n\\frac{\\partial \\mathbf{u}}{\\partial t}\n$$'
+    const live = render(<MarkdownText text={partial} streaming />)
+
+    expect(live.container.querySelector('.katex')).toBeNull()
+    expect(live.container.querySelector('.katex-error')).toBeNull()
+    expect(live.container.textContent).toContain('\\frac{\\partial \\mathbf{u}}{\\partial')
+
+    live.rerender(<MarkdownText text={complete} />)
+    expect(live.container.querySelectorAll('.katex')).toHaveLength(1)
+    expect(live.container.querySelectorAll('.katex-display')).toHaveLength(1)
+    expect(live.container.querySelector('.katex-error')).toBeNull()
+  })
 })
 
 describe('JsonBlock', () => {