dify/web/app/components/workflow/hooks/use-nodes-interactions-without-sync.ts

import type { RefObject } from 'react'
import { produce } from 'immer'
import { useCallback } from 'react'
import { useStoreApi } from 'reactflow'
import { NodeRunningStatus } from '../types'

export const useNodesInteractionsWithoutSync = (isMountedRef?: RefObject<boolean>) => {
  const store = useStoreApi()

  const handleNodeCancelRunningStatus = useCallback(() => {
    if (isMountedRef && isMountedRef.current === false)
      return

    const {
      getNodes,
      setNodes,
    } = store.getState()

    const nodes = getNodes()
    const newNodes = produce(nodes, (draft) => {
      draft.forEach((node) => {
        node.data._runningStatus = undefined
        node.data._waitingRun = false
      })
    })
    setNodes(newNodes)
  }, [store, isMountedRef])

  const handleCancelAllNodeSuccessStatus = useCallback(() => {
    if (isMountedRef && isMountedRef.current === false)
      return

    const {
      getNodes,
      setNodes,
    } = store.getState()

    const nodes = getNodes()
    const newNodes = produce(nodes, (draft) => {
      draft.forEach((node) => {
        if (node.data._runningStatus === NodeRunningStatus.Succeeded)
          node.data._runningStatus = undefined
      })
    })
    setNodes(newNodes)
  }, [store, isMountedRef])

  const handleCancelNodeSuccessStatus = useCallback((nodeId: string) => {
    if (isMountedRef && isMountedRef.current === false)
      return

    const {
      getNodes,
      setNodes,
    } = store.getState()

    const newNodes = produce(getNodes(), (draft) => {
      const node = draft.find(n => n.id === nodeId)
      if (node && node.data._runningStatus === NodeRunningStatus.Succeeded) {
        node.data._runningStatus = undefined
        node.data._waitingRun = false
      }
    })
    setNodes(newNodes)
  }, [store, isMountedRef])

  return {
    handleNodeCancelRunningStatus,
    handleCancelAllNodeSuccessStatus,
    handleCancelNodeSuccessStatus,
  }
}