feat(mermaid): Rearchitect component for robustness, security, and theming (#21281)

web/app/components/base/markdown-blocks/code-block.tsx

@@ -271,9 +271,7 @@ const CodeBlock: any = memo(({ inline, className, children = '', ...props }: any
     const content = String(children).replace(/\n$/, '')
     switch (language) {
       case 'mermaid':
-        if (isSVG)
+        return <Flowchart PrimitiveCode={content} theme={theme as 'light' | 'dark'} />
-          return <Flowchart PrimitiveCode={content} />
-        break
       case 'echarts': {
         // Loading state: show loading indicator
         if (chartState === 'loading') {
@@ -428,7 +426,7 @@ const CodeBlock: any = memo(({ inline, className, children = '', ...props }: any
       <div className='flex h-8 items-center justify-between rounded-t-[10px] border-b border-divider-subtle bg-components-input-bg-normal p-1 pl-3'>
         <div className='system-xs-semibold-uppercase text-text-secondary'>{languageShowName}</div>
         <div className='flex items-center gap-1'>
-          {(['mermaid', 'svg']).includes(language!) && <SVGBtn isSVG={isSVG} setIsSVG={setIsSVG} />}
+          {language === 'svg' && <SVGBtn isSVG={isSVG} setIsSVG={setIsSVG} />}
           <ActionButton>
             <CopyIcon content={String(children).replace(/\n$/, '')} />
           </ActionButton>

web/app/components/base/mermaid/index.tsx

@@ -1,5 +1,5 @@
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
-import mermaid from 'mermaid'
+import mermaid, { type MermaidConfig } from 'mermaid'
 import { useTranslation } from 'react-i18next'
 import { ExclamationTriangleIcon } from '@heroicons/react/24/outline'
 import { MoonIcon, SunIcon } from '@heroicons/react/24/solid'
@@ -68,14 +68,13 @@ const THEMES = {
 const initMermaid = () => {
   if (typeof window !== 'undefined' && !isMermaidInitialized) {
     try {
-      mermaid.initialize({
+      const config: MermaidConfig = {
         startOnLoad: false,
         fontFamily: 'sans-serif',
         securityLevel: 'loose',
         flowchart: {
           htmlLabels: true,
           useMaxWidth: true,
-          diagramPadding: 10,
           curve: 'basis',
           nodeSpacing: 50,
           rankSpacing: 70,
@@ -94,10 +93,10 @@ const initMermaid = () => {
         mindmap: {
           useMaxWidth: true,
           padding: 10,
-          diagramPadding: 20,
         },
         maxTextSize: 50000,
-      })
+      }
+      mermaid.initialize(config)
       isMermaidInitialized = true
     }
     catch (error) {
@@ -113,7 +112,7 @@ const Flowchart = React.forwardRef((props: {
   theme?: 'light' | 'dark'
 }, ref) => {
   const { t } = useTranslation()
-  const [svgCode, setSvgCode] = useState<string | null>(null)
+  const [svgString, setSvgString] = useState<string | null>(null)
   const [look, setLook] = useState<'classic' | 'handDrawn'>('classic')
   const [isInitialized, setIsInitialized] = useState(false)
   const [currentTheme, setCurrentTheme] = useState<'light' | 'dark'>(props.theme || 'light')
@@ -125,6 +124,7 @@ const Flowchart = React.forwardRef((props: {
   const [imagePreviewUrl, setImagePreviewUrl] = useState('')
   const [isCodeComplete, setIsCodeComplete] = useState(false)
   const codeCompletionCheckRef = useRef<NodeJS.Timeout>()
+  const prevCodeRef = useRef<string>()
 
   // Create cache key from code, style and theme
   const cacheKey = useMemo(() => {
@@ -169,50 +169,18 @@ const Flowchart = React.forwardRef((props: {
    */
   const handleRenderError = (error: any) => {
     console.error('Mermaid rendering error:', error)
-    const errorMsg = (error as Error).message
 
-    if (errorMsg.includes('getAttribute')) {
+    // On any render error, assume the mermaid state is corrupted and force a re-initialization.
-      diagramCache.clear()
+    try {
-      mermaid.initialize({
+      diagramCache.clear() // Clear cache to prevent using potentially corrupted SVGs
-        startOnLoad: false,
+      isMermaidInitialized = false // <-- THE FIX: Force re-initialization
-        securityLevel: 'loose',
+      initMermaid() // Re-initialize with the default safe configuration
-      })
     }
-    else {
+    catch (reinitError) {
-      setErrMsg(`Rendering chart failed, please refresh and try again ${look === 'handDrawn' ? 'Or try using classic mode' : ''}`)
+      console.error('Failed to re-initialize Mermaid after error:', reinitError)
-    }
-
-    if (look === 'handDrawn') {
-      try {
-        // Clear possible cache issues
-        diagramCache.delete(`${props.PrimitiveCode}-handDrawn-${currentTheme}`)
-
-        // Reset mermaid configuration
-        mermaid.initialize({
-          startOnLoad: false,
-          securityLevel: 'loose',
-          theme: 'default',
-          maxTextSize: 50000,
-        })
-
-        // Try rendering with standard mode
-        setLook('classic')
-        setErrMsg('Hand-drawn mode is not supported for this diagram. Switched to classic mode.')
-
-        // Delay error clearing
-        setTimeout(() => {
-          if (containerRef.current) {
-            // Try rendering again with standard mode, but can't call renderFlowchart directly due to circular dependency
-            // Instead set state to trigger re-render
-            setIsCodeComplete(true) // This will trigger useEffect re-render
-          }
-        }, 500)
-      }
-      catch (e) {
-        console.error('Reset after handDrawn error failed:', e)
-      }
     }
 
+    setErrMsg(`Rendering failed: ${(error as Error).message || 'Unknown error. Please check the console.'}`)
     setIsLoading(false)
   }
 
@@ -223,51 +191,23 @@ const Flowchart = React.forwardRef((props: {
       setIsInitialized(true)
   }, [])
 
-  // Update theme when prop changes
+  // Update theme when prop changes, but allow internal override.
+  const prevThemeRef = useRef<string>()
   useEffect(() => {
-    if (props.theme)
+    // Only react if the theme prop from the outside has actually changed.
+    if (props.theme && props.theme !== prevThemeRef.current) {
+      // When the global theme prop changes, it should act as the source of truth,
+      // overriding any local theme selection.
+      diagramCache.clear()
+      setSvgString(null)
       setCurrentTheme(props.theme)
+      // Reset look to classic for a consistent state after a global change.
+      setLook('classic')
+    }
+    // Update the ref to the current prop value for the next render.
+    prevThemeRef.current = props.theme
   }, [props.theme])
 
-  // Validate mermaid code and check for completeness
-  useEffect(() => {
-    if (codeCompletionCheckRef.current)
-      clearTimeout(codeCompletionCheckRef.current)
-
-    // Reset code complete status when code changes
-    setIsCodeComplete(false)
-
-    // If no code or code is extremely short, don't proceed
-    if (!props.PrimitiveCode || props.PrimitiveCode.length < 10)
-      return
-
-    // Check if code already in cache - if so we know it's valid
-    if (diagramCache.has(cacheKey)) {
-      setIsCodeComplete(true)
-      return
-    }
-
-    // Initial check using the extracted isMermaidCodeComplete function
-    const isComplete = isMermaidCodeComplete(props.PrimitiveCode)
-    if (isComplete) {
-      setIsCodeComplete(true)
-      return
-    }
-
-    // Set a delay to check again in case code is still being generated
-    codeCompletionCheckRef.current = setTimeout(() => {
-      setIsCodeComplete(isMermaidCodeComplete(props.PrimitiveCode))
-    }, 300)
-
-    return () => {
-      if (codeCompletionCheckRef.current)
-        clearTimeout(codeCompletionCheckRef.current)
-    }
-  }, [props.PrimitiveCode, cacheKey])
-
-  /**
-   * Renders flowchart based on provided code
-   */
   const renderFlowchart = useCallback(async (primitiveCode: string) => {
     if (!isInitialized || !containerRef.current) {
       setIsLoading(false)
@@ -275,15 +215,11 @@ const Flowchart = React.forwardRef((props: {
       return
     }
 
-    // Don't render if code is not complete yet
-    if (!isCodeComplete) {
-      setIsLoading(true)
-      return
-    }
-
     // Return cached result if available
+    const cacheKey = `${primitiveCode}-${look}-${currentTheme}`
     if (diagramCache.has(cacheKey)) {
-      setSvgCode(diagramCache.get(cacheKey) || null)
+      setErrMsg('')
+      setSvgString(diagramCache.get(cacheKey) || null)
       setIsLoading(false)
       return
     }
@@ -294,17 +230,45 @@ const Flowchart = React.forwardRef((props: {
     try {
       let finalCode: string
 
-      // Check if it's a gantt chart or mindmap
+      const trimmedCode = primitiveCode.trim()
-      const isGanttChart = primitiveCode.trim().startsWith('gantt')
+      const isGantt = trimmedCode.startsWith('gantt')
-      const isMindMap = primitiveCode.trim().startsWith('mindmap')
+      const isMindMap = trimmedCode.startsWith('mindmap')
+      const isSequence = trimmedCode.startsWith('sequenceDiagram')
 
-      if (isGanttChart || isMindMap) {
+      if (isGantt || isMindMap || isSequence) {
-        // For gantt charts and mindmaps, ensure each task is on its own line
+        if (isGantt) {
-        // and preserve exact whitespace/format
+          finalCode = trimmedCode
-        finalCode = primitiveCode.trim()
+            .split('\n')
+            .map((line) => {
+              // Gantt charts have specific syntax needs.
+              const taskMatch = line.match(/^\s*([^:]+?)\s*:\s*(.*)/)
+              if (!taskMatch)
+                return line // Not a task line, return as is.
+
+              const taskName = taskMatch[1].trim()
+              let paramsStr = taskMatch[2].trim()
+
+              // Rule 1: Correct multiple "after" dependencies ONLY if they exist.
+              // This is a common mistake, e.g., "..., after task1, after task2, ..."
+              const afterCount = (paramsStr.match(/after /g) || []).length
+              if (afterCount > 1)
+                paramsStr = paramsStr.replace(/,\s*after\s+/g, ' ')
+
+              // Rule 2: Normalize spacing between parameters for consistency.
+              const finalParams = paramsStr.replace(/\s*,\s*/g, ', ').trim()
+              return `${taskName} :${finalParams}`
+            })
+            .join('\n')
+        }
+        else {
+          // For mindmap and sequence charts, which are sensitive to syntax,
+          // pass the code through directly.
+          finalCode = trimmedCode
+        }
       }
       else {
         // Step 1: Clean and prepare Mermaid code using the extracted prepareMermaidCode function
+        // This function handles flowcharts appropriately.
         finalCode = prepareMermaidCode(primitiveCode, look)
       }
 
@@ -319,13 +283,12 @@ const Flowchart = React.forwardRef((props: {
         THEMES,
       )
 
-      // Step 4: Clean SVG code and convert to base64 using the extracted functions
+      // Step 4: Clean up SVG code
       const cleanedSvg = cleanUpSvgCode(processedSvg)
-      const base64Svg = await svgToBase64(cleanedSvg)
 
-      if (base64Svg && typeof base64Svg === 'string') {
+      if (cleanedSvg && typeof cleanedSvg === 'string') {
-        diagramCache.set(cacheKey, base64Svg)
+        diagramCache.set(cacheKey, cleanedSvg)
-        setSvgCode(base64Svg)
+        setSvgString(cleanedSvg)
       }
 
       setIsLoading(false)
@@ -334,12 +297,9 @@ const Flowchart = React.forwardRef((props: {
       // Error handling
       handleRenderError(error)
     }
-  }, [chartId, isInitialized, cacheKey, isCodeComplete, look, currentTheme, t])
+  }, [chartId, isInitialized, look, currentTheme, t])
 
-  /**
+  const configureMermaid = useCallback((primitiveCode: string) => {
-   * Configure mermaid based on selected style and theme
-   */
-  const configureMermaid = useCallback(() => {
     if (typeof window !== 'undefined' && isInitialized) {
       const themeVars = THEMES[currentTheme]
       const config: any = {
@@ -361,23 +321,37 @@ const Flowchart = React.forwardRef((props: {
         mindmap: {
           useMaxWidth: true,
           padding: 10,
-          diagramPadding: 20,
         },
       }
 
+      const isFlowchart = primitiveCode.trim().startsWith('graph') || primitiveCode.trim().startsWith('flowchart')
+
       if (look === 'classic') {
         config.theme = currentTheme === 'dark' ? 'dark' : 'neutral'
-        config.flowchart = {
+
-          htmlLabels: true,
+        if (isFlowchart) {
-          useMaxWidth: true,
+          config.flowchart = {
-          diagramPadding: 12,
+            htmlLabels: true,
-          nodeSpacing: 60,
+            useMaxWidth: true,
-          rankSpacing: 80,
+            nodeSpacing: 60,
-          curve: 'linear',
+            rankSpacing: 80,
-          ranker: 'tight-tree',
+            curve: 'linear',
+            ranker: 'tight-tree',
+          }
+        }
+
+        if (currentTheme === 'dark') {
+          config.themeVariables = {
+            background: themeVars.background,
+            primaryColor: themeVars.primaryColor,
+            primaryBorderColor: themeVars.primaryBorderColor,
+            primaryTextColor: themeVars.primaryTextColor,
+            secondaryColor: themeVars.secondaryColor,
+            tertiaryColor: themeVars.tertiaryColor,
+          }
         }
       }
-      else {
+      else { // look === 'handDrawn'
         config.theme = 'default'
         config.themeCSS = `
           .node rect { fill-opacity: 0.85; }
@@ -389,27 +363,17 @@ const Flowchart = React.forwardRef((props: {
         config.themeVariables = {
           fontSize: '14px',
           fontFamily: 'sans-serif',
+          primaryBorderColor: currentTheme === 'dark' ? THEMES.dark.connectionColor : THEMES.light.connectionColor,
         }
-        config.flowchart = {
-          htmlLabels: true,
-          useMaxWidth: true,
-          diagramPadding: 10,
-          nodeSpacing: 40,
-          rankSpacing: 60,
-          curve: 'basis',
-        }
-        config.themeVariables.primaryBorderColor = currentTheme === 'dark' ? THEMES.dark.connectionColor : THEMES.light.connectionColor
-      }
 
-      if (currentTheme === 'dark' && !config.themeVariables) {
+        if (isFlowchart) {
-        config.themeVariables = {
+          config.flowchart = {
-          background: themeVars.background,
+            htmlLabels: true,
-          primaryColor: themeVars.primaryColor,
+            useMaxWidth: true,
-          primaryBorderColor: themeVars.primaryBorderColor,
+            nodeSpacing: 40,
-          primaryTextColor: themeVars.primaryTextColor,
+            rankSpacing: 60,
-          secondaryColor: themeVars.secondaryColor,
+            curve: 'basis',
-          tertiaryColor: themeVars.tertiaryColor,
+          }
-          fontFamily: 'sans-serif',
         }
       }
 
@@ -425,44 +389,50 @@ const Flowchart = React.forwardRef((props: {
     return false
   }, [currentTheme, isInitialized, look])
 
-  // Effect for theme and style configuration
+  // This is the main rendering effect.
+  // It triggers whenever the code, theme, or style changes.
   useEffect(() => {
-    if (diagramCache.has(cacheKey)) {
+    if (!isInitialized)
-      setSvgCode(diagramCache.get(cacheKey) || null)
+      return
-      setIsLoading(false)
+
-      return
+    // Don't render if code is too short
-    }
+    if (!props.PrimitiveCode || props.PrimitiveCode.length < 10) {
-
-    if (configureMermaid() && containerRef.current && isCodeComplete)
-      renderFlowchart(props.PrimitiveCode)
-  }, [look, props.PrimitiveCode, renderFlowchart, isInitialized, cacheKey, currentTheme, isCodeComplete, configureMermaid])
-
-  // Effect for rendering with debounce
-  useEffect(() => {
-    if (diagramCache.has(cacheKey)) {
-      setSvgCode(diagramCache.get(cacheKey) || null)
       setIsLoading(false)
+      setSvgString(null)
       return
     }
 
+    // Use a timeout to handle streaming code and debounce rendering
     if (renderTimeoutRef.current)
       clearTimeout(renderTimeoutRef.current)
 
-    if (isCodeComplete) {
+    setIsLoading(true)
-      renderTimeoutRef.current = setTimeout(() => {
+
-        if (isInitialized)
+    renderTimeoutRef.current = setTimeout(() => {
-          renderFlowchart(props.PrimitiveCode)
+      // Final validation before rendering
-      }, 300)
+      if (!isMermaidCodeComplete(props.PrimitiveCode)) {
-    }
+        setIsLoading(false)
-    else {
+        setErrMsg('Diagram code is not complete or invalid.')
-      setIsLoading(true)
+        return
-    }
+      }
+
+      const cacheKey = `${props.PrimitiveCode}-${look}-${currentTheme}`
+      if (diagramCache.has(cacheKey)) {
+        setErrMsg('')
+        setSvgString(diagramCache.get(cacheKey) || null)
+        setIsLoading(false)
+        return
+      }
+
+      if (configureMermaid(props.PrimitiveCode))
+        renderFlowchart(props.PrimitiveCode)
+    }, 300) // 300ms debounce
 
     return () => {
       if (renderTimeoutRef.current)
         clearTimeout(renderTimeoutRef.current)
     }
-  }, [props.PrimitiveCode, renderFlowchart, isInitialized, cacheKey, isCodeComplete])
+  }, [props.PrimitiveCode, look, currentTheme, isInitialized, configureMermaid, renderFlowchart])
 
   // Cleanup on unmount
   useEffect(() => {
@@ -471,14 +441,22 @@ const Flowchart = React.forwardRef((props: {
         containerRef.current.innerHTML = ''
       if (renderTimeoutRef.current)
         clearTimeout(renderTimeoutRef.current)
-      if (codeCompletionCheckRef.current)
-        clearTimeout(codeCompletionCheckRef.current)
     }
   }, [])
 
+  const handlePreviewClick = async () => {
+    if (svgString) {
+      const base64 = await svgToBase64(svgString)
+      setImagePreviewUrl(base64)
+    }
+  }
+
   const toggleTheme = () => {
-    setCurrentTheme(prevTheme => prevTheme === 'light' ? Theme.dark : Theme.light)
+    const newTheme = currentTheme === 'light' ? 'dark' : 'light'
+    // Ensure a full, clean re-render cycle, consistent with global theme change.
     diagramCache.clear()
+    setSvgString(null)
+    setCurrentTheme(newTheme)
   }
 
   // Style classes for theme-dependent elements
@@ -527,14 +505,26 @@ const Flowchart = React.forwardRef((props: {
             <div
               key='classic'
               className={getLookButtonClass('classic')}
-              onClick={() => setLook('classic')}
+              onClick={() => {
+                if (look !== 'classic') {
+                  diagramCache.clear()
+                  setSvgString(null)
+                  setLook('classic')
+                }
+              }}
             >
               <div className="msh-segmented-item-label">{t('app.mermaid.classic')}</div>
             </div>
             <div
               key='handDrawn'
               className={getLookButtonClass('handDrawn')}
-              onClick={() => setLook('handDrawn')}
+              onClick={() => {
+                if (look !== 'handDrawn') {
+                  diagramCache.clear()
+                  setSvgString(null)
+                  setLook('handDrawn')
+                }
+              }}
             >
               <div className="msh-segmented-item-label">{t('app.mermaid.handDrawn')}</div>
             </div>
@@ -544,7 +534,7 @@ const Flowchart = React.forwardRef((props: {
 
       <div ref={containerRef} style={{ position: 'absolute', visibility: 'hidden', height: 0, overflow: 'hidden' }} />
 
-      {isLoading && !svgCode && (
+      {isLoading && !svgString && (
         <div className='px-[26px] py-4'>
           <LoadingAnim type='text'/>
           {!isCodeComplete && (
@@ -555,8 +545,8 @@ const Flowchart = React.forwardRef((props: {
         </div>
       )}
 
-      {svgCode && (
+      {svgString && (
-        <div className={themeClasses.mermaidDiv} style={{ objectFit: 'cover' }} onClick={() => setImagePreviewUrl(svgCode)}>
+        <div className={themeClasses.mermaidDiv} style={{ objectFit: 'cover' }} onClick={handlePreviewClick}>
           <div className="absolute bottom-2 left-2 z-[100]">
             <button
               onClick={(e) => {
@@ -571,11 +561,9 @@ const Flowchart = React.forwardRef((props: {
             </button>
           </div>
 
-          <img
+          <div
-            src={svgCode}
-            alt="mermaid_chart"
             style={{ maxWidth: '100%' }}
-            onError={() => { setErrMsg('Chart rendering failed, please refresh and retry') }}
+            dangerouslySetInnerHTML={{ __html: svgString }}
           />
         </div>
       )}

web/app/components/base/mermaid/utils.ts

@@ -3,52 +3,31 @@ export function cleanUpSvgCode(svgCode: string): string {
 }
 
 /**
- * Preprocesses mermaid code to fix common syntax issues
+ * Prepares mermaid code for rendering by sanitizing common syntax issues.
+ * @param {string} mermaidCode - The mermaid code to prepare
+ * @param {'classic' | 'handDrawn'} style - The rendering style
+ * @returns {string} - The prepared mermaid code
  */
-export function preprocessMermaidCode(code: string): string {
+export const prepareMermaidCode = (mermaidCode: string, style: 'classic' | 'handDrawn'): string => {
-  if (!code || typeof code !== 'string')
+  if (!mermaidCode || typeof mermaidCode !== 'string')
     return ''
 
-  // First check if this is a gantt chart
+  let code = mermaidCode.trim()
-  if (code.trim().startsWith('gantt')) {
-    // For gantt charts, we need to ensure each task is on its own line
-    // Split the code into lines and process each line separately
-    const lines = code.split('\n').map(line => line.trim())
-    return lines.join('\n')
-  }
 
-  return code
+  // Security: Sanitize against javascript: protocol in click events (XSS vector)
-    // Replace English colons with Chinese colons in section nodes to avoid parsing issues
+  code = code.replace(/(\bclick\s+\w+\s+")javascript:[^"]*(")/g, '$1#$2')
-    .replace(/section\s+([^:]+):/g, (match, sectionName) => `section ${sectionName}：`)
-    // Fix common syntax issues
-    .replace(/fifopacket/g, 'rect')
-    // Ensure graph has direction
-    .replace(/^graph\s+((?:TB|BT|RL|LR)*)/, (match, direction) => {
-      return direction ? match : 'graph TD'
-    })
-    // Clean up empty lines and extra spaces
-    .trim()
-}
 
-/**
+  // Convenience: Basic BR replacement. This is a common and safe operation.
- * Prepares mermaid code based on selected style
+  code = code.replace(/<br\s*\/?>/g, '\n')
- */
-export function prepareMermaidCode(code: string, style: 'classic' | 'handDrawn'): string {
-  let finalCode = preprocessMermaidCode(code)
 
-  // Special handling for gantt charts and mindmaps
+  let finalCode = code
-  if (finalCode.trim().startsWith('gantt') || finalCode.trim().startsWith('mindmap')) {
-    // For gantt charts and mindmaps, preserve the structure exactly as is
-    return finalCode
-  }
 
+  // Hand-drawn style requires some specific clean-up.
   if (style === 'handDrawn') {
     finalCode = finalCode
-      // Remove style definitions that interfere with hand-drawn style
       .replace(/style\s+[^\n]+/g, '')
       .replace(/linkStyle\s+[^\n]+/g, '')
       .replace(/^flowchart/, 'graph')
-      // Remove any styles that might interfere with hand-drawn style
       .replace(/class="[^"]*"/g, '')
       .replace(/fill="[^"]*"/g, '')
       .replace(/stroke="[^"]*"/g, '')
@@ -82,7 +61,6 @@ export function svgToBase64(svgGraph: string): Promise<string> {
     })
   }
   catch (error) {
-    console.error('Error converting SVG to base64:', error)
     return Promise.resolve('')
   }
 }
@@ -115,13 +93,11 @@ export function processSvgForTheme(
     }
     else {
       let i = 0
-      themes.dark.nodeColors.forEach(() => {
+      const nodeColorRegex = /fill="#[a-fA-F0-9]{6}"[^>]*class="node-[^"]*"/g
-        const regex = /fill="#[a-fA-F0-9]{6}"[^>]*class="node-[^"]*"/g
+      processedSvg = processedSvg.replace(nodeColorRegex, (match: string) => {
-        processedSvg = processedSvg.replace(regex, (match: string) => {
+        const colorIndex = i % themes.dark.nodeColors.length
-          const colorIndex = i % themes.dark.nodeColors.length
+        i++
-          i++
+        return match.replace(/fill="#[a-fA-F0-9]{6}"/, `fill="${themes.dark.nodeColors[colorIndex].bg}"`)
-          return match.replace(/fill="#[a-fA-F0-9]{6}"/, `fill="${themes.dark.nodeColors[colorIndex].bg}"`)
-        })
       })
 
       processedSvg = processedSvg
@@ -139,14 +115,12 @@ export function processSvgForTheme(
         .replace(/stroke-width="1"/g, 'stroke-width="1.5"')
     }
     else {
-      themes.light.nodeColors.forEach(() => {
+      let i = 0
-        const regex = /fill="#[a-fA-F0-9]{6}"[^>]*class="node-[^"]*"/g
+      const nodeColorRegex = /fill="#[a-fA-F0-9]{6}"[^>]*class="node-[^"]*"/g
-        let i = 0
+      processedSvg = processedSvg.replace(nodeColorRegex, (match: string) => {
-        processedSvg = processedSvg.replace(regex, (match: string) => {
+        const colorIndex = i % themes.light.nodeColors.length
-          const colorIndex = i % themes.light.nodeColors.length
+        i++
-          i++
+        return match.replace(/fill="#[a-fA-F0-9]{6}"/, `fill="${themes.light.nodeColors[colorIndex].bg}"`)
-          return match.replace(/fill="#[a-fA-F0-9]{6}"/, `fill="${themes.light.nodeColors[colorIndex].bg}"`)
-        })
       })
 
       processedSvg = processedSvg
@@ -187,24 +161,10 @@ export function isMermaidCodeComplete(code: string): boolean {
     // Check for basic syntax structure
     const hasValidStart = /^(graph|flowchart|sequenceDiagram|classDiagram|classDef|class|stateDiagram|gantt|pie|er|journey|requirementDiagram|mindmap)/.test(trimmedCode)
 
-    // Check for balanced brackets and parentheses
+    // The balanced bracket check was too strict and produced false negatives for valid
-    const isBalanced = (() => {
+    // mermaid syntax like the asymmetric shape `A>B]`. Relying on Mermaid's own
-      const stack = []
+    // parser is more robust.
-      const pairs = { '{': '}', '[': ']', '(': ')' }
+    const isBalanced = true
-
-      for (const char of trimmedCode) {
-        if (char in pairs) {
-          stack.push(char)
-        }
-        else if (Object.values(pairs).includes(char)) {
-          const last = stack.pop()
-          if (pairs[last as keyof typeof pairs] !== char)
-            return false
-        }
-      }
-
-      return stack.length === 0
-    })()
 
     // Check for common syntax errors
     const hasNoSyntaxErrors = !trimmedCode.includes('undefined')
@@ -215,7 +175,7 @@ export function isMermaidCodeComplete(code: string): boolean {
     return hasValidStart && isBalanced && hasNoSyntaxErrors
   }
   catch (error) {
-    console.debug('Mermaid code validation error:', error)
+    console.error('Mermaid code validation error:', error)
     return false
   }
 }