Fix/workflow use nodes hooks (#21822)

web/app/components/workflow-app/components/workflow-header/features-trigger.tsx

@@ -3,7 +3,7 @@ import {
   useCallback,
   useMemo,
 } from 'react'
-import { useNodes } from 'reactflow'
+import { useStore as useReactflowStore } from 'reactflow'
 import { RiApps2AddLine } from '@remixicon/react'
 import { useTranslation } from 'react-i18next'
 import {
@@ -22,7 +22,6 @@ import {
   BlockEnum,
   InputVarType,
 } from '@/app/components/workflow/types'
-import type { StartNodeType } from '@/app/components/workflow/nodes/start/types'
 import { useToastContext } from '@/app/components/base/toast'
 import { usePublishWorkflow, useResetWorkflowVersionHistory } from '@/service/use-workflow'
 import type { PublishWorkflowParams } from '@/types/workflow'
@@ -42,9 +41,9 @@ const FeaturesTrigger = () => {
   const publishedAt = useStore(s => s.publishedAt)
   const draftUpdatedAt = useStore(s => s.draftUpdatedAt)
   const toolPublished = useStore(s => s.toolPublished)
-  const nodes = useNodes<StartNodeType>()
-  const startNode = nodes.find(node => node.data.type === BlockEnum.Start)
-  const startVariables = startNode?.data.variables
+  const startVariables = useReactflowStore(
+    s => s.getNodes().find(node => node.data.type === BlockEnum.Start)?.data.variables,
+  )
   const fileSettings = useFeatures(s => s.features.file)
   const variables = useMemo(() => {
     const data = startVariables || []

web/app/components/workflow/nodes/_base/components/next-step/index.tsx

@@ -1,10 +1,10 @@
 import { memo, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
+import { isEqual } from 'lodash-es'
 import {
   getConnectedEdges,
   getOutgoers,
-  useEdges,
-  useStoreApi,
+  useStore,
 } from 'reactflow'
 import { useToolIcon } from '../../../../hooks'
 import BlockIcon from '../../../../block-icon'
@@ -26,12 +26,21 @@ const NextStep = ({
   const { t } = useTranslation()
   const data = selectedNode.data
   const toolIcon = useToolIcon(data)
-  const store = useStoreApi()
   const branches = useMemo(() => {
     return data._targetBranches || []
   }, [data])
-  const edges = useEdges()
-  const outgoers = getOutgoers(selectedNode as Node, store.getState().getNodes(), edges)
+  const edges = useStore(s => s.edges.map(edge => ({
+    id: edge.id,
+    source: edge.source,
+    sourceHandle: edge.sourceHandle,
+    target: edge.target,
+    targetHandle: edge.targetHandle,
+  })), isEqual)
+  const nodes = useStore(s => s.getNodes().map(node => ({
+    id: node.id,
+    data: node.data,
+  })), isEqual)
+  const outgoers = getOutgoers(selectedNode as Node, nodes as Node[], edges)
   const connectedEdges = getConnectedEdges([selectedNode] as Node[], edges).filter(edge => edge.source === selectedNode!.id)
 
   const list = useMemo(() => {

web/app/components/workflow/nodes/_base/components/node-position.tsx

@@ -1,30 +1,39 @@
 import { memo } from 'react'
 import { useTranslation } from 'react-i18next'
+import { useShallow } from 'zustand/react/shallow'
 import { RiCrosshairLine } from '@remixicon/react'
-import type { XYPosition } from 'reactflow'
-import { useReactFlow, useStoreApi } from 'reactflow'
+import { useReactFlow, useStore } from 'reactflow'
 import TooltipPlus from '@/app/components/base/tooltip'
 import { useNodesSyncDraft } from '@/app/components/workflow-app/hooks'
 
 type NodePositionProps = {
-  nodePosition: XYPosition,
-  nodeWidth?: number | null,
-  nodeHeight?: number | null,
+  nodeId: string
 }
 const NodePosition = ({
-  nodePosition,
-  nodeWidth,
-  nodeHeight,
+  nodeId,
 }: NodePositionProps) => {
   const { t } = useTranslation()
   const reactflow = useReactFlow()
-  const store = useStoreApi()
   const { doSyncWorkflowDraft } = useNodesSyncDraft()
+  const {
+    nodePosition,
+    nodeWidth,
+    nodeHeight,
+  } = useStore(useShallow((s) => {
+    const nodes = s.getNodes()
+    const currentNode = nodes.find(node => node.id === nodeId)!
+
+    return {
+      nodePosition: currentNode.position,
+      nodeWidth: currentNode.width,
+      nodeHeight: currentNode.height,
+    }
+  }))
+  const transform = useStore(s => s.transform)
 
   if (!nodePosition || !nodeWidth || !nodeHeight) return null
 
   const workflowContainer = document.getElementById('workflow-container')
-  const { transform } = store.getState()
   const zoom = transform[2]
 
   const { clientWidth, clientHeight } = workflowContainer!

web/app/components/workflow/nodes/_base/components/workflow-panel/index.tsx

@@ -62,15 +62,14 @@ import { Stop } from '@/app/components/base/icons/src/vender/line/mediaAndDevice
 
 type BasePanelProps = {
   children: ReactNode
-} & Node
+  id: Node['id']
+  data: Node['data']
+}
 
 const BasePanel: FC<BasePanelProps> = ({
   id,
   data,
   children,
-  position,
-  width,
-  height,
 }) => {
   const { t } = useTranslation()
   const { showMessageLogModal } = useAppStore(useShallow(state => ({
@@ -330,7 +329,7 @@ const BasePanel: FC<BasePanelProps> = ({
                   </Tooltip>
                 )
               }
-              <NodePosition nodePosition={position} nodeWidth={width} nodeHeight={height}></NodePosition>
+              <NodePosition nodeId={id}></NodePosition>
               <HelpLink nodeType={data.type} />
               <PanelOperator id={id} data={data} showHelpLink={false} />
               <div className='mx-3 h-3.5 w-[1px] bg-divider-regular' />

web/app/components/workflow/nodes/_base/node.tsx

@@ -48,7 +48,9 @@ import useInspectVarsCrud from '../../hooks/use-inspect-vars-crud'
 
 type BaseNodeProps = {
   children: ReactElement
-} & NodeProps
+  id: NodeProps['id']
+  data: NodeProps['data']
+}
 
 const BaseNode: FC<BaseNodeProps> = ({
   id,

web/app/components/workflow/nodes/index.tsx

@@ -14,11 +14,14 @@ import BasePanel from './_base/components/workflow-panel'
 
 const CustomNode = (props: NodeProps) => {
   const nodeData = props.data
-  const NodeComponent = NodeComponentMap[nodeData.type]
+  const NodeComponent = useMemo(() => NodeComponentMap[nodeData.type], [nodeData.type])
 
   return (
     <>
-      <BaseNode {...props}>
+      <BaseNode
+        id={props.id}
+        data={props.data}
+      >
         <NodeComponent />
       </BaseNode>
     </>
@@ -26,7 +29,12 @@ const CustomNode = (props: NodeProps) => {
 }
 CustomNode.displayName = 'CustomNode'
 
-export const Panel = memo((props: Node) => {
+export type PanelProps = {
+  type: Node['type']
+  id: Node['id']
+  data: Node['data']
+}
+export const Panel = memo((props: PanelProps) => {
   const nodeClass = props.type
   const nodeData = props.data
   const PanelComponent = useMemo(() => {
@@ -38,7 +46,11 @@ export const Panel = memo((props: Node) => {
 
   if (nodeClass === CUSTOM_NODE) {
     return (
-      <BasePanel key={props.id} {...props}>
+      <BasePanel
+        key={props.id}
+        id={props.id}
+        data={props.data}
+      >
         <PanelComponent />
       </BasePanel>
     )

web/app/components/workflow/panel/index.tsx

@@ -1,7 +1,7 @@
 import type { FC } from 'react'
+import { useShallow } from 'zustand/react/shallow'
 import { memo, useCallback, useEffect, useRef } from 'react'
-import { useNodes } from 'reactflow'
-import type { CommonNodeType } from '../types'
+import { useStore as useReactflow } from 'reactflow'
 import { Panel as NodePanel } from '../nodes'
 import { useStore } from '../store'
 import EnvPanel from './env-panel'
@@ -61,8 +61,18 @@ const useResizeObserver = (
 const Panel: FC<PanelProps> = ({
   components,
 }) => {
-  const nodes = useNodes<CommonNodeType>()
-  const selectedNode = nodes.find(node => node.data.selected)
+  const selectedNode = useReactflow(useShallow((s) => {
+    const nodes = s.getNodes()
+    const currentNode = nodes.find(node => node.data.selected)
+
+    if (currentNode) {
+      return {
+        id: currentNode.id,
+        type: currentNode.type,
+        data: currentNode.data,
+      }
+    }
+  }))
   const showEnvPanel = useStore(s => s.showEnvPanel)
   const isRestoring = useStore(s => s.isRestoring)
   const showWorkflowVersionHistoryPanel = useStore(s => s.showWorkflowVersionHistoryPanel)