diff --git a/web/app/components/workflow/hooks/use-find-node.ts b/web/app/components/workflow/hooks/use-find-node.ts index f4da759469..1e2cfbf15d 100644 --- a/web/app/components/workflow/hooks/use-find-node.ts +++ b/web/app/components/workflow/hooks/use-find-node.ts @@ -1,16 +1,36 @@ +import { useMemo } from 'react' import { useStore } from 'reactflow' +import { useShallow } from 'zustand/react/shallow' import { isSystemVar } from '@/app/components/workflow/nodes/_base/components/variable/utils' import { BlockEnum } from '@/app/components/workflow/types' import type { Node, ValueSelector } from '@/app/components/workflow/types' export const useFindNode = (valueSelector: ValueSelector = [], nodes: Node[] = []) => { const nodeFromOuterNodes = nodes.find(node => node.id === valueSelector[0]) - const node = useStore((s) => { - const nodes = s.getNodes() - if (isSystemVar(valueSelector)) - return nodes.find(node => node.data.type === BlockEnum.Start) - return nodes.find(node => node.id === valueSelector[0]) - }) + const node = useStore(useShallow((s) => { + if (isSystemVar(valueSelector)) { + for (const n of s.nodeInternals.values()) { + if (n?.data?.type === BlockEnum.Start) { + return { + id: n.id, + data: n.data, + } + } + } + } + else { + if (!!valueSelector.length) { + const id = s.nodeInternals.get(valueSelector[0])?.id + const data = s.nodeInternals.get(valueSelector[0])?.data + if (id && data) { + return { + id, + data, + } + } + } + } + })) - return nodeFromOuterNodes || node + return useMemo(() => nodeFromOuterNodes || node, [nodeFromOuterNodes, node]) } diff --git a/web/app/components/workflow/index.tsx b/web/app/components/workflow/index.tsx index 579b3b6d62..75c4d51390 100644 --- a/web/app/components/workflow/index.tsx +++ b/web/app/components/workflow/index.tsx @@ -61,6 +61,7 @@ import { CUSTOM_SIMPLE_NODE } from './simple-node/constants' import CustomDataSourceEmptyNode from './nodes/data-source-empty' import { CUSTOM_DATA_SOURCE_EMPTY_NODE } from './nodes/data-source-empty/constants' import Operator from './operator' +import { useWorkflowSearch } from './hooks/use-workflow-search' import Control from './operator/control' import CustomEdge from './custom-edge' import CustomConnectionLine from './custom-connection-line' @@ -289,7 +290,7 @@ export const Workflow: FC = memo(({ useShortcuts() // Initialize workflow node search functionality - // useWorkflowSearch() + useWorkflowSearch() // Set up scroll to node event listener using the utility function useEffect(() => {