chore: workflow performance

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])
 }

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<WorkflowProps> = memo(({
 
   useShortcuts()
   // Initialize workflow node search functionality
-  // useWorkflowSearch()
+  useWorkflowSearch()
 
   // Set up scroll to node event listener using the utility function
   useEffect(() => {