From 104e83109fd26bbe60206f349167f0245a611274 Mon Sep 17 00:00:00 2001
From: Tianyi Cui <53024+tianyicui@users.noreply.github.com>
Date: Sun, 26 Jul 2026 10:43:52 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20address=20review=20=E2=80=94=20plain=20f?=
 =?UTF-8?q?ences=20while=20streaming?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ds-review-bot finding: a growing fence retokenized on every chunk
(quadratic main-thread work). MarkdownText gains a streaming flag —
the streaming partial renders fences through the plain arm and the
finalize swap highlights once; AssistantMarkdown threads its existing
flag. (The zh Agent Note pair the review also flagged landed earlier
on this branch.) New spec pins plain-while-streaming and
highlighted-after-finalize.
---
 .../src/client/chat/AssistantMarkdown.tsx     |  2 +-
 .../src/markdown/MarkdownText.tsx             | 44 ++++++++++++-------
 .../ui-primitives/tests/markdown.spec.tsx     | 10 +++++
 3 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/packages/client/ui-conversation/src/client/chat/AssistantMarkdown.tsx b/packages/client/ui-conversation/src/client/chat/AssistantMarkdown.tsx
index 90eb3e3bee..0e91afcc07 100644
--- a/packages/client/ui-conversation/src/client/chat/AssistantMarkdown.tsx
+++ b/packages/client/ui-conversation/src/client/chat/AssistantMarkdown.tsx
@@ -44,7 +44,7 @@ export const AssistantMarkdown = memo(function AssistantMarkdown({ blocks, strea
     <div className={css.root} data-streaming={streaming || undefined}>
       {blocks.map((block, i) => {
         switch (block.kind) {
-          case 'text': return <MarkdownText key={i} text={block.text} />
+          case 'text': return <MarkdownText key={i} text={block.text} streaming={streaming} />
           case 'reasoning': return <ThinkRow key={i} text={block.text} running={streaming && i === last} />
           // Tool-call heads render as tool rows in the chat view's grouping pass.
           case 'tool-call': return null
diff --git a/packages/client/ui-primitives/src/markdown/MarkdownText.tsx b/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
index f74e939246..79ebc5f1b2 100644
--- a/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
+++ b/packages/client/ui-primitives/src/markdown/MarkdownText.tsx
@@ -24,7 +24,9 @@ function sanitizeUrl(url: string): string {
 
 const safeUrl: UrlTransform = url => sanitizeUrl(url)
 
-const components: Components = {
+/** Build the component table; while `streaming`, fences render the plain arm (see CodeBlock). */
+function buildComponents(streaming: boolean): Components {
+  return {
   a: ({ href = '', children }) => {
     const safeHref = sanitizeUrl(href)
     if (safeHref === '') return <>{children}</>
@@ -44,32 +46,40 @@ const components: Components = {
       <table>{children}</table>
     </div>
   ),
-  // Fenced blocks route through the shared CodeBlock (shiki for registered
-  // grammars, identical-geometry plain fallback for unknown/absent languages);
-  // inline code keeps the default <code> path (the :not(pre) rule styles it).
-  pre: ({ children }) => {
-    const child = isValidElement<{ className?: string; children?: unknown }>(children) ? children : undefined
-    const raw = child?.props.children
-    const text = typeof raw === 'string' ? raw : Array.isArray(raw) && typeof raw[0] === 'string' ? raw[0] : undefined
-    // A fence whose content isn't one plain string (never produced by the
-    // markdown pipeline) keeps the stock <pre> rather than guessing.
-    if (text === undefined) return <pre>{children}</pre>
-    const lang = /language-([\w-]+)/.exec(child?.props.className ?? '')?.[1]
-    return <CodeBlock code={text} lang={lang} />
-  },
+    // Fenced blocks route through the shared CodeBlock (shiki for registered
+    // grammars, identical-geometry plain fallback for unknown/absent
+    // languages); inline code keeps the default <code> path (the :not(pre)
+    // rule styles it). While the message streams, the fence renders the
+    // plain arm — retokenizing a growing fence on every chunk is quadratic
+    // main-thread work; the finalize swap highlights it once.
+    pre: ({ children }) => {
+      const child = isValidElement<{ className?: string; children?: unknown }>(children) ? children : undefined
+      const raw = child?.props.children
+      const text = typeof raw === 'string' ? raw : Array.isArray(raw) && typeof raw[0] === 'string' ? raw[0] : undefined
+      // A fence whose content isn't one plain string (never produced by the
+      // markdown pipeline) keeps the stock <pre> rather than guessing.
+      if (text === undefined) return <pre>{children}</pre>
+      const lang = /language-([\w-]+)/.exec(child?.props.className ?? '')?.[1]
+      return <CodeBlock code={text} lang={streaming ? undefined : lang} />
+    },
+  }
 }
 
+const staticComponents = buildComponents(false)
+const streamingComponents = buildComponents(true)
+
 /**
  * Render untrusted assistant-authored Markdown as semantic React elements.
- * @param props - Markdown source text preserved by the session projection.
+ * @param props - Markdown source text preserved by the session projection;
+ * `streaming` renders fences plain (highlighting lands on the finalize swap).
  * @returns A GFM document with raw HTML, relative links, unsafe protocols, and remote images disabled.
  */
-export function MarkdownText({ text }: { text: string }) {
+export function MarkdownText({ text, streaming = false }: { text: string; streaming?: boolean }) {
   return (
     <div className={css.markdown}>
       <ReactMarkdown
         remarkPlugins={remarkPlugins}
-        components={components}
+        components={streaming ? streamingComponents : staticComponents}
         urlTransform={safeUrl}
       >
         {text}
diff --git a/packages/client/ui-primitives/tests/markdown.spec.tsx b/packages/client/ui-primitives/tests/markdown.spec.tsx
index dcbf613005..1bd629a7d0 100644
--- a/packages/client/ui-primitives/tests/markdown.spec.tsx
+++ b/packages/client/ui-primitives/tests/markdown.spec.tsx
@@ -64,6 +64,16 @@ describe('MarkdownText', () => {
     expect(screen.getByRole('link', { name: 'https://deepseek.com' })).toBeTruthy()
   })
 
+  it('streaming renders fences plain; the finalize swap highlights them', () => {
+    const fence = '```ts\nconst answer = 42\n```'
+    const live = render(<MarkdownText text={fence} streaming />)
+    expect(live.container.querySelector('pre.shiki')).toBeNull()
+    expect(live.container.querySelector('pre code')?.textContent).toContain('const answer = 42')
+    live.unmount()
+    const done = render(<MarkdownText text={fence} />)
+    expect(done.container.querySelector('pre.shiki')).not.toBeNull()
+  })
+
   it('neutralizes raw HTML, unsafe or relative links, and remote images', () => {
     const markdown = [
       '<script>globalThis.compromised = true</script>',