From 7bddf323e138781abf560081d48ffedb1b088a42 Mon Sep 17 00:00:00 2001
From: zxhlyh <jasonapring2015@outlook.com>
Date: Tue, 19 Aug 2025 16:42:07 +0800
Subject: [PATCH] fix: load balancing

---
 .../model-provider-page/declarations.ts       |  5 ++
 .../model-auth/authorized/credential-item.tsx | 12 +++-
 .../model-auth/config-model.tsx               | 26 +++++---
 .../switch-credential-in-load-balancing.tsx   | 59 ++++++++++---------
 .../provider-added-card/model-list-item.tsx   | 15 +++--
 .../model-load-balancing-configs.tsx          | 28 +++++----
 .../model-load-balancing-modal.tsx            | 43 +++++++++++---
 .../components/plugins/plugin-auth/types.ts   |  1 +
 web/i18n/en-US/common.ts                      |  6 +-
 web/i18n/zh-Hans/common.ts                    |  6 +-
 web/service/use-models.ts                     |  1 +
 11 files changed, 135 insertions(+), 67 deletions(-)

diff --git a/web/app/components/header/account-setting/model-provider-page/declarations.ts b/web/app/components/header/account-setting/model-provider-page/declarations.ts
index afb1cbce48..945806928a 100644
--- a/web/app/components/header/account-setting/model-provider-page/declarations.ts
+++ b/web/app/components/header/account-setting/model-provider-page/declarations.ts
@@ -86,6 +86,7 @@ export enum ModelStatusEnum {
   quotaExceeded = 'quota-exceeded',
   noPermission = 'no-permission',
   disabled = 'disabled',
+  credentialRemoved = 'credential-removed',
 }
 
 export const MODEL_STATUS_TEXT: { [k: string]: TypeWithI18N } = {
@@ -185,6 +186,8 @@ export type QuotaConfiguration = {
 export type Credential = {
   credential_id: string
   credential_name?: string
+  from_enterprise?: boolean
+  allowed_to_use?: boolean
 }
 
 export type CustomModel = {
@@ -316,4 +319,6 @@ export type ModelCredential = {
   credentials: Record<string, any>
   load_balancing: ModelLoadBalancingConfig
   available_credentials: Credential[]
+  current_credential_id?: string
+  current_credential_name?: string
 }
diff --git a/web/app/components/header/account-setting/model-provider-page/model-auth/authorized/credential-item.tsx b/web/app/components/header/account-setting/model-provider-page/model-auth/authorized/credential-item.tsx
index c7dc37031f..4329adba4c 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-auth/authorized/credential-item.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-auth/authorized/credential-item.tsx
@@ -13,6 +13,7 @@ import ActionButton from '@/app/components/base/action-button'
 import Tooltip from '@/app/components/base/tooltip'
 import cn from '@/utils/classnames'
 import type { Credential } from '../../declarations'
+import Badge from '@/app/components/base/badge'
 
 type CredentialItemProps = {
   credential: Credential
@@ -49,7 +50,9 @@ const CredentialItem = ({
       className={cn(
         'group flex h-8 items-center rounded-lg p-1 hover:bg-state-base-hover',
       )}
-      onClick={() => onItemClick?.(credential)}
+      onClick={() => {
+        onItemClick?.(credential)
+      }}
     >
       <div className='flex w-0 grow items-center space-x-1.5'>
         {
@@ -71,6 +74,13 @@ const CredentialItem = ({
           {credential.credential_name}
         </div>
       </div>
+      {
+        credential.from_enterprise && (
+          <Badge className='shrink-0'>
+            Enterprise
+          </Badge>
+        )
+      }
       {
         showAction && (
           <div className='ml-2 hidden shrink-0 items-center group-hover:flex'>
diff --git a/web/app/components/header/account-setting/model-provider-page/model-auth/config-model.tsx b/web/app/components/header/account-setting/model-provider-page/model-auth/config-model.tsx
index c9f1744cdb..083a3e238c 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-auth/config-model.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-auth/config-model.tsx
@@ -5,30 +5,32 @@ import {
 } from '@remixicon/react'
 import { useTranslation } from 'react-i18next'
 import Button from '@/app/components/base/button'
+import Indicator from '@/app/components/header/indicator'
 import cn from '@/utils/classnames'
 
 type ConfigModelProps = {
-  className?: string
   onClick?: () => void
   loadBalancingEnabled?: boolean
   loadBalancingInvalid?: boolean
+  credentialRemoved?: boolean
 }
 const ConfigModel = ({
-  className,
   onClick,
   loadBalancingEnabled,
   loadBalancingInvalid,
+  credentialRemoved,
 }: ConfigModelProps) => {
   const { t } = useTranslation()
 
-  if (loadBalancingEnabled && loadBalancingInvalid) {
+  if (loadBalancingEnabled && loadBalancingInvalid && !credentialRemoved) {
     return (
       <div
-        className='system-2xs-medium-uppercase flex h-[18px] items-center rounded-[5px] border border-text-warning bg-components-badge-bg-dimm px-1.5'
+        className='system-2xs-medium-uppercase relative flex h-[18px] items-center rounded-[5px] border border-text-warning bg-components-badge-bg-dimm px-1.5 text-text-warning'
         onClick={onClick}
       >
         <RiScales3Line className='mr-0.5 h-3 w-3' />
         {t('common.modelProvider.auth.authorizationError')}
+        <Indicator color='orange' className='absolute right-[-1px] top-[-1px] h-1.5 w-1.5' />
       </div>
     )
   }
@@ -38,13 +40,21 @@ const ConfigModel = ({
       variant='secondary'
       size='small'
       className={cn(
-        'shrink-0',
-        className,
+        'hidden shrink-0 group-hover:flex',
+        credentialRemoved && 'flex',
       )}
       onClick={onClick}
     >
       {
-        !loadBalancingEnabled && (
+        credentialRemoved && (
+          <>
+            {t('common.modelProvider.auth.credentialRemoved')}
+            <Indicator color='red' className='ml-2' />
+          </>
+        )
+      }
+      {
+        !loadBalancingEnabled && !credentialRemoved && (
           <>
             <RiEqualizer2Line className='mr-1 h-4 w-4' />
             {t('common.operation.config')}
@@ -52,7 +62,7 @@ const ConfigModel = ({
         )
       }
       {
-        loadBalancingEnabled && !loadBalancingInvalid && (
+        loadBalancingEnabled && !loadBalancingInvalid && !credentialRemoved && (
           <>
             <RiScales3Line className='mr-1 h-4 w-4' />
             {t('common.modelProvider.auth.configLoadBalancing')}
diff --git a/web/app/components/header/account-setting/model-provider-page/model-auth/switch-credential-in-load-balancing.tsx b/web/app/components/header/account-setting/model-provider-page/model-auth/switch-credential-in-load-balancing.tsx
index fca3e76d1f..4a45cd761f 100644
--- a/web/app/components/header/account-setting/model-provider-page/model-auth/switch-credential-in-load-balancing.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/model-auth/switch-credential-in-load-balancing.tsx
@@ -7,40 +7,38 @@ import { useTranslation } from 'react-i18next'
 import { RiArrowDownSLine } from '@remixicon/react'
 import Button from '@/app/components/base/button'
 import Indicator from '@/app/components/header/indicator'
-import Badge from '@/app/components/base/badge'
 import Authorized from './authorized'
 import type {
-  ModelLoadBalancingConfig,
+  Credential,
+  CustomModel,
   ModelProvider,
 } from '../declarations'
 import { ConfigurationMethodEnum } from '@/app/components/header/account-setting/model-provider-page/declarations'
-import { useCredentialStatus } from './hooks'
 import cn from '@/utils/classnames'
 
 type SwitchCredentialInLoadBalancingProps = {
   provider: ModelProvider
-  draftConfig?: ModelLoadBalancingConfig
-  setDraftConfig: Dispatch<SetStateAction<ModelLoadBalancingConfig | undefined>>
+  model: CustomModel
+  credentials?: Credential[]
+  customModelCredential?: Credential
+  setCustomModelCredential: Dispatch<SetStateAction<Credential | undefined>>
 }
 const SwitchCredentialInLoadBalancing = ({
   provider,
-  draftConfig,
+  model,
+  customModelCredential,
+  setCustomModelCredential,
+  credentials,
 }: SwitchCredentialInLoadBalancingProps) => {
   const { t } = useTranslation()
-  const {
-    available_credentials,
-    current_credential_name,
-  } = useCredentialStatus(provider)
 
-  const handleItemClick = useCallback(() => {
-    console.log('handleItemClick', draftConfig)
-  }, [])
+  const handleItemClick = useCallback((credential: Credential) => {
+    setCustomModelCredential(credential)
+  }, [setCustomModelCredential])
 
   const renderTrigger = useCallback(() => {
-    const selectedCredentialId = draftConfig?.configs.find(config => config.name === '__inherit__')?.credential_id
-    const selectedCredential = available_credentials?.find(credential => credential.credential_id === selectedCredentialId)
-    const name = selectedCredential?.credential_name || current_credential_name
-    const authRemoved = !!selectedCredentialId && !selectedCredential
+    const selectedCredentialId = customModelCredential?.credential_id
+    const authRemoved = !selectedCredentialId && !!credentials?.length
     return (
       <Button
         variant='secondary'
@@ -54,17 +52,12 @@ const SwitchCredentialInLoadBalancing = ({
           color={authRemoved ? 'red' : 'green'}
         />
         {
-          authRemoved ? t('common.model.authRemoved') : name
-        }
-        {
-          !authRemoved && (
-            <Badge>enterprise</Badge>
-          )
+          authRemoved ? t('common.modelProvider.auth.authRemoved') : customModelCredential?.credential_name
         }
         <RiArrowDownSLine className='h-4 w-4' />
       </Button>
     )
-  }, [current_credential_name, t, draftConfig, available_credentials])
+  }, [customModelCredential, t, credentials])
 
   return (
     <Authorized
@@ -72,17 +65,25 @@ const SwitchCredentialInLoadBalancing = ({
       configurationMethod={ConfigurationMethodEnum.customizableModel}
       items={[
         {
-          model: {
-            model: t('common.modelProvider.modelCredentials'),
-          } as any,
-          credentials: available_credentials || [],
+          title: t('common.modelProvider.auth.modelCredentials'),
+          model,
+          credentials: credentials || [],
         },
       ]}
       renderTrigger={renderTrigger}
       onItemClick={handleItemClick}
       isModelCredential
       enableAddModelCredential
-      bottomAddModelCredentialText={t('common.modelProvider.addModelCredential')}
+      bottomAddModelCredentialText={t('common.modelProvider.auth.addModelCredential')}
+      selectedCredential={
+        customModelCredential
+          ? {
+            credential_id: customModelCredential?.credential_id || '',
+            credential_name: customModelCredential?.credential_name || '',
+          }
+          : undefined
+      }
+      showItemSelectedIcon
     />
   )
 }
diff --git a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-list-item.tsx b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-list-item.tsx
index 5093904b2f..020cf83e43 100644
--- a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-list-item.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-list-item.tsx
@@ -3,7 +3,6 @@ import { useTranslation } from 'react-i18next'
 import { useDebounceFn } from 'ahooks'
 import type { ModelItem, ModelProvider } from '../declarations'
 import { ModelStatusEnum } from '../declarations'
-import ModelBadge from '../model-badge'
 import ModelIcon from '../model-icon'
 import ModelName from '../model-name'
 import classNames from '@/utils/classnames'
@@ -15,6 +14,7 @@ import { disableModel, enableModel } from '@/service/common'
 import { Plan } from '@/app/components/billing/type'
 import { useAppContext } from '@/context/app-context'
 import { ConfigModel } from '../model-auth'
+import Badge from '@/app/components/base/badge'
 
 export type ModelListItemProps = {
   model: ModelItem
@@ -63,21 +63,20 @@ const ModelListItem = ({ model, provider, isConfigurable, onModifyLoadBalancing
         showMode
         showContextSize
       >
-        {modelLoadBalancingEnabled && !model.deprecated && model.load_balancing_enabled && (
-          <ModelBadge className='ml-1 border-text-accent-secondary uppercase text-text-accent-secondary'>
-            <Balance className='mr-0.5 h-3 w-3' />
-            {t('common.modelProvider.loadBalancingHeadline')}
-          </ModelBadge>
-        )}
       </ModelName>
       <div className='flex shrink-0 items-center'>
+        {modelLoadBalancingEnabled && !model.deprecated && model.load_balancing_enabled && (
+          <Badge className='mr-1 h-[18px] w-[18px] items-center justify-center border-text-accent-secondary p-0'>
+            <Balance className='h-3 w-3 text-text-accent-secondary' />
+          </Badge>
+        )}
         {
           (isCurrentWorkspaceManager && (modelLoadBalancingEnabled || plan.type === Plan.sandbox) && !model.deprecated && [ModelStatusEnum.active, ModelStatusEnum.disabled].includes(model.status)) && (
             <ConfigModel
-              className='hidden group-hover:flex'
               onClick={() => onModifyLoadBalancing?.(model)}
               loadBalancingEnabled={model.load_balancing_enabled}
               loadBalancingInvalid={model.has_invalid_load_balancing_configs}
+              credentialRemoved={model.status === ModelStatusEnum.credentialRemoved}
             />
           )
         }
diff --git a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-configs.tsx b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-configs.tsx
index 28fc5d39b7..5402ba4d73 100644
--- a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-configs.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-configs.tsx
@@ -1,5 +1,5 @@
 import type { Dispatch, SetStateAction } from 'react'
-import { useCallback } from 'react'
+import { useCallback, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
 import {
   RiDeleteBinLine,
@@ -122,6 +122,20 @@ const ModelLoadBalancingConfigs = ({
     })
   }, [updateConfigEntry])
 
+  const validDraftConfigList = useMemo(() => {
+    if (!draftConfig)
+      return []
+    return draftConfig.configs.filter((config) => {
+      if (config.name === '__inherit__')
+        return true
+
+      if (config.credential_id)
+        return true
+
+      return false
+    })
+  }, [draftConfig])
+
   if (!draftConfig)
     return null
 
@@ -165,15 +179,7 @@ const ModelLoadBalancingConfigs = ({
         </div>
         {draftConfig.enabled && (
           <div className='flex flex-col gap-1 px-3 pb-3'>
-            {draftConfig.configs.filter((config) => {
-              if (config.name === '__inherit__')
-                return true
-
-              if (config.credential_id)
-                return true
-
-              return false
-            }).map((config, index) => {
+            {validDraftConfigList.map((config, index) => {
               const isProviderManaged = config.name === '__inherit__'
               return (
                 <div key={config.id || index} className='group flex h-10 items-center rounded-lg border border-components-panel-border bg-components-panel-on-panel-item-bg px-3 shadow-xs'>
@@ -249,7 +255,7 @@ const ModelLoadBalancingConfigs = ({
           </div>
         )}
         {
-          draftConfig.enabled && draftConfig.configs.length < 2 && (
+          draftConfig.enabled && validDraftConfigList.length < 2 && (
             <div className='flex h-[34px] items-center rounded-b-xl border-t border-t-divider-subtle bg-components-panel-bg px-6 text-xs text-text-secondary'>
               <AlertTriangle className='mr-1 h-3 w-3 text-[#f79009]' />
               {t('common.modelProvider.loadBalancingLeastKeyWarning')}
diff --git a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-modal.tsx b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-modal.tsx
index 7f2b88a186..9228fa5c93 100644
--- a/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-modal.tsx
+++ b/web/app/components/header/account-setting/model-provider-page/provider-added-card/model-load-balancing-modal.tsx
@@ -19,7 +19,7 @@ import Modal from '@/app/components/base/modal'
 import Button from '@/app/components/base/button'
 import Loading from '@/app/components/base/loading'
 import { useToastContext } from '@/app/components/base/toast'
-// import { SwitchCredentialInLoadBalancing } from '@/app/components/header/account-setting/model-provider-page/model-auth'
+import { SwitchCredentialInLoadBalancing } from '@/app/components/header/account-setting/model-provider-page/model-auth'
 import {
   useGetModelCredential,
   useUpdateModelLoadBalancingConfig,
@@ -59,6 +59,9 @@ const ModelLoadBalancingModal = ({
   const modelCredential = data
   const {
     load_balancing,
+    current_credential_id,
+    available_credentials,
+    current_credential_name,
   } = modelCredential ?? {}
   const originalConfig = load_balancing
   const [draftConfig, setDraftConfig] = useState<ModelLoadBalancingConfig>()
@@ -102,11 +105,21 @@ const ModelLoadBalancingModal = ({
   }, [extendedSecretFormSchemas, originalConfigMap])
 
   const { mutateAsync: updateModelLoadBalancingConfig } = useUpdateModelLoadBalancingConfig(provider.provider)
+  const initialCustomModelCredential = useMemo(() => {
+    if (!current_credential_id)
+      return undefined
+    return {
+      credential_id: current_credential_id,
+      credential_name: current_credential_name,
+    }
+  }, [current_credential_id, current_credential_name])
+  const [customModelCredential, setCustomModelCredential] = useState<Credential | undefined>(initialCustomModelCredential)
   const handleSave = async () => {
     try {
       setLoading(true)
       const res = await updateModelLoadBalancingConfig(
         {
+          credential_id: customModelCredential?.credential_id || current_credential_id,
           config_from: configFrom,
           model: model.model,
           model_type: model.model_type,
@@ -178,14 +191,28 @@ const ModelLoadBalancingModal = ({
                     )}
                   </div>
                   <div className='grow'>
-                    <div className='text-sm text-text-secondary'>{t('common.modelProvider.providerManaged')}</div>
-                    <div className='text-xs text-text-tertiary'>{t('common.modelProvider.providerManagedDescription')}</div>
+                    <div className='text-sm text-text-secondary'>{
+                      providerFormSchemaPredefined
+                       ? t('common.modelProvider.auth.providerManaged')
+                       : t('common.modelProvider.auth.specifyModelCredential')
+                    }</div>
+                    <div className='text-xs text-text-tertiary'>{
+                      providerFormSchemaPredefined
+                        ? t('common.modelProvider.auth.providerManagedTip')
+                        : t('common.modelProvider.auth.specifyModelCredentialTip')
+                    }</div>
                   </div>
-                  {/* <SwitchCredentialInLoadBalancing
-                    draftConfig={draftConfig}
-                    setDraftConfig={setDraftConfig}
-                    provider={provider}
-                  /> */}
+                  {
+                    !providerFormSchemaPredefined && (
+                      <SwitchCredentialInLoadBalancing
+                        provider={provider}
+                        customModelCredential={initialCustomModelCredential ?? customModelCredential}
+                        setCustomModelCredential={setCustomModelCredential}
+                        model={model}
+                        credentials={available_credentials}
+                      />
+                    )
+                  }
                 </div>
               </div>
               {
diff --git a/web/app/components/plugins/plugin-auth/types.ts b/web/app/components/plugins/plugin-auth/types.ts
index df774d5d25..93ec4c8c1a 100644
--- a/web/app/components/plugins/plugin-auth/types.ts
+++ b/web/app/components/plugins/plugin-auth/types.ts
@@ -23,4 +23,5 @@ export type Credential = {
   credentials?: Record<string, any>
   isWorkspaceDefault?: boolean
   from_enterprise?: boolean
+  allowed_to_use?: boolean
 }
diff --git a/web/i18n/en-US/common.ts b/web/i18n/en-US/common.ts
index 7226c9850f..f4622c3068 100644
--- a/web/i18n/en-US/common.ts
+++ b/web/i18n/en-US/common.ts
@@ -467,7 +467,7 @@ const translation = {
     loadPresets: 'Load Presets',
     parameters: 'PARAMETERS',
     loadBalancing: 'Load balancing',
-    loadBalancingDescription: 'Reduce pressure with multiple sets of credentials.',
+    loadBalancingDescription: 'Configure multiple credentials for the model and invoke them automatically. ',
     loadBalancingHeadline: 'Load Balancing',
     configLoadBalancing: 'Config Load Balancing',
     modelHasBeenDeprecated: 'This model has been deprecated',
@@ -499,6 +499,10 @@ const translation = {
       configModel: 'Config model',
       configLoadBalancing: 'Config Load Balancing',
       authorizationError: 'Authorization error',
+      specifyModelCredential: 'Specify model credential',
+      specifyModelCredentialTip: 'Use a configured model credential.',
+      providerManaged: 'Provider managed',
+      providerManagedTip: 'The current configuration is hosted by the provider.',
     },
   },
   dataSource: {
diff --git a/web/i18n/zh-Hans/common.ts b/web/i18n/zh-Hans/common.ts
index e91407872a..fcbf825bbc 100644
--- a/web/i18n/zh-Hans/common.ts
+++ b/web/i18n/zh-Hans/common.ts
@@ -466,7 +466,7 @@ const translation = {
     loadPresets: '加载预设',
     parameters: '参数',
     loadBalancing: '负载均衡',
-    loadBalancingDescription: '为了减轻单组凭据的压力，您可以为模型调用配置多组凭据。',
+    loadBalancingDescription: '为模型配置多组凭据，并自动调用。',
     loadBalancingHeadline: '负载均衡',
     configLoadBalancing: '设置负载均衡',
     modelHasBeenDeprecated: '该模型已废弃',
@@ -499,6 +499,10 @@ const translation = {
       configModel: '配置模型',
       configLoadBalancing: '配置负载均衡',
       authorizationError: '授权错误',
+      specifyModelCredential: '指定模型凭据',
+      specifyModelCredentialTip: '使用已配置的模型凭据。',
+      providerManaged: '由模型供应商管理',
+      providerManagedTip: '使用模型供应商提供的单组凭据。',
     },
   },
   dataSource: {
diff --git a/web/service/use-models.ts b/web/service/use-models.ts
index 75321d2706..f3336dd03b 100644
--- a/web/service/use-models.ts
+++ b/web/service/use-models.ts
@@ -147,6 +147,7 @@ export const useUpdateModelLoadBalancingConfig = (provider: string) => {
       model: string
       model_type: ModelTypeEnum
       load_balancing: ModelLoadBalancingConfig
+      credential_id?: string
     }) => post<{ result: string }>(`/workspaces/current/model-providers/${provider}/models`, {
       body: data,
     }),