From 8fc576870d6e3630b8ce8bf5f12f4554d607b0da Mon Sep 17 00:00:00 2001
From: StyleZhang <jasonapring2015@outlook.com>
Date: Wed, 20 Mar 2024 20:51:59 +0800
Subject: [PATCH] fix

---
 .../workflow/hooks/use-workflow-run.ts        | 34 +++++++++----------
 .../components/workflow/panel/run-history.tsx |  1 -
 2 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-workflow-run.ts b/web/app/components/workflow/hooks/use-workflow-run.ts
index 30100a8d8c..3945918d34 100644
--- a/web/app/components/workflow/hooks/use-workflow-run.ts
+++ b/web/app/components/workflow/hooks/use-workflow-run.ts
@@ -31,7 +31,10 @@ export const useWorkflowRun = () => {
   const reactflow = useReactFlow()
   const featuresStore = useFeaturesStore()
   const nodesExtraData = useNodesExtraData()
-  const { getValidTreeNodes } = useWorkflow()
+  const {
+    getValidTreeNodes,
+    renderTreeFromRecord,
+  } = useWorkflow()
 
   const handleBackupDraft = useCallback(() => {
     const {
@@ -134,12 +137,16 @@ export const useWorkflowRun = () => {
       onNodeFinished,
       ...restCallback
     } = callback || {}
+    workflowStore.setState({ historyWorkflowData: undefined })
+    const { backupDraft } = workflowStore.getState()
     const {
-      getNodes,
       setNodes,
-      edges,
       setEdges,
     } = store.getState()
+    const {
+      nodes,
+      edges,
+    } = backupDraft!
     const appDetail = useAppStore.getState().appDetail
     const workflowContainer = document.getElementById('workflow-container')
 
@@ -176,7 +183,7 @@ export const useWorkflowRun = () => {
             }
           }))
 
-          const newNodes = produce(getNodes(), (draft) => {
+          const newNodes = produce(nodes, (draft) => {
             draft.forEach((node) => {
               node.data._runningStatus = NodeRunningStatus.Waiting
             })
@@ -216,7 +223,6 @@ export const useWorkflowRun = () => {
             } as any)
           }))
 
-          const nodes = getNodes()
           const {
             setViewport,
           } = reactflow
@@ -238,7 +244,7 @@ export const useWorkflowRun = () => {
             const edge = draft.find(edge => edge.target === data.node_id)
 
             if (edge)
-              edge.data._runned = true
+              edge.data = { ...edge.data, _runned: true } as any
           })
           setEdges(newEdges)
 
@@ -258,10 +264,10 @@ export const useWorkflowRun = () => {
               draft.tracing[currentIndex] = data as any
           }))
 
-          const newNodes = produce(getNodes(), (draft) => {
+          const newNodes = produce(nodes, (draft) => {
             const currentNode = draft.find(node => node.id === data.node_id)!
 
-            currentNode.data._runningStatus = data.status
+            currentNode.data._runningStatus = data.status as any
           })
           setNodes(newNodes)
 
@@ -284,23 +290,15 @@ export const useWorkflowRun = () => {
     const publishedWorkflow = await fetchPublishedWorkflow(`/apps/${appDetail?.id}/workflows/publish`)
 
     if (publishedWorkflow) {
-      const {
-        setNodes,
-        setEdges,
-      } = store.getState()
-      const { setViewport } = reactflow
       const nodes = publishedWorkflow.graph.nodes
       const edges = publishedWorkflow.graph.edges
       const viewport = publishedWorkflow.graph.viewport
 
-      setNodes(nodes)
-      setEdges(edges)
-      if (viewport)
-        setViewport(viewport)
+      renderTreeFromRecord(nodes, edges, viewport)
       featuresStore?.setState({ features: publishedWorkflow.features })
       workflowStore.getState().setPublishedAt(publishedWorkflow.created_at)
     }
-  }, [store, reactflow, featuresStore, workflowStore])
+  }, [featuresStore, workflowStore, renderTreeFromRecord])
 
   const handleCheckBeforePublish = useCallback(() => {
     const {
diff --git a/web/app/components/workflow/panel/run-history.tsx b/web/app/components/workflow/panel/run-history.tsx
index ccc19a4769..6364755839 100644
--- a/web/app/components/workflow/panel/run-history.tsx
+++ b/web/app/components/workflow/panel/run-history.tsx
@@ -46,7 +46,6 @@ const RunHistory = () => {
           onClick={() => {
             workflowStore.setState({
               showRunHistory: false,
-              historyWorkflowData: undefined,
             })
             setCurrentLogItem()
             setShowMessageLogModal(false)