From a4041cb40bb480c03ce1863b3c86fca87ef4015c Mon Sep 17 00:00:00 2001
From: zxhlyh <jasonapring2015@outlook.com>
Date: Wed, 5 Jun 2024 14:00:47 +0800
Subject: [PATCH] fix: end node limit in next step (#4945)

---
 .../nodes/_base/components/next-step/add.tsx        | 13 ++++++-------
 .../nodes/_base/components/next-step/index.tsx      |  4 ++--
 web/app/components/workflow/nodes/_base/panel.tsx   |  7 +++----
 web/app/components/workflow/operator/add-block.tsx  |  7 +++----
 4 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/web/app/components/workflow/nodes/_base/components/next-step/add.tsx b/web/app/components/workflow/nodes/_base/components/next-step/add.tsx
index 81b0207dec..06718a6c22 100644
--- a/web/app/components/workflow/nodes/_base/components/next-step/add.tsx
+++ b/web/app/components/workflow/nodes/_base/components/next-step/add.tsx
@@ -4,34 +4,33 @@ import {
 } from 'react'
 import { useTranslation } from 'react-i18next'
 import {
-  useNodesExtraData,
+  useAvailableBlocks,
   useNodesInteractions,
   useNodesReadOnly,
 } from '@/app/components/workflow/hooks'
 import BlockSelector from '@/app/components/workflow/block-selector'
 import { Plus } from '@/app/components/base/icons/src/vender/line/general'
 import type {
-  BlockEnum,
+  CommonNodeType,
   OnSelectBlock,
 } from '@/app/components/workflow/types'
 
 type AddProps = {
   nodeId: string
-  nodeType: BlockEnum
+  nodeData: CommonNodeType
   sourceHandle: string
   branchName?: string
 }
 const Add = ({
   nodeId,
-  nodeType,
+  nodeData,
   sourceHandle,
   branchName,
 }: AddProps) => {
   const { t } = useTranslation()
   const { handleNodeAdd } = useNodesInteractions()
-  const nodesExtraData = useNodesExtraData()
   const { nodesReadOnly } = useNodesReadOnly()
-  const availableNextNodes = nodesExtraData[nodeType].availableNextNodes
+  const { availableNextBlocks } = useAvailableBlocks(nodeData.type, nodeData.isInIteration)
 
   const handleSelect = useCallback<OnSelectBlock>((type, toolDefaultValue) => {
     handleNodeAdd(
@@ -82,7 +81,7 @@ const Add = ({
       offset={0}
       trigger={renderTrigger}
       popupClassName='!w-[328px]'
-      availableBlocksTypes={availableNextNodes}
+      availableBlocksTypes={availableNextBlocks}
     />
   )
 }
diff --git a/web/app/components/workflow/nodes/_base/components/next-step/index.tsx b/web/app/components/workflow/nodes/_base/components/next-step/index.tsx
index 73ceb50b52..a6fd940f63 100644
--- a/web/app/components/workflow/nodes/_base/components/next-step/index.tsx
+++ b/web/app/components/workflow/nodes/_base/components/next-step/index.tsx
@@ -54,7 +54,7 @@ const NextStep = ({
           !nodeWithBranches && !outgoers.length && (
             <Add
               nodeId={selectedNode!.id}
-              nodeType={selectedNode!.data.type}
+              nodeData={selectedNode!.data}
               sourceHandle='source'
             />
           )
@@ -85,7 +85,7 @@ const NextStep = ({
                       <Add
                         key={branch.id}
                         nodeId={selectedNode!.id}
-                        nodeType={selectedNode!.data.type}
+                        nodeData={selectedNode!.data}
                         sourceHandle={branch.id}
                         branchName={branch.name}
                       />
diff --git a/web/app/components/workflow/nodes/_base/panel.tsx b/web/app/components/workflow/nodes/_base/panel.tsx
index a6d3d6a2aa..c5072dce12 100644
--- a/web/app/components/workflow/nodes/_base/panel.tsx
+++ b/web/app/components/workflow/nodes/_base/panel.tsx
@@ -23,8 +23,8 @@ import {
 } from '@/app/components/base/icons/src/vender/line/general'
 import BlockIcon from '@/app/components/workflow/block-icon'
 import {
+  useAvailableBlocks,
   useNodeDataUpdate,
-  useNodesExtraData,
   useNodesInteractions,
   useNodesReadOnly,
   useNodesSyncDraft,
@@ -57,8 +57,7 @@ const BasePanel: FC<BasePanelProps> = ({
   const { handleNodeSelect } = useNodesInteractions()
   const { handleSyncWorkflowDraft } = useNodesSyncDraft()
   const { nodesReadOnly } = useNodesReadOnly()
-  const nodesExtraData = useNodesExtraData()
-  const availableNextNodes = nodesExtraData[data.type].availableNextNodes
+  const { availableNextBlocks } = useAvailableBlocks(data.type, data.isInIteration)
   const toolIcon = useToolIcon(data)
 
   const handleResize = useCallback((width: number) => {
@@ -157,7 +156,7 @@ const BasePanel: FC<BasePanelProps> = ({
           {cloneElement(children, { id, data })}
         </div>
         {
-          !!availableNextNodes.length && (
+          !!availableNextBlocks.length && (
             <div className='p-4 border-t-[0.5px] border-t-black/5'>
               <div className='flex items-center mb-1 text-gray-700 text-[13px] font-semibold'>
                 {t('workflow.panel.nextStep').toLocaleUpperCase()}
diff --git a/web/app/components/workflow/operator/add-block.tsx b/web/app/components/workflow/operator/add-block.tsx
index 5c1514e4c9..019f33bdae 100644
--- a/web/app/components/workflow/operator/add-block.tsx
+++ b/web/app/components/workflow/operator/add-block.tsx
@@ -11,7 +11,7 @@ import {
   generateNewNode,
 } from '../utils'
 import {
-  useNodesExtraData,
+  useAvailableBlocks,
   useNodesReadOnly,
   usePanelInteractions,
 } from '../hooks'
@@ -38,11 +38,10 @@ const AddBlock = ({
   const { t } = useTranslation()
   const store = useStoreApi()
   const workflowStore = useWorkflowStore()
-  const nodesExtraData = useNodesExtraData()
   const { nodesReadOnly } = useNodesReadOnly()
   const { handlePaneContextmenuCancel } = usePanelInteractions()
   const [open, setOpen] = useState(false)
-  const availableNextNodes = nodesExtraData[BlockEnum.Start].availableNextNodes
+  const { availableNextBlocks } = useAvailableBlocks(BlockEnum.Start, false)
 
   const handleOpenChange = useCallback((open: boolean) => {
     setOpen(open)
@@ -102,7 +101,7 @@ const AddBlock = ({
       }}
       trigger={renderTrigger || renderTriggerElement}
       popupClassName='!min-w-[256px]'
-      availableBlocksTypes={availableNextNodes}
+      availableBlocksTypes={availableNextBlocks}
     />
   )
 }