Fix thinking state management and UI improvements for chat messages

- Reset expansion state on new thinking - Prevent content leakage during thinking - Update thinkingTime type to allow null - Add memory leak prevention cleanup - Add missing i18n translations
2025-09-08 18:47:41 +08:00 · 2025-09-08 18:47:41 +08:00 · eeff0d5c67
commit eeff0d5c67
parent 2ba10dbb5b
8 changed files with 112 additions and 43 deletions
--- a/lightrag_webui/src/api/lightrag.ts
+++ b/lightrag_webui/src/api/lightrag.ts
@ -101,7 +101,7 @@ export type Message = {
  content: string
  thinkingContent?: string
  displayContent?: string
-  thinkingTime?: number
+  thinkingTime?: number | null
 }

 export type QueryRequest = {
--- a/lightrag_webui/src/components/retrieval/ChatMessage.tsx
+++ b/lightrag_webui/src/components/retrieval/ChatMessage.tsx
@ -39,11 +39,21 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
  // Directly use props passed from the parent.
  const { thinkingContent, displayContent, thinkingTime, isThinking } = message

+  // Reset expansion state when new thinking starts
+  useEffect(() => {
+    if (isThinking) {
+      // When thinking starts, always reset to collapsed state
+      setIsThinkingExpanded(false)
+    }
+  }, [isThinking, message.id])
+
  // The content to display is now non-ambiguous.
  const finalThinkingContent = thinkingContent
  // For user messages, displayContent will be undefined, so we fall back to content.
-  // For assistant messages, we prefer displayContent.
-  const finalDisplayContent = message.role === 'user' ? message.content : displayContent
+  // For assistant messages, we prefer displayContent and don't fallback to avoid content leakage during thinking
+  const finalDisplayContent = message.role === 'user'
+    ? message.content
+    : displayContent || ''

  // Load KaTeX dynamically
  useEffect(() => {
@ -106,7 +116,12 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
        <div className="mb-2">
          <div
            className="flex items-center text-gray-500 dark:text-gray-400 hover:text-gray-700 dark:hover:text-gray-200 transition-colors duration-200 text-sm cursor-pointer select-none"
-            onClick={() => setIsThinkingExpanded(!isThinkingExpanded)}
+            onClick={() => {
+              // Allow expansion when there's thinking content, even during thinking process
+              if (finalThinkingContent && finalThinkingContent.trim() !== '') {
+                setIsThinkingExpanded(!isThinkingExpanded)
+              }
+            }}
          >
            {isThinking ? (
              <>
@ -116,10 +131,17 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
            ) : (
              typeof thinkingTime === 'number' && <span>{t('retrievePanel.chatMessage.thinkingTime', { time: thinkingTime })}</span>
            )}
-            {finalThinkingContent && <ChevronDownIcon className={`ml-2 size-4 shrink-0 transition-transform ${isThinkingExpanded ? 'rotate-180' : ''}`} />}
+            {/* Show chevron when there's thinking content, even during thinking process */}
+            {finalThinkingContent && finalThinkingContent.trim() !== '' && <ChevronDownIcon className={`ml-2 size-4 shrink-0 transition-transform ${isThinkingExpanded ? 'rotate-180' : ''}`} />}
          </div>
-          {isThinkingExpanded && finalThinkingContent && (
+          {/* Show thinking content when expanded and content exists, even during thinking process */}
+          {isThinkingExpanded && finalThinkingContent && finalThinkingContent.trim() !== '' && (
            <div className="mt-2 pl-4 border-l-2 border-primary/20 text-sm prose dark:prose-invert max-w-none break-words prose-p:my-1 prose-headings:my-2">
+              {isThinking && (
+                <div className="mb-2 text-xs text-gray-400 dark:text-gray-500 italic">
+                  {t('retrievePanel.chatMessage.thinkingInProgress', 'Thinking in progress...')}
+                </div>
+              )}
              <ReactMarkdown
                remarkPlugins={[remarkGfm, remarkMath]}
                rehypePlugins={[
@ -170,7 +192,12 @@ export const ChatMessage = ({ message }: { message: MessageWithError }) => { //
          )}
        </div>
      )}
-      {message.content === '' && !isThinking && !thinkingTime && <LoaderIcon className="animate-spin duration-2000" />} {/* Check for empty string specifically */}
+      {(() => {
+        // More comprehensive loading state check
+        const hasVisibleContent = finalDisplayContent && finalDisplayContent.trim() !== '';
+        const isLoadingState = !hasVisibleContent && !isThinking && !thinkingTime;
+        return isLoadingState && <LoaderIcon className="animate-spin duration-2000" />;
+      })()}
    </div>
  )
 }
--- a/lightrag_webui/src/features/RetrievalTesting.tsx
+++ b/lightrag_webui/src/features/RetrievalTesting.tsx
@ -68,6 +68,16 @@ export default function RetrievalTesting() {
  const messagesEndRef = useRef<HTMLDivElement>(null)
  const messagesContainerRef = useRef<HTMLDivElement>(null)

+  // Add cleanup effect for memory leak prevention
+  useEffect(() => {
+    // Component cleanup - reset timer state to prevent memory leaks
+    return () => {
+      if (thinkingStartTime.current) {
+        thinkingStartTime.current = null;
+      }
+    };
+  }, []);
+
  // Scroll to bottom function - restored smooth scrolling with better handling
  const scrollToBottom = useCallback(() => {
    // Set flag to indicate this is a programmatic scroll
@ -116,6 +126,9 @@ export default function RetrievalTesting() {
      // Clear error message
      setInputError('')

+      // Reset thinking timer state for new query to prevent confusion
+      thinkingStartTime.current = null
+
      // Create messages
      // Save the original input (with prefix if any) in userMessage.content for display
      const userMessage: MessageWithError = {
@ -128,7 +141,11 @@ export default function RetrievalTesting() {
        id: generateUniqueId(), // Use browser-compatible ID generation
        content: '',
        role: 'assistant',
-        mermaidRendered: false
+        mermaidRendered: false,
+        thinkingTime: null,        // Explicitly initialize to null
+        thinkingContent: undefined, // Explicitly initialize to undefined
+        displayContent: undefined,  // Explicitly initialize to undefined
+        isThinking: false          // Explicitly initialize to false
      }

      const prevMessages = [...messages]
@ -160,10 +177,10 @@ export default function RetrievalTesting() {
        }

        // Real-time parsing for streaming
-        const thinkStartTag = '<think>';
-        const thinkEndTag = '</think>';
-        const thinkStartIndex = assistantMessage.content.indexOf(thinkStartTag);
-        const thinkEndIndex = assistantMessage.content.indexOf(thinkEndTag);
+        const thinkStartTag = '<think>'
+        const thinkEndTag = '</think>'
+        const thinkStartIndex = assistantMessage.content.indexOf(thinkStartTag)
+        const thinkEndIndex = assistantMessage.content.indexOf(thinkEndTag)

        if (thinkStartIndex !== -1) {
          if (thinkEndIndex !== -1) {
@ -173,20 +190,21 @@ export default function RetrievalTesting() {
              const duration = (Date.now() - thinkingStartTime.current) / 1000
              assistantMessage.thinkingTime = parseFloat(duration.toFixed(2))
            }
-            assistantMessage.thinkingContent = assistantMessage.content.substring(thinkStartIndex + thinkStartTag.length, thinkEndIndex).trim()
+            assistantMessage.thinkingContent = assistantMessage.content
+              .substring(thinkStartIndex + thinkStartTag.length, thinkEndIndex)
+              .trim()
            assistantMessage.displayContent = assistantMessage.content.substring(thinkEndIndex + thinkEndTag.length).trim()
          } else {
            // Still thinking
-            assistantMessage.isThinking = true;
-            assistantMessage.thinkingContent = assistantMessage.content.substring(thinkStartIndex + thinkStartTag.length);
-            assistantMessage.displayContent = '';
+            assistantMessage.isThinking = true
+            assistantMessage.thinkingContent = assistantMessage.content.substring(thinkStartIndex + thinkStartTag.length)
+            assistantMessage.displayContent = ''
          }
        } else {
-          assistantMessage.isThinking = false;
-          assistantMessage.displayContent = assistantMessage.content;
+          assistantMessage.isThinking = false
+          assistantMessage.displayContent = assistantMessage.content
        }

-
        // Detect if the assistant message contains a complete mermaid code block
        // Simple heuristic: look for ```mermaid ... ```
        const mermaidBlockRegex = /```mermaid\s+([\s\S]+?)```/g
@ -201,16 +219,21 @@ export default function RetrievalTesting() {
        }
        assistantMessage.mermaidRendered = mermaidRendered

+        // Single unified update to avoid race conditions
        setMessages((prev) => {
          const newMessages = [...prev]
          const lastMessage = newMessages[newMessages.length - 1]
-          if (lastMessage.role === 'assistant') {
-            lastMessage.content = assistantMessage.content;
-            lastMessage.thinkingContent = assistantMessage.thinkingContent;
-            lastMessage.displayContent = assistantMessage.displayContent;
-            lastMessage.isThinking = assistantMessage.isThinking;
-            lastMessage.isError = isError;
-            lastMessage.mermaidRendered = assistantMessage.mermaidRendered;
+          if (lastMessage && lastMessage.id === assistantMessage.id) {
+            // Update all properties at once to maintain consistency
+            Object.assign(lastMessage, {
+              content: assistantMessage.content,
+              thinkingContent: assistantMessage.thinkingContent,
+              displayContent: assistantMessage.displayContent,
+              isThinking: assistantMessage.isThinking,
+              isError: isError,
+              mermaidRendered: assistantMessage.mermaidRendered,
+              thinkingTime: assistantMessage.thinkingTime
+            })
          }
          return newMessages
        })
@ -261,21 +284,29 @@ export default function RetrievalTesting() {
        setIsLoading(false)
        isReceivingResponseRef.current = false

-        // Final calculation for thinking time, only if not already calculated
-        if (assistantMessage.thinkingContent && thinkingStartTime.current && !assistantMessage.thinkingTime) {
-          const duration = (Date.now() - thinkingStartTime.current) / 1000
-          assistantMessage.thinkingTime = parseFloat(duration.toFixed(2))
-        }
-        // Ensure isThinking is false at the very end
-        assistantMessage.isThinking = false;
-        // Always reset the timer at the end of a query
-        if (thinkingStartTime.current) {
+        // Enhanced cleanup with error handling to prevent memory leaks
+        try {
+          // Final calculation for thinking time, only if not already calculated
+          if (assistantMessage.thinkingContent && thinkingStartTime.current && !assistantMessage.thinkingTime) {
+            const duration = (Date.now() - thinkingStartTime.current) / 1000
+            assistantMessage.thinkingTime = parseFloat(duration.toFixed(2))
+          }
+        } catch (error) {
+          console.error('Error calculating thinking time:', error)
+        } finally {
+          // Ensure cleanup happens regardless of errors
+          assistantMessage.isThinking = false;
          thinkingStartTime.current = null;
        }

-        useSettingsStore
-          .getState()
-          .setRetrievalHistory([...prevMessages, userMessage, assistantMessage])
+        // Save history with error handling
+        try {
+          useSettingsStore
+            .getState()
+            .setRetrievalHistory([...prevMessages, userMessage, assistantMessage])
+        } catch (error) {
+          console.error('Error saving retrieval history:', error)
+        }
      }
    },
    [inputValue, isLoading, messages, setMessages, t, scrollToBottom]
--- a/lightrag_webui/src/locales/ar.json
+++ b/lightrag_webui/src/locales/ar.json
@ -337,7 +337,10 @@
  "retrievePanel": {
    "chatMessage": {
      "copyTooltip": "نسخ إلى الحافظة",
-      "copyError": "فشل نسخ النص إلى الحافظة"
+      "copyError": "فشل نسخ النص إلى الحافظة",
+      "thinking": "جاري التفكير...",
+      "thinkingTime": "وقت التفكير {{time}} ثانية",
+      "thinkingInProgress": "التفكير قيد التقدم..."
    },
    "retrieval": {
      "startPrompt": "ابدأ الاسترجاع بكتابة استفسارك أدناه",
--- a/lightrag_webui/src/locales/en.json
+++ b/lightrag_webui/src/locales/en.json
@ -339,7 +339,8 @@
      "copyTooltip": "Copy to clipboard",
      "copyError": "Failed to copy text to clipboard",
      "thinking": "Thinking...",
-      "thinkingTime": "Thinking time {{time}}s"
+      "thinkingTime": "Thinking time {{time}}s",
+      "thinkingInProgress": "Thinking in progress..."
    },
    "retrieval": {
      "startPrompt": "Start a retrieval by typing your query below",
--- a/lightrag_webui/src/locales/fr.json
+++ b/lightrag_webui/src/locales/fr.json
@ -337,7 +337,10 @@
  "retrievePanel": {
    "chatMessage": {
      "copyTooltip": "Copier dans le presse-papiers",
-      "copyError": "Échec de la copie du texte dans le presse-papiers"
+      "copyError": "Échec de la copie du texte dans le presse-papiers",
+      "thinking": "Réflexion en cours...",
+      "thinkingTime": "Temps de réflexion {{time}}s",
+      "thinkingInProgress": "Réflexion en cours..."
    },
    "retrieval": {
      "startPrompt": "Démarrez une récupération en tapant votre requête ci-dessous",
--- a/lightrag_webui/src/locales/zh.json
+++ b/lightrag_webui/src/locales/zh.json
@ -339,7 +339,8 @@
      "copyTooltip": "复制到剪贴板",
      "copyError": "复制文本到剪贴板失败",
      "thinking": "正在思考...",
-      "thinkingTime": "思考用时 {{time}} 秒"
+      "thinkingTime": "思考用时 {{time}} 秒",
+      "thinkingInProgress": "思考进行中..."
    },
    "retrieval": {
      "startPrompt": "输入查询开始检索",
--- a/lightrag_webui/src/locales/zh_TW.json
+++ b/lightrag_webui/src/locales/zh_TW.json
@ -337,7 +337,10 @@
  "retrievePanel": {
    "chatMessage": {
      "copyTooltip": "複製到剪貼簿",
-      "copyError": "複製文字到剪貼簿失敗"
+      "copyError": "複製文字到剪貼簿失敗",
+      "thinking": "正在思考...",
+      "thinkingTime": "思考用時 {{time}} 秒",
+      "thinkingInProgress": "思考進行中..."
    },
    "retrieval": {
      "startPrompt": "輸入查詢開始檢索",