diff --git a/lightrag_webui/src/components/retrieval/ChatMessage.tsx b/lightrag_webui/src/components/retrieval/ChatMessage.tsx
index 81888889..1647c2a2 100644
--- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx
+++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx
@@ -1,4 +1,4 @@
-import { ReactNode, useCallback, useEffect, useRef } from 'react'
+import { ReactNode, useCallback, useEffect, useMemo, useRef, memo, useState } from 'react' // Import useMemo
 import { Message } from '@/api/lightrag'
 import useTheme from '@/hooks/useTheme'
 import Button from '@/components/ui/Button'
@@ -19,10 +19,17 @@ import { LoaderIcon, CopyIcon } from 'lucide-react'
 import { useTranslation } from 'react-i18next'
 
 export type MessageWithError = Message & {
+  id: string // Unique identifier for stable React keys
   isError?: boolean
+  /**
+   * Indicates if the mermaid diagram in this message has been rendered.
+   * Used to persist the rendering state across updates and prevent flickering.
+   */
+  mermaidRendered?: boolean
 }
 
-export const ChatMessage = ({ message, isComplete = true }: { message: MessageWithError, isComplete?: boolean }) => {
+// Restore original component definition and export
+export const ChatMessage = ({ message }: { message: MessageWithError }) => { // Remove isComplete prop
   const { t } = useTranslation()
   const handleCopyMarkdown = useCallback(async () => {
     if (message.content) {
@@ -50,17 +57,23 @@ export const ChatMessage = ({ message, isComplete = true }: { message: MessageWi
           remarkPlugins={[remarkGfm, remarkMath]}
           rehypePlugins={[rehypeReact]}
           skipHtml={false}
-          components={{
-            code: (props) => <CodeHighlight {...props} isComplete={isComplete} />,
-            p: ({ children }) => <p className="my-2">{children}</p>,
-            h1: ({ children }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
-            h2: ({ children }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
-            h3: ({ children }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
-            h4: ({ children }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
-            ul: ({ children }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
-            ol: ({ children }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
-            li: ({ children }) => <li className="my-1">{children}</li>
-          }}
+          // Memoize the components object to prevent unnecessary re-renders of ReactMarkdown children
+          components={useMemo(() => ({
+            code: (props: any) => ( // Add type annotation if needed, e.g., props: CodeProps from 'react-markdown/lib/ast-to-react'
+              <CodeHighlight
+                {...props}
+                renderAsDiagram={message.mermaidRendered ?? false}
+              />
+            ),
+            p: ({ children }: { children?: ReactNode }) => <p className="my-2">{children}</p>,
+            h1: ({ children }: { children?: ReactNode }) => <h1 className="text-xl font-bold mt-4 mb-2">{children}</h1>,
+            h2: ({ children }: { children?: ReactNode }) => <h2 className="text-lg font-bold mt-4 mb-2">{children}</h2>,
+            h3: ({ children }: { children?: ReactNode }) => <h3 className="text-base font-bold mt-3 mb-2">{children}</h3>,
+            h4: ({ children }: { children?: ReactNode }) => <h4 className="text-base font-semibold mt-3 mb-2">{children}</h4>,
+            ul: ({ children }: { children?: ReactNode }) => <ul className="list-disc pl-5 my-2">{children}</ul>,
+            ol: ({ children }: { children?: ReactNode }) => <ol className="list-decimal pl-5 my-2">{children}</ol>,
+            li: ({ children }: { children?: ReactNode }) => <li className="my-1">{children}</li>
+          }), [message.mermaidRendered])} // Dependency ensures update if mermaid state changes
         >
           {message.content}
         </ReactMarkdown>
@@ -81,12 +94,14 @@ export const ChatMessage = ({ message, isComplete = true }: { message: MessageWi
   )
 }
 
+// Remove the incorrect memo export line
+
 interface CodeHighlightProps {
   inline?: boolean
   className?: string
   children?: ReactNode
   node?: Element // Keep node for inline check
-  isComplete?: boolean // Flag to indicate if the message is complete
+  renderAsDiagram?: boolean // Flag to indicate if rendering as diagram should be attempted
 }
 
 // Helper function remains the same
@@ -101,8 +116,10 @@ const isInlineCode = (node?: Element): boolean => {
 };
 
 
-const CodeHighlight = ({ className, children, node, isComplete = true, ...props }: CodeHighlightProps) => {
+// Memoize the CodeHighlight component
+const CodeHighlight = memo(({ className, children, node, renderAsDiagram = false, ...props }: CodeHighlightProps) => {
   const { theme } = useTheme();
+  const [hasRendered, setHasRendered] = useState(false); // State to track successful render
   const match = className?.match(/language-(\w+)/);
   const language = match ? match[1] : undefined;
   const inline = isInlineCode(node); // Use the helper function
@@ -111,35 +128,37 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
 
   // Handle Mermaid rendering with debounce
   useEffect(() => {
-    // Clear any existing timer when dependencies change
-    if (debounceTimerRef.current) {
-      clearTimeout(debounceTimerRef.current);
-    }
+    // Effect should run when renderAsDiagram becomes true or hasRendered changes.
+    // The actual rendering logic inside checks language and hasRendered state.
+    if (renderAsDiagram && !hasRendered && language === 'mermaid' && mermaidRef.current) {
+      const container = mermaidRef.current; // Capture ref value
 
-    if (language === 'mermaid' && mermaidRef.current) {
-      const container = mermaidRef.current; // Capture ref value for use inside timeout/callbacks
+      // Clear previous timer if dependencies change before timeout (e.g., renderAsDiagram flips quickly)
+      if (debounceTimerRef.current) {
+        clearTimeout(debounceTimerRef.current);
+      }
 
-      // Set a new timer to render after a short delay
       debounceTimerRef.current = setTimeout(() => {
-        // Ensure container still exists when timer fires
-        if (!container) return;
+        if (!container) return; // Container might have unmounted
+
+        // Double check hasRendered state inside timeout, in case it changed rapidly
+        if (hasRendered) return;
 
         try {
-          // Initialize mermaid config (safe to call multiple times)
+          // Initialize mermaid config
           mermaid.initialize({
             startOnLoad: false,
             theme: theme === 'dark' ? 'dark' : 'default',
             securityLevel: 'loose',
           });
 
-          // Show loading indicator while processing
+          // Show loading indicator
           container.innerHTML = '<div class="flex justify-center items-center p-4"><svg class="animate-spin h-5 w-5 text-primary" xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24"><circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle><path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path></svg></div>';
 
           // Preprocess mermaid content
-          const rawContent = String(children).replace(/\n$/, '').trim(); // Trim whitespace as well
+          const rawContent = String(children).replace(/\n$/, '').trim();
 
           // Heuristic check for potentially complete graph definition
-          // Looks for graph type declaration and some content beyond it.
           const looksPotentiallyComplete = rawContent.length > 10 && (
             rawContent.startsWith('graph') ||
             rawContent.startsWith('sequenceDiagram') ||
@@ -151,19 +170,17 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
             rawContent.startsWith('erDiagram')
           );
 
-
           if (!looksPotentiallyComplete) {
             console.log('Mermaid content might be incomplete, skipping render attempt:', rawContent);
-            // Keep loading indicator or show a message
+            // Optionally keep loading indicator or show a message
             // container.innerHTML = '<p class="text-sm text-muted-foreground">Waiting for complete diagram...</p>';
-            return; // Don't attempt to render potentially incomplete content
+            return;
           }
 
           const processedContent = rawContent
             .split('\n')
             .map(line => {
               const trimmedLine = line.trim();
-              // Keep subgraph processing
               if (trimmedLine.startsWith('subgraph')) {
                 const parts = trimmedLine.split(' ');
                 if (parts.length > 1) {
@@ -173,26 +190,25 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
               }
               return trimmedLine;
             })
-            .filter(line => !line.trim().startsWith('linkStyle')) // Keep filtering linkStyle
+            .filter(line => !line.trim().startsWith('linkStyle'))
             .join('\n');
 
           const mermaidId = `mermaid-${Date.now()}`;
           mermaid.render(mermaidId, processedContent)
             .then(({ svg, bindFunctions }) => {
-              // Check ref again inside async callback
-              // Ensure the container is still the one we intended to update
-              if (mermaidRef.current === container) {
+              // Check ref and hasRendered state again inside async callback
+              if (mermaidRef.current === container && !hasRendered) {
                 container.innerHTML = svg;
+                setHasRendered(true); // Mark as rendered successfully
                 if (bindFunctions) {
-                  try { // Add try-catch around bindFunctions as it can also throw
+                  try {
                     bindFunctions(container);
                   } catch (bindError) {
                     console.error('Mermaid bindFunctions error:', bindError);
-                    // Optionally display a message in the container
                     container.innerHTML += '<p class="text-orange-500 text-xs">Diagram interactions might be limited.</p>';
                   }
                 }
-              } else {
+              } else if (mermaidRef.current !== container) {
                 console.log('Mermaid container changed before rendering completed.');
               }
             })
@@ -201,11 +217,10 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
               console.error('Failed content (debounced):', processedContent);
               if (mermaidRef.current === container) {
                 const errorMessage = error instanceof Error ? error.message : String(error);
-                // Make error display more robust
                 const errorPre = document.createElement('pre');
                 errorPre.className = 'text-red-500 text-xs whitespace-pre-wrap break-words';
                 errorPre.textContent = `Mermaid diagram error: ${errorMessage}\n\nContent:\n${processedContent}`;
-                container.innerHTML = ''; // Clear previous content
+                container.innerHTML = '';
                 container.appendChild(errorPre);
               }
             });
@@ -218,24 +233,28 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
             const errorPre = document.createElement('pre');
             errorPre.className = 'text-red-500 text-xs whitespace-pre-wrap break-words';
             errorPre.textContent = `Mermaid diagram setup error: ${errorMessage}`;
-            container.innerHTML = ''; // Clear previous content
+            container.innerHTML = '';
             container.appendChild(errorPre);
           }
         }
-      }, 300); // 300ms debounce delay
+      }, 300); // Debounce delay
     }
 
-    // Cleanup function to clear the timer
+    // Cleanup function to clear the timer on unmount or before re-running effect
     return () => {
       if (debounceTimerRef.current) {
         clearTimeout(debounceTimerRef.current);
       }
     };
-  }, [language, children, theme]); // Dependencies
+  // Dependencies: renderAsDiagram ensures effect runs when diagram should be shown.
+  // children, language, theme trigger re-render if code/context changes.
+  // Dependencies are minimal: only run when the intent to render changes or the rendered state changes.
+  // Access children, theme, language inside the effect when needed.
+  }, [renderAsDiagram, hasRendered, language]); // Keep language to ensure it IS mermaid
 
   // Render based on language type
-  // If it's a mermaid language block and the message is not complete, display as plain text
-  if (language === 'mermaid' && !isComplete) {
+  // If it's a mermaid language block and rendering as diagram is not requested (e.g., incomplete stream), display as plain text
+  if (language === 'mermaid' && !renderAsDiagram) {
     return (
       <SyntaxHighlighter
         style={theme === 'dark' ? oneDark : oneLight}
@@ -273,4 +292,7 @@ const CodeHighlight = ({ className, children, node, isComplete = true, ...props
       {children}
     </code>
   );
-};
+});
+
+// Assign display name for React DevTools
+CodeHighlight.displayName = 'CodeHighlight';
diff --git a/lightrag_webui/src/features/RetrievalTesting.tsx b/lightrag_webui/src/features/RetrievalTesting.tsx
index a2e5f569..6e7bdf7e 100644
--- a/lightrag_webui/src/features/RetrievalTesting.tsx
+++ b/lightrag_webui/src/features/RetrievalTesting.tsx
@@ -2,7 +2,7 @@ import Input from '@/components/ui/Input'
 import Button from '@/components/ui/Button'
 import { useCallback, useEffect, useRef, useState } from 'react'
 import { throttle } from '@/lib/utils'
-import { queryText, queryTextStream, Message } from '@/api/lightrag'
+import { queryText, queryTextStream } from '@/api/lightrag'
 import { errorMessage } from '@/lib/utils'
 import { useSettingsStore } from '@/stores/settings'
 import { useDebounce } from '@/hooks/useDebounce'
@@ -14,9 +14,18 @@ import type { QueryMode } from '@/api/lightrag'
 
 export default function RetrievalTesting() {
   const { t } = useTranslation()
-  const [messages, setMessages] = useState<MessageWithError[]>(
-    () => useSettingsStore.getState().retrievalHistory || []
-  )
+  const [messages, setMessages] = useState<MessageWithError[]>(() => {
+    const history = useSettingsStore.getState().retrievalHistory || []
+    // Ensure each message from history has a unique ID and mermaidRendered status
+    return history.map((msg, index) => {
+      const msgWithError = msg as MessageWithError // Cast to access potential properties
+      return {
+        ...msg,
+        id: msgWithError.id || `hist-${Date.now()}-${index}`, // Add ID if missing
+        mermaidRendered: msgWithError.mermaidRendered ?? true // Assume historical mermaid is rendered
+      }
+    })
+  })
   const [inputValue, setInputValue] = useState('')
   const [isLoading, setIsLoading] = useState(false)
   const [inputError, setInputError] = useState('') // Error message for input
@@ -81,14 +90,17 @@ export default function RetrievalTesting() {
 
       // Create messages
       // Save the original input (with prefix if any) in userMessage.content for display
-      const userMessage: Message = {
+      const userMessage: MessageWithError = {
+        id: crypto.randomUUID(), // Add unique ID
         content: inputValue,
         role: 'user'
       }
 
-      const assistantMessage: Message = {
+      const assistantMessage: MessageWithError = {
+        id: crypto.randomUUID(), // Add unique ID
         content: '',
-        role: 'assistant'
+        role: 'assistant',
+        mermaidRendered: false
       }
 
       const prevMessages = [...messages]
@@ -113,12 +125,28 @@ export default function RetrievalTesting() {
       // Create a function to update the assistant's message
       const updateAssistantMessage = (chunk: string, isError?: boolean) => {
         assistantMessage.content += chunk
+
+        // Detect if the assistant message contains a complete mermaid code block
+        // Simple heuristic: look for ```mermaid ... ```
+        const mermaidBlockRegex = /```mermaid\s+([\s\S]+?)```/g
+        let mermaidRendered = false
+        let match
+        while ((match = mermaidBlockRegex.exec(assistantMessage.content)) !== null) {
+          // If the block is not too short, consider it complete
+          if (match[1] && match[1].trim().length > 10) {
+            mermaidRendered = true
+            break
+          }
+        }
+        assistantMessage.mermaidRendered = mermaidRendered
+
         setMessages((prev) => {
           const newMessages = [...prev]
           const lastMessage = newMessages[newMessages.length - 1]
           if (lastMessage.role === 'assistant') {
             lastMessage.content = assistantMessage.content
             lastMessage.isError = isError
+            lastMessage.mermaidRendered = assistantMessage.mermaidRendered
           }
           return newMessages
         })
@@ -279,20 +307,14 @@ export default function RetrievalTesting() {
                   {t('retrievePanel.retrieval.startPrompt')}
                 </div>
               ) : (
-                messages.map((message, idx) => {
-                  // Determine if this message is complete:
-                  // 1. If it's not the last message, it's complete
-                  // 2. If it's the last message but we're not receiving a streaming response, it's complete
-                  // 3. If it's the last message and we're receiving a streaming response, it's not complete
-                  const isLastMessage = idx === messages.length - 1;
-                  const isMessageComplete = !isLastMessage || !isReceivingResponseRef.current;
-                  
+                messages.map((message) => { // Remove unused idx
+                  // isComplete logic is now handled internally based on message.mermaidRendered
                   return (
                     <div
-                      key={idx}
+                      key={message.id} // Use stable ID for key
                       className={`flex ${message.role === 'user' ? 'justify-end' : 'justify-start'}`}
                     >
-                      {<ChatMessage message={message} isComplete={isMessageComplete} />}
+                      {<ChatMessage message={message} />}
                     </div>
                   );
                 })