From 8ae46a8a1428a689187a87e722e33b4c6eea2c30 Mon Sep 17 00:00:00 2001
From: StyleZhang <jasonapring2015@outlook.com>
Date: Tue, 12 Mar 2024 13:25:57 +0800
Subject: [PATCH] fix

---
 web/app/components/workflow/hooks.ts          | 32 +++++++++++++++++--
 .../components/class-list.tsx                 |  9 ++++--
 .../nodes/question-classifier/panel.tsx       |  1 +
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/web/app/components/workflow/hooks.ts b/web/app/components/workflow/hooks.ts
index e2951a82df..89d0a2b639 100644
--- a/web/app/components/workflow/hooks.ts
+++ b/web/app/components/workflow/hooks.ts
@@ -80,7 +80,7 @@ export const useWorkflow = () => {
 
     if (appId) {
       const features = featuresStore!.getState().features
-      const nodes = produce(getNodes(), (draft) => {
+      const producedNodes = produce(getNodes(), (draft) => {
         draft.forEach((node) => {
           Object.keys(node.data).forEach((key) => {
             if (key.startsWith('_'))
@@ -88,12 +88,17 @@ export const useWorkflow = () => {
           })
         })
       })
+      const producedEdges = produce(edges, (draft) => {
+        draft.forEach((edge) => {
+          delete edge.data
+        })
+      })
       syncWorkflowDraft({
         url: `/apps/${appId}/workflows/draft`,
         params: {
           graph: {
-            nodes,
-            edges,
+            nodes: producedNodes,
+            edges: producedEdges,
             viewport: getViewport(),
           },
           features: {
@@ -599,6 +604,26 @@ export const useWorkflow = () => {
     setEdges(newEdges)
   }, [store])
 
+  const handleEdgeDeleteByDeleteBranch = useCallback((nodeId: string, branchId: string) => {
+    const { runningStatus } = useStore.getState()
+
+    if (runningStatus)
+      return
+
+    const {
+      edges,
+      setEdges,
+    } = store.getState()
+    const newEdges = produce(edges, (draft) => {
+      const index = draft.findIndex(edge => edge.source === nodeId && edge.sourceHandle === branchId)
+
+      if (index > -1)
+        draft.splice(index, 1)
+    })
+    setEdges(newEdges)
+    handleSyncWorkflowDraft()
+  }, [store, handleSyncWorkflowDraft])
+
   const handleEdgeDelete = useCallback(() => {
     const { runningStatus } = useStore.getState()
 
@@ -670,6 +695,7 @@ export const useWorkflow = () => {
 
     handleEdgeEnter,
     handleEdgeLeave,
+    handleEdgeDeleteByDeleteBranch,
     handleEdgeDelete,
     handleEdgesChange,
 
diff --git a/web/app/components/workflow/nodes/question-classifier/components/class-list.tsx b/web/app/components/workflow/nodes/question-classifier/components/class-list.tsx
index 53f9c0079b..2fd1b8adae 100644
--- a/web/app/components/workflow/nodes/question-classifier/components/class-list.tsx
+++ b/web/app/components/workflow/nodes/question-classifier/components/class-list.tsx
@@ -3,6 +3,7 @@ import type { FC } from 'react'
 import React, { useCallback } from 'react'
 import produce from 'immer'
 import { useTranslation } from 'react-i18next'
+import { useWorkflow } from '../../../hooks'
 import AddButton from '../../_base/components/add-button'
 import Item from './class-item'
 import type { Topic } from '@/app/components/workflow/nodes/question-classifier/types'
@@ -10,15 +11,18 @@ import type { Topic } from '@/app/components/workflow/nodes/question-classifier/
 const i18nPrefix = 'workflow.nodes.questionClassifiers'
 
 type Props = {
+  id: string
   list: Topic[]
   onChange: (list: Topic[]) => void
 }
 
 const ClassList: FC<Props> = ({
+  id,
   list,
   onChange,
 }) => {
   const { t } = useTranslation()
+  const { handleEdgeDeleteByDeleteBranch } = useWorkflow()
 
   const handleClassChange = useCallback((index: number) => {
     return (value: Topic) => {
@@ -34,16 +38,17 @@ const ClassList: FC<Props> = ({
       draft.push({ id: `${Date.now()}`, name: '' })
     })
     onChange(newList)
-  }, [list, onChange, t])
+  }, [list, onChange])
 
   const handleRemoveClass = useCallback((index: number) => {
     return () => {
+      handleEdgeDeleteByDeleteBranch(id, list[index].id)
       const newList = produce(list, (draft) => {
         draft.splice(index, 1)
       })
       onChange(newList)
     }
-  }, [list, onChange])
+  }, [list, onChange, handleEdgeDeleteByDeleteBranch, id])
 
   // Todo Remove; edit topic name
   return (
diff --git a/web/app/components/workflow/nodes/question-classifier/panel.tsx b/web/app/components/workflow/nodes/question-classifier/panel.tsx
index 07a220d379..210ae992cc 100644
--- a/web/app/components/workflow/nodes/question-classifier/panel.tsx
+++ b/web/app/components/workflow/nodes/question-classifier/panel.tsx
@@ -72,6 +72,7 @@ const Panel: FC<NodePanelProps<QuestionClassifierNodeType>> = ({
           title={t(`${i18nPrefix}.class`)}
         >
           <ClassList
+            id={id}
             list={inputs.classes}
             onChange={handleTopicsChange} />
         </Field>