dify/web/app/components/workflow/nodes/agent-v2/agent-soul-config.ts

import type {
  AgentSoulConfig,
  WorkflowAgentComposerResponse,
} from '@dify/contracts/api/console/apps/types.gen'
import type { DefaultModelResponse } from '@/app/components/header/account-setting/model-provider-page/declarations'
import { useMutation, useQueryClient } from '@tanstack/react-query'
import { debounce } from 'es-toolkit/compat'
import isEqual from 'fast-deep-equal'
import { useStore as useJotaiStore, useSetAtom } from 'jotai'
import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import { useHooksStore } from '@/app/components/workflow/hooks-store'
import { useSerialAsyncCallback } from '@/app/components/workflow/hooks/use-serial-async-callback'
import {
  agentSoulConfigToFormState,
  formStateToAgentSoulConfig,
} from '@/features/agent-v2/agent-composer/conversions'
import {
  agentComposerDraftAtom,
  agentComposerOriginalConfigAtom,
  agentComposerOriginalDraftAtom,
  isAgentComposerDirtyAtom,
} from '@/features/agent-v2/agent-composer/store'
import { consoleQuery } from '@/service/client'
import { FlowType } from '@/types/common'

const DRAFT_AUTOSAVE_WAIT = 5000

function getModelProviderPluginId(provider: string) {
  const [organization, pluginName] = provider.split('/').filter(Boolean)

  if (organization && pluginName) return `${organization}/${pluginName}`

  return provider ? `langgenius/${provider}` : ''
}

export function getDefaultAgentSoul(defaultModel?: DefaultModelResponse): AgentSoulConfig {
  const baseConfig: AgentSoulConfig = {
    schema_version: 1,
    prompt: {
      system_prompt: '',
    },
  }

  if (!defaultModel) return baseConfig

  const modelProvider = defaultModel.provider.provider

  return {
    ...baseConfig,
    model: {
      model_provider: modelProvider,
      model: defaultModel.model,
      plugin_id: getModelProviderPluginId(modelProvider),
    },
  }
}

export function useWorkflowInlineAgentConfigureSync({
  nodeId,
  baseConfig,
  currentModel,
  autoSaveEnabled = true,
  onDraftSaved,
  enabled,
}: {
  nodeId: string
  baseConfig?: AgentSoulConfig
  currentModel?: {
    provider: string
    model: string
    plugin_id?: string
  }
  autoSaveEnabled?: boolean
  onDraftSaved?: (composerState: WorkflowAgentComposerResponse) => void
  enabled: boolean
}) {
  const queryClient = useQueryClient()
  const configsMap = useHooksStore((state) => state.configsMap)
  const store = useJotaiStore()
  const setOriginalConfig = useSetAtom(agentComposerOriginalConfigAtom)
  const setOriginalDraft = useSetAtom(agentComposerOriginalDraftAtom)
  const [draftSavedAt, setDraftSavedAt] = useState<number | undefined>(undefined)
  const baseConfigRef = useRef(baseConfig)
  const currentModelRef = useRef(currentModel)
  const enabledRef = useRef(enabled)
  const onDraftSavedRef = useRef(onDraftSaved)
  const lastAutosavedDraftKeyRef = useRef<string | undefined>(undefined)
  const { mutateAsync: saveAppComposer } = useMutation(
    consoleQuery.apps.byAppId.workflows.draft.nodes.byNodeId.agentComposer.put.mutationOptions(),
  )
  const { mutateAsync: saveSnippetComposer } = useMutation(
    consoleQuery.snippets.bySnippetId.workflows.draft.nodes.byNodeId.agentComposer.put.mutationOptions(),
  )

  baseConfigRef.current = baseConfig
  currentModelRef.current = currentModel
  enabledRef.current = enabled
  onDraftSavedRef.current = onDraftSaved

  const getAgentSoulDraft = useCallback(
    () =>
      formStateToAgentSoulConfig({
        baseConfig: baseConfigRef.current,
        formState: store.get(agentComposerDraftAtom),
        currentModel: currentModelRef.current,
      }),
    [store],
  )

  const saveComposer = useSerialAsyncCallback(
    async (configSnapshot: AgentSoulConfig): Promise<WorkflowAgentComposerResponse | undefined> => {
      if (
        !configsMap?.flowId ||
        (configsMap.flowType !== FlowType.appFlow && configsMap.flowType !== FlowType.snippet)
      )
        return

      const savedDraftKey = JSON.stringify(configSnapshot)
      const body = {
        variant: 'workflow' as const,
        save_strategy: 'node_job_only' as const,
        agent_soul: configSnapshot,
      }
      const composerState =
        configsMap.flowType === FlowType.snippet
          ? await saveSnippetComposer({
              params: {
                snippet_id: configsMap.flowId,
                node_id: nodeId,
              },
              body,
            })
          : await saveAppComposer({
              params: {
                app_id: configsMap.flowId,
                node_id: nodeId,
              },
              body,
            })

      if (configsMap.flowType === FlowType.snippet) {
        queryClient.setQueryData(
          consoleQuery.snippets.bySnippetId.workflows.draft.nodes.byNodeId.agentComposer.get.queryKey(
            {
              input: {
                params: {
                  snippet_id: configsMap.flowId,
                  node_id: nodeId,
                },
              },
            },
          ),
          composerState,
        )
      } else {
        queryClient.setQueryData(
          consoleQuery.apps.byAppId.workflows.draft.nodes.byNodeId.agentComposer.get.queryKey({
            input: {
              params: {
                app_id: configsMap.flowId,
                node_id: nodeId,
              },
            },
          }),
          composerState,
        )
      }
      setOriginalConfig(composerState.agent_soul)
      setOriginalDraft(agentSoulConfigToFormState(composerState.agent_soul))
      setDraftSavedAt(Date.now())
      lastAutosavedDraftKeyRef.current = savedDraftKey
      onDraftSavedRef.current?.(composerState)
      return composerState
    },
  )

  const latestDraftSaveRef = useRef<() => void>(() => undefined)
  latestDraftSaveRef.current = () => {
    void saveComposer(getAgentSoulDraft())
  }

  const debouncedSaveDraft = useMemo(
    () =>
      debounce(() => {
        latestDraftSaveRef.current()
      }, DRAFT_AUTOSAVE_WAIT),
    [],
  )

  const saveDraft = useCallback(async () => {
    if (!enabledRef.current) return

    const configSnapshot = getAgentSoulDraft()
    const hasEffectiveModelChange = !isEqual(configSnapshot.model, baseConfigRef.current?.model)
    debouncedSaveDraft.cancel?.()
    if (!store.get(isAgentComposerDirtyAtom) && !hasEffectiveModelChange) return

    return saveComposer(configSnapshot)
  }, [debouncedSaveDraft, getAgentSoulDraft, saveComposer, store])
  const saveAgentSoulConfig = useCallback(
    async (agentSoulConfig: AgentSoulConfig) => {
      debouncedSaveDraft.cancel?.()
      return saveComposer(agentSoulConfig)
    },
    [debouncedSaveDraft, saveComposer],
  )

  useEffect(() => {
    return store.sub(agentComposerDraftAtom, () => {
      const agentSoulDraft = getAgentSoulDraft()
      const agentSoulDraftKey = JSON.stringify(agentSoulDraft)

      if (
        !enabledRef.current ||
        !autoSaveEnabled ||
        !store.get(isAgentComposerDirtyAtom) ||
        lastAutosavedDraftKeyRef.current === agentSoulDraftKey
      ) {
        return
      }

      debouncedSaveDraft()
    })
  }, [autoSaveEnabled, debouncedSaveDraft, getAgentSoulDraft, store])

  useEffect(() => {
    return () => {
      if (autoSaveEnabled) debouncedSaveDraft.flush?.()
    }
  }, [autoSaveEnabled, debouncedSaveDraft])

  return {
    draftSavedAt,
    saveAgentSoulConfig,
    saveDraft,
  }
}