From 9f024835aa1f7d79a29a10f232feedefbdabe42e Mon Sep 17 00:00:00 2001
From: StyleZhang <jasonapring2015@outlook.com>
Date: Tue, 19 Mar 2024 20:34:10 +0800
Subject: [PATCH] chat

---
 .../workflow/hooks/use-workflow-run.ts        | 38 ++++++++++++++++---
 .../workflow/panel/debug-and-preview/hooks.ts |  8 ++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-workflow-run.ts b/web/app/components/workflow/hooks/use-workflow-run.ts
index a40dbb7334..7d1b694a51 100644
--- a/web/app/components/workflow/hooks/use-workflow-run.ts
+++ b/web/app/components/workflow/hooks/use-workflow-run.ts
@@ -102,7 +102,17 @@ export const useWorkflowRun = () => {
     }
   }, [store, handleLoadBackupDraft, handleBackupDraft, workflowStore])
 
-  const handleRun = useCallback((params: any, callback?: IOtherOptions) => {
+  const handleRun = useCallback((
+    params: any,
+    callback?: IOtherOptions,
+  ) => {
+    const {
+      onWorkflowStarted,
+      onWorkflowFinished,
+      onNodeStarted,
+      onNodeFinished,
+      ...restCallback
+    } = callback || {}
     const {
       getNodes,
       setNodes,
@@ -130,7 +140,8 @@ export const useWorkflowRun = () => {
         body: params,
       },
       {
-        onWorkflowStarted: ({ task_id, workflow_run_id, data }) => {
+        onWorkflowStarted: (params) => {
+          const { task_id, workflow_run_id, data } = params
           workflowStore.setState({ runningStatus: WorkflowRunningStatus.Running })
           workflowStore.setState({ taskId: task_id })
           workflowStore.setState({ currentSequenceNumber: data.sequence_number })
@@ -141,11 +152,19 @@ export const useWorkflowRun = () => {
             })
           })
           setNodes(newNodes)
+
+          if (onWorkflowStarted)
+            onWorkflowStarted(params)
         },
-        onWorkflowFinished: ({ data }) => {
+        onWorkflowFinished: (params) => {
+          const { data } = params
           workflowStore.setState({ runningStatus: data.status as WorkflowRunningStatus })
+
+          if (onWorkflowFinished)
+            onWorkflowFinished(params)
         },
-        onNodeStarted: ({ data }) => {
+        onNodeStarted: (params) => {
+          const { data } = params
           const nodes = getNodes()
           const {
             setViewport,
@@ -171,16 +190,23 @@ export const useWorkflowRun = () => {
               edge.data._runned = true
           })
           setEdges(newEdges)
+
+          if (onNodeStarted)
+            onNodeStarted(params)
         },
-        onNodeFinished: ({ data }) => {
+        onNodeFinished: (params) => {
+          const { data } = params
           const newNodes = produce(getNodes(), (draft) => {
             const currentNode = draft.find(node => node.id === data.node_id)!
 
             currentNode.data._runningStatus = data.status
           })
           setNodes(newNodes)
+
+          if (onNodeFinished)
+            onNodeFinished(params)
         },
-        ...callback,
+        ...restCallback,
       },
     )
   }, [store, reactflow, workflowStore])
diff --git a/web/app/components/workflow/panel/debug-and-preview/hooks.ts b/web/app/components/workflow/panel/debug-and-preview/hooks.ts
index f8efb2695e..dfa4ce841d 100644
--- a/web/app/components/workflow/panel/debug-and-preview/hooks.ts
+++ b/web/app/components/workflow/panel/debug-and-preview/hooks.ts
@@ -261,6 +261,14 @@ export const useChat = (
           })
           handleUpdateChatList(newChatList)
         },
+        onWorkflowStarted: () => {
+        },
+        onWorkflowFinished: () => {
+        },
+        onNodeStarted: () => {
+        },
+        onNodeFinished: () => {
+        },
       },
     )
   }, [handleRun, handleResponding, handleUpdateChatList, notify, t, updateCurrentQA, config.suggested_questions_after_answer?.enabled])