Fix Mermaid diagrams that stay on error after replies finish (#26480)

* Fix Mermaid diagram recovery after streaming replies

* Clarify when the Mermaid render queue releases tasks

src/renderer/src/components/editor/MermaidBlock.test.tsx

@@ -0,0 +1,108 @@
+// @vitest-environment happy-dom
+import '@testing-library/jest-dom/vitest'
+import { act, cleanup, render, screen, waitFor } from '@testing-library/react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mermaid = vi.hoisted(() => ({
+  initialize: vi.fn(),
+  render: vi.fn<(id: string, content: string) => Promise<{ svg: string }>>(),
+  sanitize: vi.fn<(svg: string) => string>()
+}))
+vi.mock('mermaid', () => ({ default: mermaid }))
+vi.mock('dompurify', () => ({ default: { sanitize: mermaid.sanitize } }))
+
+import MermaidBlock from './MermaidBlock'
+
+const partial = 'flowchart TD\n A["Sync with main '
+const complete = `${partial}once"]`
+const diagram = '<svg><text>Sync with main once</text></svg>'
+
+beforeEach(() => {
+  mermaid.render.mockReset()
+  mermaid.render.mockResolvedValue({ svg: diagram })
+  mermaid.sanitize.mockReset().mockImplementation((svg) => svg)
+})
+afterEach(cleanup)
+
+describe('MermaidBlock rendering', () => {
+  it('replaces an invalid partial source error with the completed diagram', async () => {
+    mermaid.render.mockRejectedValueOnce(new Error('Parse error on line 2'))
+    const { container, rerender } = render(<MermaidBlock content={partial} isDark={false} />)
+    await screen.findByText(/Diagram error: Parse error on line 2/)
+    expect(container.querySelector('pre')?.textContent).toBe(partial)
+
+    rerender(<MermaidBlock content={complete} isDark={false} />)
+    await waitFor(() =>
+      expect(container.querySelector('svg')).toHaveTextContent('Sync with main once')
+    )
+    expect(screen.queryByText(/Diagram error:/)).not.toBeInTheDocument()
+    expect(container.querySelector('pre')).toBeNull()
+  })
+
+  it.each(['success', 'failure'])(
+    'ignores an obsolete render %s after content changes',
+    async (outcome) => {
+      const obsolete = Promise.withResolvers<{ svg: string }>()
+      const current = Promise.withResolvers<{ svg: string }>()
+      mermaid.render.mockReturnValueOnce(obsolete.promise).mockReturnValueOnce(current.promise)
+      const { container, rerender } = render(<MermaidBlock content={partial} isDark={false} />)
+      await waitFor(() => expect(mermaid.render).toHaveBeenCalledTimes(1))
+      rerender(<MermaidBlock content={complete} isDark={false} />)
+
+      await act(async () => {
+        if (outcome === 'success') {
+          obsolete.resolve({ svg: '<svg><text>Obsolete</text></svg>' })
+        } else {
+          obsolete.reject(new Error('Obsolete parse error'))
+        }
+        await obsolete.promise.catch(() => {})
+      })
+      await waitFor(() => expect(mermaid.render).toHaveBeenCalledTimes(2))
+      expect(container).not.toHaveTextContent('Obsolete')
+      await act(async () => current.resolve({ svg: diagram }))
+      await waitFor(() =>
+        expect(container.querySelector('svg')).toHaveTextContent('Sync with main once')
+      )
+      expect(container.querySelector('.mermaid-error')).toBeNull()
+    }
+  )
+
+  it('sanitizes the successful SVG before displaying it', async () => {
+    const unsafeSvg =
+      '<svg onload="alert(1)"><script>alert(1)</script><text>Safe diagram</text></svg>'
+    mermaid.render.mockResolvedValue({ svg: unsafeSvg })
+    mermaid.sanitize.mockReturnValueOnce('<svg><text>Safe diagram</text></svg>')
+    const { container } = render(<MermaidBlock content={complete} isDark={false} />)
+    await screen.findByText('Safe diagram')
+    expect(mermaid.sanitize).toHaveBeenCalledWith(unsafeSvg, { USE_PROFILES: { svg: true } })
+    expect(container.querySelector('script')).toBeNull()
+    expect(container.querySelector('svg')).not.toHaveAttribute('onload')
+  })
+
+  it('keeps new diagrams queued behind a render already in progress', async () => {
+    const first = Promise.withResolvers<{ svg: string }>()
+    const second = Promise.withResolvers<{ svg: string }>()
+    mermaid.render.mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise)
+    const { rerender } = render(
+      <>
+        <MermaidBlock content="first" isDark={false} />
+        <MermaidBlock content="second" isDark={false} />
+      </>
+    )
+    await waitFor(() => expect(mermaid.render).toHaveBeenCalledTimes(1))
+    await act(async () => first.resolve({ svg: diagram }))
+    await waitFor(() => expect(mermaid.render).toHaveBeenCalledTimes(2))
+
+    rerender(
+      <>
+        <MermaidBlock content="first" isDark={false} />
+        <MermaidBlock content="second" isDark={false} />
+        <MermaidBlock content="third" isDark={false} />
+      </>
+    )
+    await act(async () => {})
+    expect(mermaid.render).toHaveBeenCalledTimes(2)
+    await act(async () => second.resolve({ svg: diagram }))
+    await waitFor(() => expect(mermaid.render).toHaveBeenCalledTimes(3))
+  })
+})

src/renderer/src/components/editor/MermaidBlock.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useId, useRef, useState } from 'react'
+import React, { useEffect, useId, useState } from 'react'
 import type mermaidNamespace from 'mermaid'
 import DOMPurify from 'dompurify'
 import { getMermaidConfig } from './mermaid-config'
@@ -29,18 +29,21 @@ type MermaidBlockProps = {
 // where one render can clobber another's temporary DOM node. Serializing all
 // render calls through a single promise chain avoids this.
 //
-// The queue is replaced with a fresh promise after each render completes so
-// that old .then() closures (which capture containerRef, content, and id)
+// The queue is replaced with a fresh promise once all waiting renders complete so
+// that old .then() closures (which capture content and id)
 // become unreachable and can be GC'd. Without this, the chain grows with
 // every MermaidBlock mount/unmount cycle for the lifetime of the renderer.
 let renderQueue: Promise<void> = Promise.resolve()
 
 function enqueueRender(fn: () => Promise<void>): void {
-  renderQueue = renderQueue.then(fn, fn).then(() => {
+  const tail = renderQueue.then(fn, fn).then(() => {
     // Why: collapse the chain back to a single resolved promise so previous
     // closures do not remain reachable through a growing .then() chain.
-    renderQueue = Promise.resolve()
+    if (renderQueue === tail) {
+      renderQueue = Promise.resolve()
+    }
   })
+  renderQueue = tail
 }
 
 /**
@@ -53,8 +56,7 @@ export default function MermaidBlock({
   htmlLabels = false
 }: MermaidBlockProps): React.JSX.Element {
   const id = useId().replace(/:/g, '_')
-  const containerRef = useRef<HTMLDivElement>(null)
-  const [error, setError] = useState<string | null>(null)
+  const [result, setResult] = useState<{ svg: string } | { error: string } | null>(null)
 
   useEffect(() => {
     let cancelled = false
@@ -72,18 +74,17 @@ export default function MermaidBlock({
         // broken foreignObject label path again.
         mermaid.initialize(getMermaidConfig(isDark, htmlLabels))
         const { svg } = await mermaid.render(`mermaid-${id}`, content)
-        if (!cancelled && containerRef.current) {
+        if (!cancelled) {
           // Why: although mermaid uses DOMPurify internally, we add an explicit
           // sanitization pass as defense-in-depth against XSS in case upstream
           // behaviour changes or a mermaid version ships without sanitization.
-          containerRef.current.innerHTML = DOMPurify.sanitize(svg, {
-            USE_PROFILES: { svg: true }
+          setResult({
+            svg: DOMPurify.sanitize(svg, { USE_PROFILES: { svg: true } })
           })
-          setError(null)
         }
       } catch (err) {
         if (!cancelled) {
-          setError(err instanceof Error ? err.message : 'Invalid mermaid syntax')
+          setResult({ error: err instanceof Error ? err.message : 'Invalid mermaid syntax' })
           // Mermaid leaves an error element in the DOM on failure — clean it up.
           const errorEl = document.getElementById(`d${`mermaid-${id}`}`)
           errorEl?.remove()
@@ -99,11 +100,12 @@ export default function MermaidBlock({
     }
   }, [content, htmlLabels, isDark, id])
 
-  if (error) {
+  if (result && 'error' in result) {
     return (
       <div className="mermaid-block">
         <div className="mermaid-error">
-          {translate('auto.components.editor.MermaidBlock.dcc132e691', 'Diagram error:')} {error}
+          {translate('auto.components.editor.MermaidBlock.dcc132e691', 'Diagram error:')}{' '}
+          {result.error}
         </div>
         <pre>
           <code>{content}</code>
@@ -112,5 +114,10 @@ export default function MermaidBlock({
     )
   }
 
-  return <div className="mermaid-block" ref={containerRef} />
+  return (
+    <div
+      className="mermaid-block"
+      dangerouslySetInnerHTML={result ? { __html: result.svg } : undefined}
+    />
+  )
 }

src/renderer/src/components/native-chat/NativeChatMarkdown.mermaid.test.tsx

@@ -0,0 +1,70 @@
+// @vitest-environment happy-dom
+import '@testing-library/jest-dom/vitest'
+import { cleanup, render, screen, waitFor } from '@testing-library/react'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const mermaid = vi.hoisted(() => ({ initialize: vi.fn(), render: vi.fn() }))
+vi.mock('mermaid', () => ({ default: mermaid }))
+vi.mock('dompurify', () => ({ default: { sanitize: (svg: string) => svg } }))
+vi.mock('@/store', () => ({
+  useAppStore: (selector: (state: { settings: { theme: string } }) => unknown) =>
+    selector({ settings: { theme: 'light' } })
+}))
+
+import { NativeChatMarkdown } from './NativeChatMarkdown'
+import { NativeChatCodeBlock } from './NativeChatCodeBlock'
+import { TooltipProvider } from '@/components/ui/tooltip'
+
+const partial = 'flowchart TD\n A["Sync with main '
+const complete = `${partial}once"]`
+
+function reply(source: string, streaming: boolean): React.JSX.Element {
+  return (
+    <TooltipProvider>
+      <NativeChatMarkdown
+        content={`Here is the plan:\n\n\`\`\`mermaid\n${source}${streaming ? '' : '\n```'}`}
+        variant="document"
+        renderCodeBlock={NativeChatCodeBlock}
+        streaming={streaming}
+      />
+    </TooltipProvider>
+  )
+}
+
+beforeEach(() => {
+  mermaid.render.mockReset()
+  mermaid.render.mockResolvedValue({ svg: '<svg><text>Sync with main once</text></svg>' })
+})
+afterEach(cleanup)
+
+describe('native chat Mermaid fences', () => {
+  it('shows partial source while streaming and renders the diagram when the reply finishes', async () => {
+    const { container, rerender } = render(reply(partial, true))
+    expect(container.querySelector('[data-native-chat-code-content]')?.textContent).toContain(
+      partial
+    )
+    expect(container.querySelector('[data-code-language="mermaid"]')).toBeInTheDocument()
+    expect(container.querySelector('.mermaid-block')).toBeNull()
+    expect(mermaid.render).not.toHaveBeenCalled()
+
+    rerender(reply(complete, true))
+    expect(container.querySelector('pre')?.textContent).toContain(complete)
+    expect(mermaid.render).not.toHaveBeenCalled()
+
+    rerender(reply(complete, false))
+    await waitFor(() =>
+      expect(container.querySelector('svg')).toHaveTextContent('Sync with main once')
+    )
+    expect(mermaid.render).toHaveBeenCalledWith(expect.any(String), complete)
+    expect(container.querySelector('pre')).toBeNull()
+    expect(screen.queryByText(/Diagram error:/)).not.toBeInTheDocument()
+  })
+
+  it('still shows the existing error and source for an invalid finished reply', async () => {
+    mermaid.render.mockRejectedValue(new Error('Invalid mermaid syntax'))
+    const { container } = render(reply(partial, false))
+    await screen.findByText(/Diagram error: Invalid mermaid syntax/)
+    expect(container.querySelector('pre')?.textContent).toBe(partial.trimEnd())
+    expect(container.querySelector('svg')).toBeNull()
+  })
+})

src/renderer/src/components/native-chat/NativeChatMarkdown.tsx

@@ -27,6 +27,7 @@ export function NativeChatMarkdown({
         extension && streaming ? withoutPendingNativeChatVisualDirectiveTail(content) : content
       }
       extension={extension}
+      renderMermaid={!streaming}
       className={cn('native-chat-markdown', className)}
     />
   )

src/renderer/src/components/sidebar/CommentMarkdown.tsx

@@ -220,6 +220,7 @@ type CommentMarkdownProps = React.ComponentPropsWithoutRef<'div'> & {
   linkifyFilePaths?: boolean
   expandImages?: boolean
   renderCodeBlock?: DocumentCodeBlockRenderer
+  renderMermaid?: boolean
   extension?: CommentMarkdownExtension
 }
 
@@ -238,6 +239,7 @@ const CommentMarkdown = React.memo(
       linkifyFilePaths = false,
       expandImages = false,
       renderCodeBlock,
+      renderMermaid = true,
       extension,
       ...rest
     },
@@ -246,17 +248,17 @@ const CommentMarkdown = React.memo(
     const baseComponents = React.useMemo(() => {
       if (!onLinkClick) {
         return variant === 'document'
-          ? renderCodeBlock
-            ? createDocumentCommentMarkdownComponents(undefined, renderCodeBlock)
+          ? renderCodeBlock || !renderMermaid
+            ? createDocumentCommentMarkdownComponents(undefined, renderCodeBlock, renderMermaid)
             : documentCommentMarkdownComponents
           : expandImages
             ? createCompactCommentMarkdownComponents(undefined, true)
             : compactCommentMarkdownComponents
       }
       return variant === 'document'
-        ? createDocumentCommentMarkdownComponents(onLinkClick, renderCodeBlock)
+        ? createDocumentCommentMarkdownComponents(onLinkClick, renderCodeBlock, renderMermaid)
         : createCompactCommentMarkdownComponents(onLinkClick, expandImages)
-    }, [expandImages, renderCodeBlock, variant, onLinkClick])
+    }, [expandImages, renderCodeBlock, renderMermaid, variant, onLinkClick])
     const components = React.useMemo(
       () => (extension ? { ...baseComponents, ...extension.components } : baseComponents),
       [baseComponents, extension]

src/renderer/src/components/sidebar/comment-markdown-element-renderers.tsx

@@ -250,7 +250,8 @@ export function createCompactCommentMarkdownComponents(
 
 export function createDocumentCommentMarkdownComponents(
   onLinkClick?: CommentMarkdownLinkClickHandler,
-  renderCodeBlock?: DocumentCodeBlockRenderer
+  renderCodeBlock?: DocumentCodeBlockRenderer,
+  renderMermaid = true
 ): Components {
   return {
     p: ({ children }) => <p className="my-2 first:mt-0 last:mb-0">{children}</p>,
@@ -272,7 +273,7 @@ export function createDocumentCommentMarkdownComponents(
         </a>
       ),
     code: ({ className, children }) =>
-      isMermaidFence(className) ? (
+      renderMermaid && isMermaidFence(className) ? (
         renderMermaidFence(
           children,
           'my-3 min-w-0 max-w-full overflow-x-auto rounded-md border border-border/60 p-3 [&_.mermaid-block]:min-w-0 [&_.mermaid-block_pre]:my-0 [&_.mermaid-block_pre]:max-h-80 [&_.mermaid-block_pre]:max-w-full [&_.mermaid-block_pre]:overflow-x-auto [&_.mermaid-block_pre]:rounded-md [&_.mermaid-block_pre]:bg-accent [&_.mermaid-block_pre]:p-3 [&_.mermaid-block_pre]:font-mono [&_.mermaid-block_pre]:text-[12px]'
@@ -284,7 +285,7 @@ export function createDocumentCommentMarkdownComponents(
       ),
     // Mermaid fences render a <div>, which is invalid inside <pre>, so unwrap them.
     pre: ({ children }) =>
-      isMermaidPre(children) ? (
+      renderMermaid && isMermaidPre(children) ? (
         <>{children}</>
       ) : renderCodeBlock ? (
         renderCodeBlock({ children, language: extractCodeFenceLanguage(children) })