dify/web/features/agent-v2/agent-composer/store-modules/tools.ts

import type {
  AgentCliTool,
  AgentProviderTool,
  AgentSoulConfigFormState,
  AgentTool,
} from '../form-state'
import type { DraftFieldUpdate } from './utils'
import type { ToolDefaultValue } from '@/app/components/workflow/block-selector/types'
import { atom } from 'jotai'
import { syncCliToolReferenceLabels } from '../reference-labels'
import { agentComposerDraftAtom } from '../store'
import { resolveDraftFieldUpdate } from './utils'

export type AgentProviderToolDefaultValue = ToolDefaultValue & {
  allowDelete?: boolean
  credentialType?: AgentProviderTool['credentialType']
  credentialRequired?: boolean
}

export const agentComposerToolsAtom = atom(
  (get) => get(agentComposerDraftAtom).tools,
  (get, set, toolsUpdate: DraftFieldUpdate<AgentTool[]>) => {
    const draft = get(agentComposerDraftAtom)
    const tools = resolveDraftFieldUpdate(draft.tools, toolsUpdate)

    set(agentComposerDraftAtom, {
      ...draft,
      prompt: syncCliToolReferenceLabels({
        prompt: draft.prompt,
        currentTools: draft.tools,
        nextTools: tools,
      }),
      tools,
    })
  },
)

const toProviderToolAction = (tool: AgentProviderToolDefaultValue) => ({
  id: `${tool.provider_id}:${tool.tool_name}`,
  name: tool.tool_label || tool.title || tool.tool_name,
  toolName: tool.tool_name,
  description: tool.tool_description || '',
})

const getCredentialVariant = (tool: AgentProviderToolDefaultValue) => {
  if (!tool.credentialRequired) return 'none' as const

  if (!tool.allowDelete)
    return tool.credential_id ? ('authorized' as const) : ('unauthorized' as const)

  return tool.is_team_authorization ? ('authorized' as const) : ('unauthorized' as const)
}

const getCredentialType = (tool: AgentProviderToolDefaultValue) => {
  if (!tool.credentialRequired) return undefined

  if (tool.credentialType === 'oauth2') return 'oauth2' as const

  if (!tool.allowDelete)
    return tool.credential_id ? ('api-key' as const) : ('unauthorized' as const)

  return tool.is_team_authorization ? ('api-key' as const) : ('unauthorized' as const)
}

export const addProviderTools = (
  currentTools: AgentTool[],
  selectedTools: AgentProviderToolDefaultValue[],
): AgentTool[] => {
  if (selectedTools.length === 0) return currentTools

  const nextTools = [...currentTools]

  selectedTools.forEach((selectedTool) => {
    const action = toProviderToolAction(selectedTool)
    const existingToolIndex = nextTools.findIndex(
      (tool) => tool.kind === 'provider' && tool.id === selectedTool.provider_id,
    )
    const existingTool = nextTools[existingToolIndex]

    if (existingTool?.kind === 'provider') {
      if (
        existingTool.actions.some((existingAction) => existingAction.toolName === action.toolName)
      )
        return

      nextTools[existingToolIndex] = {
        ...existingTool,
        displayName: existingTool.displayName ?? selectedTool.provider_show_name,
        icon: existingTool.icon ?? selectedTool.provider_icon,
        iconDark: existingTool.iconDark ?? selectedTool.provider_icon_dark,
        allowDelete: existingTool.allowDelete ?? selectedTool.allowDelete,
        actions: [...existingTool.actions, action],
      }
      return
    }

    nextTools.push({
      id: selectedTool.provider_id,
      name: selectedTool.provider_name,
      kind: 'provider',
      displayName: selectedTool.provider_show_name,
      iconClassName: 'i-custom-public-other-default-tool-icon text-text-tertiary',
      icon: selectedTool.provider_icon,
      iconDark: selectedTool.provider_icon_dark,
      providerType: selectedTool.provider_type,
      allowDelete: selectedTool.allowDelete,
      credentialId: selectedTool.credential_id,
      credentialKey: selectedTool.is_team_authorization
        ? 'agentDetail.configure.tools.credential.authOne'
        : undefined,
      credentialType: getCredentialType(selectedTool),
      credentialVariant: getCredentialVariant(selectedTool),
      actions: [action],
    })
  })

  return nextTools
}

export const addProviderToolsAtom = atom(
  null,
  (_get, set, selectedTools: AgentProviderToolDefaultValue[]) => {
    set(agentComposerToolsAtom, (tools) => addProviderTools(tools, selectedTools))
  },
)

export const saveCliToolAtom = atom(null, (_get, set, cliTool: AgentCliTool) => {
  set(agentComposerToolsAtom, (tools) =>
    tools.some((tool) => tool.kind === 'cli' && tool.id === cliTool.id)
      ? tools.map((tool) => (tool.id === cliTool.id ? cliTool : tool))
      : [...tools, cliTool],
  )
})

export const removeCliToolAtom = atom(null, (_get, set, toolId: string) => {
  set(agentComposerToolsAtom, (tools) => tools.filter((tool) => tool.id !== toolId))
})

export const agentComposerToolSettingsAtom = atom(
  (get) => get(agentComposerDraftAtom).toolSettings,
  (get, set, toolSettingsUpdate: DraftFieldUpdate<Record<string, Record<string, unknown>>>) => {
    const draft = get(agentComposerDraftAtom)

    set(agentComposerDraftAtom, {
      ...draft,
      toolSettings: resolveDraftFieldUpdate(draft.toolSettings, toolSettingsUpdate),
    })
  },
)

const omitToolSettings = (
  toolSettings: AgentSoulConfigFormState['toolSettings'],
  actionIds: string[],
) => {
  const nextToolSettings = { ...toolSettings }

  actionIds.forEach((actionId) => {
    delete nextToolSettings[actionId]
  })

  return nextToolSettings
}

export const removeProviderToolAtom = atom(null, (get, set, toolId: string) => {
  const draft = get(agentComposerDraftAtom)
  const toolToRemove = draft.tools.find((tool) => tool.kind === 'provider' && tool.id === toolId)
  const actionIds =
    toolToRemove?.kind === 'provider' ? toolToRemove.actions.map((action) => action.id) : []

  set(agentComposerDraftAtom, {
    ...draft,
    tools: draft.tools.filter((tool) => tool.id !== toolId),
    toolSettings: omitToolSettings(draft.toolSettings, actionIds),
  })
})

export const removeProviderToolActionAtom = atom(
  null,
  (
    get,
    set,
    {
      toolId,
      actionId,
    }: {
      toolId: string
      actionId: string
    },
  ) => {
    const draft = get(agentComposerDraftAtom)

    set(agentComposerDraftAtom, {
      ...draft,
      tools: draft.tools.flatMap((tool) => {
        if (tool.kind !== 'provider' || tool.id !== toolId) return [tool]

        const nextActions = tool.actions.filter((action) => action.id !== actionId)
        return nextActions.length > 0 ? [{ ...tool, actions: nextActions }] : []
      }),
      toolSettings: omitToolSettings(draft.toolSettings, [actionId]),
    })
  },
)

export const setProviderToolCredentialAtom = atom(
  null,
  (
    _get,
    set,
    {
      toolId,
      credentialId,
      credentialType,
    }: {
      toolId: string
      credentialId?: string
      credentialType?: AgentProviderTool['credentialType']
    },
  ) => {
    set(agentComposerToolsAtom, (tools) =>
      tools.map((tool) => {
        if (tool.kind !== 'provider' || tool.id !== toolId) return tool

        const nextCredentialType =
          credentialType === 'oauth2' || tool.credentialType === 'oauth2' ? 'oauth2' : 'api-key'

        return {
          ...tool,
          credentialId,
          credentialType: nextCredentialType,
          credentialVariant: 'authorized',
        }
      }),
    )
  },
)

export const saveProviderToolActionSettingsAtom = atom(
  null,
  (
    _get,
    set,
    {
      actionId,
      value,
    }: {
      actionId: string
      value: Record<string, unknown>
    },
  ) => {
    set(agentComposerToolSettingsAtom, (toolSettings) => ({
      ...toolSettings,
      [actionId]: value,
    }))
  },
)