dify/web/app/components/sub-graph/components/sub-graph-main.tsx

import type { FC } from 'react'
import type { Viewport } from 'reactflow'
import type { SyncWorkflowDraft, SyncWorkflowDraftCallback } from '../types'
import type { Shape as HooksStoreShape } from '@/app/components/workflow/hooks-store'
import type { NestedNodeConfig } from '@/app/components/workflow/nodes/_base/types'
import type { Edge, Node } from '@/app/components/workflow/types'
import { useCallback, useEffect, useMemo } from 'react'
import { useStoreApi } from 'reactflow'
import { InteractionMode, WorkflowWithInnerContext } from '@/app/components/workflow'
import { useNodesInteractions } from '@/app/components/workflow/hooks'
import { useSetWorkflowVarsWithValue } from '@/app/components/workflow/hooks/use-fetch-workflow-inspect-vars'
import { useInspectVarsCrudCommon } from '@/app/components/workflow/hooks/use-inspect-vars-crud-common'
import { useWorkflowStore } from '@/app/components/workflow/store'
import { BlockEnum } from '@/app/components/workflow/types'
import { FlowType } from '@/types/common'
import { useAvailableNodesMetaData } from '../hooks'
import SubGraphChildren from './sub-graph-children'

type SubGraphMainBaseProps = {
  nodes: Node[]
  edges: Edge[]
  viewport: Viewport
  title: string
  extractorNodeId: string
  configsMap?: HooksStoreShape['configsMap']
  selectableNodeTypes?: BlockEnum[]
  onSave?: (nodes: Node[], edges: Edge[]) => void
  onSyncWorkflowDraft?: SyncWorkflowDraft
  isOpen: boolean
  pendingSingleRun?: boolean
  onPendingSingleRunHandled?: () => void
}

type SubGraphMainProps
  = | (SubGraphMainBaseProps & {
    variant: 'agent'
    nestedNodeConfig: NestedNodeConfig
    onNestedNodeConfigChange: (config: NestedNodeConfig) => void
  })
  | (SubGraphMainBaseProps & {
    variant: 'assemble'
    nestedNodeConfig: NestedNodeConfig
    onNestedNodeConfigChange: (config: NestedNodeConfig) => void
  })

const SubGraphMain: FC<SubGraphMainProps> = (props) => {
  const {
    nodes,
    edges,
    viewport,
    variant,
    title,
    extractorNodeId,
    configsMap,
    selectableNodeTypes,
    onSave,
    onSyncWorkflowDraft,
    isOpen,
    pendingSingleRun,
    onPendingSingleRunHandled,
  } = props
  const reactFlowStore = useStoreApi()
  const workflowStore = useWorkflowStore()
  const { handleNodeSelect } = useNodesInteractions()
  const availableNodesMetaData = useAvailableNodesMetaData()
  const flowType = configsMap?.flowType ?? FlowType.appFlow
  const flowId = configsMap?.flowId ?? ''
  const { fetchInspectVars } = useSetWorkflowVarsWithValue({
    flowType,
    flowId,
    interactionMode: InteractionMode.Subgraph,
  })
  const inspectVarsCrud = useInspectVarsCrudCommon({
    flowType,
    flowId,
    interactionMode: InteractionMode.Subgraph,
  })

  const handleSyncSubGraphDraft = useCallback(async () => {
    const { getNodes, edges } = reactFlowStore.getState()
    await onSave?.(getNodes() as Node[], edges as Edge[])
  }, [onSave, reactFlowStore])

  const handleSyncWorkflowDraft = useCallback(async (
    notRefreshWhenSyncError?: boolean,
    callback?: SyncWorkflowDraftCallback,
  ) => {
    try {
      await handleSyncSubGraphDraft()
      if (onSyncWorkflowDraft) {
        await onSyncWorkflowDraft(notRefreshWhenSyncError, callback)
        return
      }
      callback?.onSuccess?.()
    }
    catch {
      callback?.onError?.()
    }
    finally {
      callback?.onSettled?.()
    }
  }, [handleSyncSubGraphDraft, onSyncWorkflowDraft])

  useEffect(() => {
    if (!isOpen || !pendingSingleRun)
      return

    const { getNodes } = reactFlowStore.getState()
    const currentNodes = getNodes()
    const hasExtractorNode = currentNodes.some(node => node.id === extractorNodeId)
    if (!hasExtractorNode)
      return

    // NodePanel listens for pendingSingleRun only when the extractor node is selected in this subgraph.
    handleNodeSelect(extractorNodeId, false, true)

    // Defer run until the selection is applied and the panel is ready.
    const frame = requestAnimationFrame(() => {
      const store = workflowStore.getState()
      store.setPendingSingleRun({
        nodeId: extractorNodeId,
        action: 'run',
      })
      onPendingSingleRunHandled?.()
    })

    return () => cancelAnimationFrame(frame)
  }, [extractorNodeId, handleNodeSelect, isOpen, onPendingSingleRunHandled, pendingSingleRun, reactFlowStore, workflowStore])

  const resolvedSelectableTypes = useMemo(() => {
    if (selectableNodeTypes && selectableNodeTypes.length > 0)
      return selectableNodeTypes
    return variant === 'agent' ? [BlockEnum.LLM] : [BlockEnum.Code]
  }, [selectableNodeTypes, variant])

  const hooksStore = useMemo(() => ({
    interactionMode: InteractionMode.Subgraph,
    subGraphSelectableNodeTypes: resolvedSelectableTypes,
    availableNodesMetaData,
    configsMap,
    fetchInspectVars,
    ...inspectVarsCrud,
    doSyncWorkflowDraft: handleSyncWorkflowDraft,
    syncWorkflowDraftWhenPageClose: handleSyncSubGraphDraft,
  }), [availableNodesMetaData, configsMap, fetchInspectVars, handleSyncSubGraphDraft, handleSyncWorkflowDraft, inspectVarsCrud, resolvedSelectableTypes])

  const subGraphChildren = variant === 'agent'
    ? (
        <SubGraphChildren
          variant="agent"
          title={title}
          extractorNodeId={extractorNodeId}
          nestedNodeConfig={props.nestedNodeConfig}
          onNestedNodeConfigChange={props.onNestedNodeConfigChange}
        />
      )
    : (
        <SubGraphChildren
          variant="assemble"
          title={title}
          extractorNodeId={extractorNodeId}
          nestedNodeConfig={props.nestedNodeConfig}
          onNestedNodeConfigChange={props.onNestedNodeConfigChange}
        />
      )

  return (
    <WorkflowWithInnerContext
      nodes={nodes}
      edges={edges}
      viewport={viewport}
      hooksStore={hooksStore}
      allowSelectionWhenReadOnly
      canvasReadOnly
      interactionMode={InteractionMode.Subgraph}
    >
      {subGraphChildren}
    </WorkflowWithInnerContext>
  )
}

export default SubGraphMain