fix(workflow): keep preview run state on panel close

web/app/components/workflow/hooks/use-workflow-interactions.ts

@@ -18,7 +18,7 @@ import {
   useWorkflowReadOnly,
 } from '../hooks'
 import { useStore, useWorkflowStore } from '../store'
-import { BlockEnum, ControlMode } from '../types'
+import { BlockEnum, ControlMode, WorkflowRunningStatus } from '../types'
 import {
   getLayoutByDagre,
   getLayoutForChildNodes,
@@ -36,12 +36,17 @@ export const useWorkflowInteractions = () => {
   const { handleEdgeCancelRunningStatus } = useEdgesInteractionsWithoutSync()
 
   const handleCancelDebugAndPreviewPanel = useCallback(() => {
+    const { workflowRunningData } = workflowStore.getState()
+    const runningStatus = workflowRunningData?.result?.status
+    const isActiveRun = runningStatus === WorkflowRunningStatus.Running || runningStatus === WorkflowRunningStatus.Waiting
     workflowStore.setState({
       showDebugAndPreviewPanel: false,
-      workflowRunningData: undefined,
+      ...(isActiveRun ? {} : { workflowRunningData: undefined }),
     })
-    handleNodeCancelRunningStatus()
+    if (!isActiveRun) {
-    handleEdgeCancelRunningStatus()
+      handleNodeCancelRunningStatus()
+      handleEdgeCancelRunningStatus()
+    }
   }, [workflowStore, handleNodeCancelRunningStatus, handleEdgeCancelRunningStatus])
 
   return {

web/app/components/workflow/panel/debug-and-preview/hooks/use-chat-flow-control.ts

@@ -1,23 +1,21 @@
-import type { RefObject } from 'react'
+import { useCallback } from 'react'
-import { useCallback, useRef } from 'react'
 import { DEFAULT_ITER_TIMES, DEFAULT_LOOP_TIMES } from '../../../constants'
 import { useStore, useWorkflowStore } from '../../../store'
 
 type UseChatFlowControlParams = {
   stopChat?: (taskId: string) => void
-  suggestedQuestionsAbortControllerRef: RefObject<AbortController | null>
 }
 
 export function useChatFlowControl({
   stopChat,
-  suggestedQuestionsAbortControllerRef,
 }: UseChatFlowControlParams) {
   const workflowStore = useWorkflowStore()
   const setIsResponding = useStore(s => s.setIsResponding)
   const resetChatPreview = useStore(s => s.resetChatPreview)
-
+  const setActiveTaskId = useStore(s => s.setActiveTaskId)
-  const hasStopResponded = useRef(false)
+  const setHasStopResponded = useStore(s => s.setHasStopResponded)
-  const taskIdRef = useRef('')
+  const setSuggestedQuestionsAbortController = useStore(s => s.setSuggestedQuestionsAbortController)
+  const invalidateRun = useStore(s => s.invalidateRun)
 
   const { setIterTimes, setLoopTimes } = workflowStore.getState()
 
@@ -26,18 +24,31 @@ export function useChatFlowControl({
   }, [setIsResponding])
 
   const handleStop = useCallback(() => {
-    hasStopResponded.current = true
+    const { activeTaskId, suggestedQuestionsAbortController } = workflowStore.getState()
+    setHasStopResponded(true)
     handleResponding(false)
-    if (stopChat && taskIdRef.current)
+    if (stopChat && activeTaskId)
-      stopChat(taskIdRef.current)
+      stopChat(activeTaskId)
     setIterTimes(DEFAULT_ITER_TIMES)
     setLoopTimes(DEFAULT_LOOP_TIMES)
-    if (suggestedQuestionsAbortControllerRef.current)
+    if (suggestedQuestionsAbortController)
-      suggestedQuestionsAbortControllerRef.current.abort()
+      suggestedQuestionsAbortController.abort()
-  }, [handleResponding, setIterTimes, setLoopTimes, stopChat, suggestedQuestionsAbortControllerRef])
+    setSuggestedQuestionsAbortController(null)
+    setActiveTaskId('')
+    invalidateRun()
+  }, [
+    handleResponding,
+    setIterTimes,
+    setLoopTimes,
+    stopChat,
+    workflowStore,
+    setHasStopResponded,
+    setSuggestedQuestionsAbortController,
+    setActiveTaskId,
+    invalidateRun,
+  ])
 
   const handleRestart = useCallback(() => {
-    taskIdRef.current = ''
     handleStop()
     resetChatPreview()
     setIterTimes(DEFAULT_ITER_TIMES)
@@ -45,8 +56,6 @@ export function useChatFlowControl({
   }, [handleStop, setIterTimes, setLoopTimes, resetChatPreview])
 
   return {
-    hasStopResponded,
-    taskIdRef,
     handleResponding,
     handleStop,
     handleRestart,

web/app/components/workflow/panel/debug-and-preview/hooks/use-chat-message-sender.ts

@@ -1,9 +1,8 @@
-import type { RefObject } from 'react'
 import type { SendCallback, SendParams, UpdateCurrentQAParams } from './types'
 import type { InputForm } from '@/app/components/base/chat/chat/type'
 import type { ChatItem, ChatItemInTree, Inputs } from '@/app/components/base/chat/types'
 import { uniqBy } from 'es-toolkit/compat'
-import { useCallback, useRef } from 'react'
+import { useCallback } from 'react'
 import { useTranslation } from 'react-i18next'
 import { getProcessedInputs } from '@/app/components/base/chat/chat/utils'
 import { getProcessedFiles, getProcessedFilesFromResponse } from '@/app/components/base/file-uploader/utils'
@@ -26,9 +25,6 @@ type UseChatMessageSenderParams = {
     inputs: Inputs
     inputsForm: InputForm[]
   }
-  hasStopResponded: RefObject<boolean>
-  taskIdRef: RefObject<string>
-  suggestedQuestionsAbortControllerRef: RefObject<AbortController | null>
   handleResponding: (responding: boolean) => void
   updateCurrentQAOnTree: (params: UpdateCurrentQAParams) => void
 }
@@ -37,9 +33,6 @@ export function useChatMessageSender({
   threadMessages,
   config,
   formSettings,
-  hasStopResponded,
-  taskIdRef,
-  suggestedQuestionsAbortControllerRef,
   handleResponding,
   updateCurrentQAOnTree,
 }: UseChatMessageSenderParams) {
@@ -54,8 +47,9 @@ export function useChatMessageSender({
   const setConversationId = useStore(s => s.setConversationId)
   const setTargetMessageId = useStore(s => s.setTargetMessageId)
   const setSuggestedQuestions = useStore(s => s.setSuggestedQuestions)
-
+  const setActiveTaskId = useStore(s => s.setActiveTaskId)
-  const activeRunIdRef = useRef(0)
+  const setSuggestedQuestionsAbortController = useStore(s => s.setSuggestedQuestionsAbortController)
+  const startRun = useStore(s => s.startRun)
 
   const handleSend = useCallback((
     params: SendParams,
@@ -66,8 +60,13 @@ export function useChatMessageSender({
       return false
     }
 
-    const runId = ++activeRunIdRef.current
+    const { suggestedQuestionsAbortController } = workflowStore.getState()
-    const isCurrentRun = () => runId === activeRunIdRef.current
+    if (suggestedQuestionsAbortController)
+      suggestedQuestionsAbortController.abort()
+    setSuggestedQuestionsAbortController(null)
+
+    const runId = startRun()
+    const isCurrentRun = () => runId === workflowStore.getState().activeRunId
 
     const parentMessage = threadMessages.find(item => item.id === params.parent_message_id)
 
@@ -109,7 +108,6 @@ export function useChatMessageSender({
     }
 
     handleResponding(true)
-    hasStopResponded.current = false
 
     const { files, inputs, ...restParams } = params
     const bodyParams = {
@@ -143,6 +141,8 @@ export function useChatMessageSender({
       bodyParams,
       {
         onData: (message: string, isFirstMessage: boolean, { conversationId: newConversationId, messageId, taskId }) => {
+          if (!isCurrentRun())
+            return
           responseItem.content = responseItem.content + message
 
           if (messageId && !hasSetResponseId) {
@@ -156,7 +156,7 @@ export function useChatMessageSender({
             setConversationId(newConversationId)
 
           if (taskId)
-            taskIdRef.current = taskId
+            setActiveTaskId(taskId)
           if (messageId)
             responseItem.id = messageId
 
@@ -188,20 +188,25 @@ export function useChatMessageSender({
             return
           }
 
-          if (config?.suggested_questions_after_answer?.enabled && !hasStopResponded.current && onGetSuggestedQuestions) {
+          if (config?.suggested_questions_after_answer?.enabled && !workflowStore.getState().hasStopResponded && onGetSuggestedQuestions) {
             try {
               const result = await onGetSuggestedQuestions(
                 responseItem.id,
-                newAbortController => suggestedQuestionsAbortControllerRef.current = newAbortController,
+                newAbortController => setSuggestedQuestionsAbortController(newAbortController),
               ) as { data: string[] }
               setSuggestedQuestions(result.data)
             }
             catch {
               setSuggestedQuestions([])
             }
+            finally {
+              setSuggestedQuestionsAbortController(null)
+            }
           }
         },
         onMessageEnd: (messageEnd) => {
+          if (!isCurrentRun())
+            return
           responseItem.citation = messageEnd.metadata?.retriever_resources || []
           const processedFilesFromResponse = getProcessedFilesFromResponse(messageEnd.files || [])
           responseItem.allFiles = uniqBy([...(responseItem.allFiles || []), ...(processedFilesFromResponse || [])], 'id')
@@ -214,6 +219,8 @@ export function useChatMessageSender({
           })
         },
         onMessageReplace: (messageReplace) => {
+          if (!isCurrentRun())
+            return
           responseItem.content = messageReplace.answer
         },
         onError() {
@@ -222,17 +229,57 @@ export function useChatMessageSender({
           handleResponding(false)
         },
         onWorkflowStarted: (event) => {
-          taskIdRef.current = workflowHandlers.onWorkflowStarted(event)
+          if (!isCurrentRun())
+            return
+          const taskId = workflowHandlers.onWorkflowStarted(event)
+          if (taskId)
+            setActiveTaskId(taskId)
+        },
+        onWorkflowFinished: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onWorkflowFinished(event)
+        },
+        onIterationStart: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onIterationStart(event)
+        },
+        onIterationFinish: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onIterationFinish(event)
+        },
+        onLoopStart: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onLoopStart(event)
+        },
+        onLoopFinish: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onLoopFinish(event)
+        },
+        onNodeStarted: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onNodeStarted(event)
+        },
+        onNodeRetry: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onNodeRetry(event)
+        },
+        onNodeFinished: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onNodeFinished(event)
+        },
+        onAgentLog: (event) => {
+          if (!isCurrentRun())
+            return
+          workflowHandlers.onAgentLog(event)
         },
-        onWorkflowFinished: workflowHandlers.onWorkflowFinished,
-        onIterationStart: workflowHandlers.onIterationStart,
-        onIterationFinish: workflowHandlers.onIterationFinish,
-        onLoopStart: workflowHandlers.onLoopStart,
-        onLoopFinish: workflowHandlers.onLoopFinish,
-        onNodeStarted: workflowHandlers.onNodeStarted,
-        onNodeRetry: workflowHandlers.onNodeRetry,
-        onNodeFinished: workflowHandlers.onNodeFinished,
-        onAgentLog: workflowHandlers.onAgentLog,
       },
     )
   }, [
@@ -247,12 +294,12 @@ export function useChatMessageSender({
     setTargetMessageId,
     setConversationId,
     setSuggestedQuestions,
+    setActiveTaskId,
+    setSuggestedQuestionsAbortController,
+    startRun,
     fetchInspectVars,
     invalidAllLastRun,
     workflowStore,
-    hasStopResponded,
-    taskIdRef,
-    suggestedQuestionsAbortControllerRef,
   ])
 
   return { handleSend }

web/app/components/workflow/panel/debug-and-preview/hooks/use-chat.ts

@@ -1,7 +1,6 @@
 import type { ChatConfig } from './types'
 import type { InputForm } from '@/app/components/base/chat/chat/type'
 import type { ChatItemInTree, Inputs } from '@/app/components/base/chat/types'
-import { setAutoFreeze } from 'immer'
 import { useEffect, useRef } from 'react'
 import { useStore } from '../../../store'
 import { useChatFlowControl } from './use-chat-flow-control'
@@ -27,8 +26,6 @@ export function useChat(
   const setTargetMessageId = useStore(s => s.setTargetMessageId)
 
   const initialChatTreeRef = useRef(prevChatTree)
-  const suggestedQuestionsAbortControllerRef = useRef<AbortController | null>(null)
-
   useEffect(() => {
     const initialChatTree = initialChatTreeRef.current
     if (!initialChatTree || initialChatTree.length === 0)
@@ -36,24 +33,14 @@ export function useChat(
     updateChatTree(currentChatTree => (currentChatTree.length === 0 ? initialChatTree : currentChatTree))
   }, [updateChatTree])
 
-  useEffect(() => {
-    setAutoFreeze(false)
-    return () => {
-      setAutoFreeze(true)
-    }
-  }, [])
-
   const { updateCurrentQAOnTree } = useChatTreeOperations(updateChatTree)
 
   const {
-    hasStopResponded,
-    taskIdRef,
     handleResponding,
     handleStop,
     handleRestart,
   } = useChatFlowControl({
     stopChat,
-    suggestedQuestionsAbortControllerRef,
   })
 
   const {
@@ -70,9 +57,6 @@ export function useChat(
     threadMessages,
     config,
     formSettings,
-    hasStopResponded,
-    taskIdRef,
-    suggestedQuestionsAbortControllerRef,
     handleResponding,
     updateCurrentQAOnTree,
   })

web/app/components/workflow/store/workflow/chat-preview-slice.ts

@@ -7,6 +7,10 @@ type ChatPreviewState = {
   suggestedQuestions: string[]
   conversationId: string
   isResponding: boolean
+  activeRunId: number
+  activeTaskId: string
+  hasStopResponded: boolean
+  suggestedQuestionsAbortController: AbortController | null
 }
 
 type ChatPreviewActions = {
@@ -16,6 +20,11 @@ type ChatPreviewActions = {
   setSuggestedQuestions: (questions: string[]) => void
   setConversationId: (conversationId: string) => void
   setIsResponding: (isResponding: boolean) => void
+  setActiveTaskId: (taskId: string) => void
+  setHasStopResponded: (hasStopResponded: boolean) => void
+  setSuggestedQuestionsAbortController: (controller: AbortController | null) => void
+  startRun: () => number
+  invalidateRun: () => number
   resetChatPreview: () => void
 }
 
@@ -27,9 +36,13 @@ const initialState: ChatPreviewState = {
   suggestedQuestions: [],
   conversationId: '',
   isResponding: false,
+  activeRunId: 0,
+  activeTaskId: '',
+  hasStopResponded: false,
+  suggestedQuestionsAbortController: null,
 }
 
-export const createChatPreviewSlice: StateCreator<ChatPreviewSliceShape> = set => ({
+export const createChatPreviewSlice: StateCreator<ChatPreviewSliceShape> = (set, get) => ({
   ...initialState,
 
   setChatTree: chatTree => set({ chatTree }),
@@ -49,5 +62,35 @@ export const createChatPreviewSlice: StateCreator<ChatPreviewSliceShape> = set =
 
   setIsResponding: isResponding => set({ isResponding }),
 
-  resetChatPreview: () => set(initialState),
+  setActiveTaskId: activeTaskId => set({ activeTaskId }),
+
+  setHasStopResponded: hasStopResponded => set({ hasStopResponded }),
+
+  setSuggestedQuestionsAbortController: suggestedQuestionsAbortController => set({ suggestedQuestionsAbortController }),
+
+  startRun: () => {
+    const activeRunId = get().activeRunId + 1
+    set({
+      activeRunId,
+      activeTaskId: '',
+      hasStopResponded: false,
+      suggestedQuestionsAbortController: null,
+    })
+    return activeRunId
+  },
+
+  invalidateRun: () => {
+    const activeRunId = get().activeRunId + 1
+    set({
+      activeRunId,
+      activeTaskId: '',
+      suggestedQuestionsAbortController: null,
+    })
+    return activeRunId
+  },
+
+  resetChatPreview: () => set(state => ({
+    ...initialState,
+    activeRunId: state.activeRunId + 1,
+  })),
 })