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

import type { AgentInlineBinding } from '../../block-selector/types'
import { toast } from '@langgenius/dify-ui/toast'
import { skipToken, useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { useCallback } from 'react'
import { useTranslation } from 'react-i18next'
import { ModelTypeEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
import { useDefaultModel } from '@/app/components/header/account-setting/model-provider-page/hooks'
import { useHooksStore } from '@/app/components/workflow/hooks-store'
import { consoleQuery } from '@/service/client'
import { FlowType } from '@/types/common'
import { getDefaultAgentSoul } from './agent-soul-config'

type CreatedInlineAgentBinding = AgentInlineBinding & {
  agent_id: string
  current_snapshot_id: string
}

type CreateInlineAgentBindingOptions = {
  onError?: () => void
  onSuccess?: (binding: CreatedInlineAgentBinding) => void
}

export function useAgentRosterDetail(agentId?: string) {
  return useQuery(
    consoleQuery.agent.byAgentId.get.queryOptions({
      input: agentId
        ? {
            params: {
              agent_id: agentId,
            },
          }
        : skipToken,
    }),
  )
}

export function useWorkflowInlineAgentDetail(nodeId?: string, agentId?: string | null) {
  const configsMap = useHooksStore((state) => state.configsMap)

  return useQuery(
    consoleQuery.apps.byAppId.workflows.draft.nodes.byNodeId.agentComposer.get.queryOptions({
      input:
        configsMap?.flowId && configsMap.flowType === FlowType.appFlow && nodeId && agentId
          ? {
              params: {
                app_id: configsMap.flowId,
                node_id: nodeId,
              },
            }
          : skipToken,
    }),
  )
}

export function useCreateInlineAgentBinding() {
  const { t } = useTranslation('agentV2')
  const configsMap = useHooksStore((state) => state.configsMap)
  const { data: defaultModel } = useDefaultModel(ModelTypeEnum.textGeneration)
  const queryClient = useQueryClient()
  const { isPending, mutateAsync } = useMutation(
    consoleQuery.apps.byAppId.workflows.draft.nodes.byNodeId.agentComposer.put.mutationOptions(),
  )

  const createInlineAgentBinding = useCallback(
    async (nodeId: string, options?: CreateInlineAgentBindingOptions) => {
      if (!configsMap?.flowId || configsMap.flowType !== FlowType.appFlow) {
        toast.error(t(($) => $['roster.nodeSelector.createInlineFailed']))
        options?.onError?.()
        return
      }

      try {
        const composerState = await mutateAsync({
          params: {
            app_id: configsMap.flowId,
            node_id: nodeId,
          },
          body: {
            variant: 'workflow',
            save_strategy: 'node_job_only',
            binding: {
              binding_type: 'inline_agent',
            },
            soul_lock: {
              locked: false,
            },
            agent_soul: getDefaultAgentSoul(defaultModel),
          },
        })
        const binding = composerState.binding

        if (
          binding?.binding_type !== 'inline_agent' ||
          !binding.agent_id ||
          !binding.current_snapshot_id
        ) {
          toast.error(t(($) => $['roster.nodeSelector.createInlineFailed']))
          options?.onError?.()
          return
        }

        queryClient.setQueryData(
          consoleQuery.apps.byAppId.workflows.draft.nodes.byNodeId.agentComposer.get.queryKey({
            input: {
              params: {
                app_id: configsMap.flowId,
                node_id: nodeId,
              },
            },
          }),
          composerState,
        )
        options?.onSuccess?.({
          binding_type: 'inline_agent',
          agent_id: binding.agent_id,
          current_snapshot_id: binding.current_snapshot_id,
        })
      } catch {
        options?.onError?.()
      }
    },
    [configsMap?.flowId, configsMap?.flowType, defaultModel, mutateAsync, queryClient, t],
  )

  return {
    createInlineAgentBinding,
    isCreatingInlineAgent: isPending,
  }
}