dify/web/app/components/header/account-setting/model-provider-page/model-auth/hooks/use-auth.ts

import type {
  ConfigurationMethodEnum,
  Credential,
  CustomConfigurationModelFixedFields,
  CustomModel,
  ModelModalModeEnum,
  ModelProvider,
} from '../../declarations'
import { toast } from '@langgenius/dify-ui/toast'
import { useCallback, useRef, useState } from 'react'
import { useTranslation } from 'react-i18next'
import {
  useModelModalHandler,
  useRefreshModel,
} from '@/app/components/header/account-setting/model-provider-page/hooks'
import { useDeleteModel } from '@/service/use-models'
import { useAuthService } from './use-auth-service'

type ProviderCredentialOperationPayload = {
  credential_id: string
}

type ModelCredentialOperationPayload = {
  credential_id: string
  model: string
  model_type: CustomModel['model_type']
}

export const useAuth = (
  provider: ModelProvider,
  configurationMethod: ConfigurationMethodEnum,
  currentCustomConfigurationModelFixedFields?: CustomConfigurationModelFixedFields,
  extra: {
    isModelCredential?: boolean
    onUpdate?: (newPayload?: any, formValues?: Record<string, any>) => void
    onRemove?: (credentialId: string) => void
    mode?: ModelModalModeEnum
  } = {},
) => {
  const { isModelCredential, onUpdate, onRemove, mode } = extra
  const { t } = useTranslation()
  const {
    getDeleteCredentialService,
    getActiveCredentialService,
    getEditCredentialService,
    getAddCredentialService,
  } = useAuthService(provider.provider)
  const { mutateAsync: deleteModelService } = useDeleteModel(provider.provider)
  const handleOpenModelModal = useModelModalHandler()
  const { handleRefreshModel } = useRefreshModel()
  const pendingOperationCredentialId = useRef<string | null>(null)
  const [deleteCredentialId, setDeleteCredentialId] = useState<string | null>(null)
  const handleSetDeleteCredentialId = useCallback((credentialId: string | null) => {
    setDeleteCredentialId(credentialId)
    pendingOperationCredentialId.current = credentialId
  }, [])
  const pendingOperationModel = useRef<CustomModel | null>(null)
  const [deleteModel, setDeleteModel] = useState<CustomModel | null>(null)
  const handleSetDeleteModel = useCallback((model: CustomModel | null) => {
    setDeleteModel(model)
    pendingOperationModel.current = model
  }, [])
  const resolveModelContext = useCallback(
    (model?: CustomModel | null): CustomModel | undefined => {
      if (model) return model

      if (!currentCustomConfigurationModelFixedFields) return undefined

      return {
        model: currentCustomConfigurationModelFixedFields.__model_name,
        model_type: currentCustomConfigurationModelFixedFields.__model_type,
      }
    },
    [currentCustomConfigurationModelFixedFields],
  )
  const openConfirmDelete = useCallback((credential?: Credential, model?: CustomModel) => {
    if (credential) handleSetDeleteCredentialId(credential.credential_id)
    if (model) handleSetDeleteModel(model)
  }, [])
  const closeConfirmDelete = useCallback(() => {
    handleSetDeleteCredentialId(null)
    handleSetDeleteModel(null)
  }, [])
  const [doingAction, setDoingAction] = useState(false)
  const doingActionRef = useRef(doingAction)
  const handleSetDoingAction = useCallback((doing: boolean) => {
    doingActionRef.current = doing
    setDoingAction(doing)
  }, [])
  const handleActiveCredential = useCallback(
    async (credential: Credential, model?: CustomModel) => {
      if (doingActionRef.current) return
      try {
        handleSetDoingAction(true)
        const modelContext = model ?? (isModelCredential ? resolveModelContext() : undefined)
        if (modelContext) {
          const activeModelCredential = getActiveCredentialService(true) as (
            payload: ModelCredentialOperationPayload,
          ) => Promise<unknown>
          await activeModelCredential({
            credential_id: credential.credential_id,
            model: modelContext.model,
            model_type: modelContext.model_type,
          })
        } else {
          const activeProviderCredential = getActiveCredentialService(false) as (
            payload: ProviderCredentialOperationPayload,
          ) => Promise<unknown>
          await activeProviderCredential({
            credential_id: credential.credential_id,
          })
        }
        toast.success(t(($) => $['api.actionSuccess'], { ns: 'common' }))
        handleRefreshModel(provider, undefined, true)
      } finally {
        handleSetDoingAction(false)
      }
    },
    [getActiveCredentialService, isModelCredential, resolveModelContext, t, handleSetDoingAction],
  )
  const handleConfirmDelete = useCallback(async () => {
    if (doingActionRef.current) return
    if (!pendingOperationCredentialId.current && !pendingOperationModel.current) {
      closeConfirmDelete()
      return
    }
    try {
      handleSetDoingAction(true)
      let payload: any = {}
      if (pendingOperationCredentialId.current) {
        if (isModelCredential) {
          const modelContext = resolveModelContext(pendingOperationModel.current)
          if (!modelContext) return

          payload = {
            credential_id: pendingOperationCredentialId.current,
            model: modelContext.model,
            model_type: modelContext.model_type,
          }
          const deleteModelCredential = getDeleteCredentialService(true) as (
            payload: ModelCredentialOperationPayload,
          ) => Promise<unknown>
          await deleteModelCredential(payload)
        } else {
          payload = {
            credential_id: pendingOperationCredentialId.current,
          }
          const deleteProviderCredential = getDeleteCredentialService(false) as (
            payload: ProviderCredentialOperationPayload,
          ) => Promise<unknown>
          await deleteProviderCredential(payload)
        }
      }
      if (!pendingOperationCredentialId.current && pendingOperationModel.current) {
        payload = {
          model: pendingOperationModel.current.model,
          model_type: pendingOperationModel.current.model_type,
        }
        await deleteModelService(payload)
      }
      toast.success(t(($) => $['api.actionSuccess'], { ns: 'common' }))
      handleRefreshModel(provider, undefined, true)
      onRemove?.(pendingOperationCredentialId.current ?? '')
      closeConfirmDelete()
    } finally {
      handleSetDoingAction(false)
    }
  }, [
    t,
    handleSetDoingAction,
    getDeleteCredentialService,
    isModelCredential,
    closeConfirmDelete,
    handleRefreshModel,
    provider,
    configurationMethod,
    deleteModelService,
    resolveModelContext,
  ])
  const handleSaveCredential = useCallback(
    async (payload: Record<string, any>) => {
      if (doingActionRef.current) return
      try {
        handleSetDoingAction(true)
        let res: {
          result?: string
        } = {}
        if (payload.credential_id)
          res = await getEditCredentialService(!!isModelCredential)(payload as any)
        else res = await getAddCredentialService(!!isModelCredential)(payload as any)
        if (res.result === 'success') {
          toast.success(t(($) => $['actionMsg.modifiedSuccessfully'], { ns: 'common' }))
          handleRefreshModel(provider, undefined, !payload.credential_id)
        }
      } finally {
        handleSetDoingAction(false)
      }
    },
    [t, handleSetDoingAction, getEditCredentialService, getAddCredentialService],
  )
  const handleOpenModal = useCallback(
    (credential?: Credential, model?: CustomModel) => {
      handleOpenModelModal(
        provider,
        configurationMethod,
        currentCustomConfigurationModelFixedFields,
        {
          isModelCredential,
          credential,
          model,
          onUpdate,
          mode,
        },
      )
    },
    [
      handleOpenModelModal,
      provider,
      configurationMethod,
      currentCustomConfigurationModelFixedFields,
      isModelCredential,
      onUpdate,
      mode,
    ],
  )
  return {
    pendingOperationCredentialId,
    pendingOperationModel,
    openConfirmDelete,
    closeConfirmDelete,
    doingAction,
    handleActiveCredential,
    handleConfirmDelete,
    deleteCredentialId,
    deleteModel,
    handleSaveCredential,
    handleOpenModal,
  }
}