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

import type { WorkflowFeaturesConfigPayload } from '@dify/contracts/api/console/apps/types.gen'
import type {
  SyncDraftCallback,
  SyncDraftOptions,
  SyncDraftResult,
} from '@/app/components/workflow/hooks-store'
import { safe } from '@orpc/client'
import { skipToken, useQuery } from '@tanstack/react-query'
import { produce } from 'immer'
import { useAtomValue } from 'jotai'
import { useCallback, useEffect, useRef } from 'react'
import { useTranslation } from 'react-i18next'
import { useStoreApi } from 'reactflow'
import { useFeaturesStore } from '@/app/components/base/features/hooks'
import { collaborationManager } from '@/app/components/workflow/collaboration/core/collaboration-manager'
import {
  editorCollaborationEnabledAtom,
  workflowAppIdentityAtom,
} from '@/app/components/workflow/editor-resource/state'
import { useSerialAsyncCallback } from '@/app/components/workflow/hooks/use-serial-async-callback'
import {
  useNodesReadOnly,
  useNodesReadOnlyByCanEdit,
} from '@/app/components/workflow/hooks/use-workflow'
import {
  isAgentV2NodeData,
  needsInlineAgentBindingCreation,
} from '@/app/components/workflow/nodes/agent-v2/types'
import { useWorkflowStore } from '@/app/components/workflow/store'
import { BlockEnum } from '@/app/components/workflow/types'
import { normalizeWorkflowNodes } from '@/app/components/workflow/utils/normalize-workflow-nodes'
import { API_PREFIX } from '@/config'
import { isAppDeletingOrDeleted } from '@/service/app-deletion'
import { consoleClient, consoleQuery } from '@/service/console'
import { postWithKeepalive } from '@/service/fetch'
import { updateEnvironmentVariables } from '@/service/workflow'
import { useWorkflowRefreshDraft } from './use-workflow-refresh-draft'

const shouldSkipDraftSync = (appId: string | undefined, isWorkflowDataLoaded: boolean) =>
  !appId || !isWorkflowDataLoaded || isAppDeletingOrDeleted(appId)

const isEmptyGraph = (graph: { nodes: unknown[]; edges: unknown[] }) =>
  graph.nodes.length === 0 && graph.edges.length === 0

const useNodesSyncDraftBase = (getNodesReadOnly: () => boolean) => {
  const { t } = useTranslation(['workflow'] as const)
  const store = useStoreApi()
  const workflowStore = useWorkflowStore()
  const featuresStore = useFeaturesStore()
  const appId = useAtomValue(workflowAppIdentityAtom).appId
  const { data: appMode } = useQuery(
    consoleQuery.apps.byAppId.get.queryOptions({
      input: appId ? { params: { app_id: appId } } : skipToken,
      select: (app) => app.mode,
    }),
  )
  const { handleRefreshWorkflowDraft } = useWorkflowRefreshDraft()
  const isCollaborationEnabled = useAtomValue(editorCollaborationEnabledAtom)
  const isMountedRef = useRef(false)
  const emptySyncPendingRef = useRef(false)
  const cancelConfirmationRef = useRef<(() => void) | undefined>(undefined)

  useEffect(() => {
    isMountedRef.current = true
    return () => {
      isMountedRef.current = false
      cancelConfirmationRef.current?.()
    }
  }, [])

  const getPostParams = useCallback(() => {
    const { getNodes, edges, transform } = store.getState()
    const allNodes = getNodes()
    const nodes = allNodes.filter((node) => {
      if (node.data?.type === BlockEnum.StartPlaceholder) return false

      if (!node.data?._isTempNode) return true

      return isAgentV2NodeData(node.data) && needsInlineAgentBindingCreation(node.data)
    })
    const skippedNodeIds = new Set(
      allNodes
        .filter((node) => {
          if (node.data?.type === BlockEnum.StartPlaceholder) return true

          if (!node.data?._isTempNode) return false

          return !(isAgentV2NodeData(node.data) && needsInlineAgentBindingCreation(node.data))
        })
        .map((node) => node.id),
    )
    const [x, y, zoom] = transform
    const { conversationVariables, syncWorkflowDraftHash, isWorkflowDataLoaded } =
      workflowStore.getState()

    if (!appId || shouldSkipDraftSync(appId, isWorkflowDataLoaded)) return null

    const features = featuresStore!.getState().features
    const producedNodes = produce(nodes, (draft) => {
      draft.forEach((node) => {
        Object.keys(node.data).forEach((key) => {
          if (key.startsWith('_')) delete node.data[key]
        })
      })
    })
    const producedEdges = produce(
      edges.filter(
        (edge) =>
          !edge.data?._isTemp &&
          !skippedNodeIds.has(edge.source) &&
          !skippedNodeIds.has(edge.target),
      ),
      (draft) => {
        draft.forEach((edge) => {
          Object.keys(edge.data).forEach((key) => {
            if (key.startsWith('_')) delete edge.data[key]
          })
        })
      },
    )
    const featuresPayload: WorkflowFeaturesConfigPayload = {
      opening_statement: features.opening?.enabled ? features.opening?.opening_statement || '' : '',
      suggested_questions: features.opening?.enabled
        ? features.opening?.suggested_questions || []
        : [],
      suggested_questions_after_answer: features.suggested,
      text_to_speech: features.text2speech,
      speech_to_text: features.speech2text,
      retriever_resource: features.citation,
      sensitive_word_avoidance: features.moderation,
      file_upload: features.file,
    }

    return {
      appId,
      url: `/apps/${appId}/workflows/draft`,
      params: {
        graph: {
          nodes: normalizeWorkflowNodes(producedNodes, appMode),
          edges: producedEdges,
          viewport: {
            x,
            y,
            zoom,
          },
        },
        features: featuresPayload,
        conversation_variables: conversationVariables,
        hash: syncWorkflowDraftHash,
        ...(isCollaborationEnabled ? { _is_collaborative: true } : {}),
      },
    }
  }, [store, featuresStore, workflowStore, isCollaborationEnabled, appMode, appId])

  const syncWorkflowDraftWhenPageClose = useCallback(() => {
    if (getNodesReadOnly()) return
    if (isCollaborationEnabled && workflowStore.getState().loroDraft) return

    const canPersistOnPageClose =
      !isCollaborationEnabled ||
      collaborationManager.canFlushGraphOnPageClose() ||
      collaborationManager.canUseLocalDraftFallback()
    if (!canPersistOnPageClose) return

    const postParams = getPostParams()

    // Page-close saves cannot wait for the user's consent to clear the canvas.
    if (postParams && !isEmptyGraph(postParams.params.graph))
      postWithKeepalive(`${API_PREFIX}${postParams.url}`, postParams.params)
  }, [getPostParams, getNodesReadOnly, isCollaborationEnabled, workflowStore])

  const confirmEmptyDraft = useCallback(() => {
    const { showConfirm, setShowConfirm } = workflowStore.getState()
    if (!isMountedRef.current || showConfirm) return Promise.resolve(false)
    return new Promise<boolean>((resolve) => {
      let settled = false
      const finish = (confirmed: boolean) => {
        if (settled) return
        settled = true
        cancelConfirmationRef.current = undefined
        setShowConfirm(undefined)
        resolve(confirmed)
      }
      cancelConfirmationRef.current = () => finish(false)
      setShowConfirm({
        title: t(($) => $['common.clearCanvasConfirmTitle'], { ns: 'workflow' }),
        desc: t(($) => $['common.clearCanvasConfirmDescription'], { ns: 'workflow' }),
        onConfirm: () => finish(true),
        onCancel: () => finish(false),
      })
    })
  }, [t, workflowStore])

  const performLocalSync = useCallback(
    async (
      baseParams: NonNullable<ReturnType<typeof getPostParams>>,
      notRefreshWhenSyncError?: boolean,
      callback?: SyncDraftCallback,
      options?: SyncDraftOptions,
    ): Promise<SyncDraftResult | null> => {
      if (getNodesReadOnly()) return null
      const { isWorkflowDataLoaded } = workflowStore.getState()
      if (shouldSkipDraftSync(appId, isWorkflowDataLoaded)) {
        callback?.onSettled?.()
        return null
      }

      if (isCollaborationEnabled && !collaborationManager.canPersistLocalGraph()) {
        callback?.onSettled?.()
        return null
      }

      const emptyGraph = isEmptyGraph(baseParams.params.graph)
      const canSaveEmptyGraph = () => {
        const currentParams = getPostParams()
        return (
          isMountedRef.current &&
          !getNodesReadOnly() &&
          (!isCollaborationEnabled || collaborationManager.canPersistLocalGraph()) &&
          currentParams?.url === baseParams.url &&
          isEmptyGraph(currentParams.params.graph)
        )
      }
      if (emptyGraph && !canSaveEmptyGraph()) {
        callback?.onSettled?.()
        return null
      }

      if (emptyGraph && isCollaborationEnabled) {
        // CRDT updates do not carry the persisted revision from another collaborator's save.
        const hashBeforeRefresh = workflowStore.getState().syncWorkflowDraftHash
        const [error, draft] = await safe(
          consoleClient.apps.byAppId.workflows.draft.get(
            { params: { app_id: baseParams.appId } },
            { context: { silent: true } },
          ),
        )
        if (!canSaveEmptyGraph()) {
          callback?.onSettled?.()
          return null
        }
        if (error || !Array.isArray(draft.graph.nodes) || !Array.isArray(draft.graph.edges)) {
          callback?.onError?.()
          callback?.onSettled?.()
          return null
        }

        const { syncWorkflowDraftHash, setSyncWorkflowDraftHash } = workflowStore.getState()
        // Another hook may have saved a newer revision while this request was in flight.
        if (syncWorkflowDraftHash === hashBeforeRefresh)
          setSyncWorkflowDraftHash(draft.hash, {
            nodes: draft.graph.nodes,
            edges: draft.graph.edges,
          })
      }

      // Check after queued saves finish so a newly persisted nonempty draft is protected too.
      if (emptyGraph && workflowStore.getState().syncWorkflowDraftIsEmpty !== true) {
        const confirmed = await confirmEmptyDraft()

        if (!confirmed || !canSaveEmptyGraph()) {
          callback?.onSettled?.()
          return null
        }
      }

      const { setSyncWorkflowDraftHash, setDraftUpdatedAt } = workflowStore.getState()

      try {
        const latestHash = workflowStore.getState().syncWorkflowDraftHash

        const body = {
          ...baseParams.params,
          hash: latestHash || null,
          ...(emptyGraph ? { force: true } : {}),
          ...(options?.environmentVariablePatch
            ? {
                environment_variable_patch: {
                  environment_variables: options.environmentVariablePatch.environmentVariables,
                  deleted_environment_variable_ids:
                    options.environmentVariablePatch.deletedEnvironmentVariableIds,
                },
              }
            : {}),
        }

        const [error, res, isDefined] = await safe(
          consoleClient.apps.byAppId.workflows.draft.post(
            { params: { app_id: baseParams.appId }, body },
            { context: { silent: true } },
          ),
        )
        if (error) {
          const { isWorkflowDataLoaded } = workflowStore.getState()
          if (shouldSkipDraftSync(appId, isWorkflowDataLoaded)) return null

          if (isDefined && error.code === 'draft_workflow_not_sync' && !notRefreshWhenSyncError)
            handleRefreshWorkflowDraft(true)
          callback?.onError?.()
          return null
        }

        setSyncWorkflowDraftHash(res.hash, baseParams.params.graph)
        setDraftUpdatedAt(res.updated_at)
        callback?.onSuccess?.()
        return { hash: res.hash, updatedAt: res.updated_at }
      } finally {
        callback?.onSettled?.()
      }
    },
    [
      workflowStore,
      getNodesReadOnly,
      getPostParams,
      handleRefreshWorkflowDraft,
      isCollaborationEnabled,
      confirmEmptyDraft,
      appId,
    ],
  )

  const doSyncWorkflowDraftLocally = useSerialAsyncCallback(performLocalSync, getNodesReadOnly)
  const doSyncWorkflowDraft = useCallback(
    async (
      notRefreshWhenSyncError?: boolean,
      callback?: SyncDraftCallback,
      options?: SyncDraftOptions,
    ): Promise<SyncDraftResult | null> => {
      if (getNodesReadOnly()) return null
      const { isWorkflowDataLoaded } = workflowStore.getState()
      if (shouldSkipDraftSync(appId, isWorkflowDataLoaded)) {
        callback?.onSettled?.()
        return null
      }

      // Capture before ReactFlow resets its store during route unmount.
      const baseParams = getPostParams()
      if (!baseParams) {
        callback?.onSettled?.()
        return null
      }

      const emptyGraph = isEmptyGraph(baseParams.params.graph)
      if (
        emptyGraph &&
        (!isMountedRef.current ||
          emptySyncPendingRef.current ||
          workflowStore.getState().showConfirm)
      ) {
        callback?.onSettled?.()
        return null
      }

      if (isCollaborationEnabled && workflowStore.getState().loroDraft) {
        if (emptyGraph) emptySyncPendingRef.current = true
        try {
          if (emptyGraph && workflowStore.getState().syncWorkflowDraftIsEmpty !== true) {
            if (!(await confirmEmptyDraft())) return null
            const current = getPostParams()
            if (getNodesReadOnly() || !current || !isEmptyGraph(current.params.graph)) return null
          }
          if (options?.environmentVariablePatch)
            await updateEnvironmentVariables(
              { appId: baseParams.appId, ...options.environmentVariablePatch },
              { silent: true },
            )
          const result = await collaborationManager.syncLoroDraft(baseParams.appId, emptyGraph)
          callback?.onSuccess?.()
          return result
        } catch {
          callback?.onError?.()
          return null
        } finally {
          if (emptyGraph) emptySyncPendingRef.current = false
          callback?.onSettled?.()
        }
      }

      const shouldRequestLeader =
        !emptyGraph &&
        isCollaborationEnabled &&
        collaborationManager.isConnected() &&
        !collaborationManager.getIsLeader() &&
        !options?.forceLocal

      if (!shouldRequestLeader) {
        // Empty graphs stay in this tab so any required consent belongs to the initiating user.
        if (emptyGraph) emptySyncPendingRef.current = true
        try {
          return await doSyncWorkflowDraftLocally(
            baseParams,
            notRefreshWhenSyncError,
            callback,
            options,
          )
        } finally {
          if (emptyGraph) emptySyncPendingRef.current = false
        }
      }

      try {
        const result = await collaborationManager.requestWorkflowSync()
        const { setSyncWorkflowDraftHash, setDraftUpdatedAt } = workflowStore.getState()
        setSyncWorkflowDraftHash(result.hash, baseParams.params.graph)
        setDraftUpdatedAt(result.updatedAt)
        callback?.onSuccess?.()
        return result
      } catch {
        const { isWorkflowDataLoaded } = workflowStore.getState()
        if (!shouldSkipDraftSync(appId, isWorkflowDataLoaded)) callback?.onError?.()
        return null
      } finally {
        callback?.onSettled?.()
      }
    },
    [
      doSyncWorkflowDraftLocally,
      getNodesReadOnly,
      getPostParams,
      isCollaborationEnabled,
      workflowStore,
      appId,
      confirmEmptyDraft,
    ],
  )

  return {
    doSyncWorkflowDraft,
    syncWorkflowDraftWhenPageClose,
  }
}

export const useNodesSyncDraftByCanEdit = (canEdit: boolean) => {
  const { getNodesReadOnly } = useNodesReadOnlyByCanEdit(canEdit)

  return useNodesSyncDraftBase(getNodesReadOnly)
}

export const useNodesSyncDraft = () => {
  const { getNodesReadOnly } = useNodesReadOnly()

  return useNodesSyncDraftBase(getNodesReadOnly)
}