dify/web/app/components/workflow/hooks/use-collaborative-workflow.ts

import type { Edge, Node } from '../types'
import { useCallback, useMemo } from 'react'
import { useStoreApi } from 'reactflow'
import { collaborationManager } from '../collaboration/core/collaboration-manager'

const sanitizeNodeForBroadcast = (node: Node): Node => {
  if (!node.data) return node

  if (!Object.prototype.hasOwnProperty.call(node.data, 'selected')) return node

  const sanitizedData = { ...node.data }
  delete (sanitizedData as Record<string, unknown>).selected

  return {
    ...node,
    data: sanitizedData,
  }
}

const sanitizeEdgeForBroadcast = (edge: Edge): Edge => {
  if (!edge.data) return edge

  if (!Object.prototype.hasOwnProperty.call(edge.data, '_connectedNodeIsSelected')) return edge

  const sanitizedData = { ...edge.data }
  delete (sanitizedData as Record<string, unknown>)._connectedNodeIsSelected

  return {
    ...edge,
    data: sanitizedData,
  }
}

export const useCollaborativeWorkflow = () => {
  const store = useStoreApi()
  const { setNodes: collabSetNodes, setEdges: collabSetEdges } = collaborationManager

  const setNodes = useCallback(
    (
      newNodes: Node[],
      shouldBroadcast: boolean = true,
      source = 'use-collaborative-workflow:setNodes',
    ) => {
      const { getNodes, setNodes: reactFlowSetNodes } = store.getState()
      if (shouldBroadcast) {
        if (!collaborationManager.canApplyLocalGraphMutation()) return

        const oldNodes = getNodes()
        const sanitizedNodes = new Map<Node, Node>()
        const sanitizeNode = (node: Node) => {
          let sanitized = sanitizedNodes.get(node)
          if (!sanitized) {
            sanitized = sanitizeNodeForBroadcast(node)
            sanitizedNodes.set(node, sanitized)
          }
          return sanitized
        }
        collabSetNodes(oldNodes.map(sanitizeNode), newNodes.map(sanitizeNode), source)
      }
      reactFlowSetNodes(newNodes)
    },
    [store, collabSetNodes],
  )

  const setEdges = useCallback(
    (newEdges: Edge[], shouldBroadcast: boolean = true) => {
      const { edges, setEdges: reactFlowSetEdges } = store.getState()
      if (shouldBroadcast) {
        if (!collaborationManager.canApplyLocalGraphMutation()) return

        collabSetEdges(edges.map(sanitizeEdgeForBroadcast), newEdges.map(sanitizeEdgeForBroadcast))
      }

      reactFlowSetEdges(newEdges)
    },
    [store, collabSetEdges],
  )

  const collaborativeStore = useCallback(() => {
    const state = store.getState()
    return {
      nodes: state.getNodes(),
      edges: state.edges,

      setNodes,
      setEdges,
    }
  }, [store, setNodes, setEdges])

  return useMemo(
    () => ({
      getState: collaborativeStore,
      setNodes,
      setEdges,
    }),
    [collaborativeStore, setEdges, setNodes],
  )
}