diff --git a/internal/handler/key_handler.go b/internal/handler/key_handler.go index b0c7755dd..37d89e017 100644 --- a/internal/handler/key_handler.go +++ b/internal/handler/key_handler.go @@ -503,6 +503,11 @@ type UpdateKeyNotesRequest struct { Notes string `json:"notes"` } +// UpdateKeyEnabledRequest defines the payload for toggling a key's manual enabled switch. +type UpdateKeyEnabledRequest struct { + Enabled bool `json:"enabled"` +} + // UpdateKeyNotes handles updating the notes of a specific API key. func (s *Server) UpdateKeyNotes(c *gin.Context) { keyIDStr := c.Param("id") @@ -544,3 +549,35 @@ func (s *Server) UpdateKeyNotes(c *gin.Context) { response.Success(c, nil) } + +// UpdateKeyEnabled handles toggling whether a key can enter the active routing pool. +func (s *Server) UpdateKeyEnabled(c *gin.Context) { + keyIDStr := c.Param("id") + keyID, err := strconv.Atoi(keyIDStr) + if err != nil || keyID <= 0 { + response.Error(c, app_errors.NewAPIError(app_errors.ErrBadRequest, "invalid key ID format")) + return + } + + var req UpdateKeyEnabledRequest + if err := c.ShouldBindJSON(&req); err != nil { + response.Error(c, app_errors.NewAPIError(app_errors.ErrInvalidJSON, err.Error())) + return + } + + key, err := s.KeyService.SetKeyEnabled(uint(keyID), req.Enabled) + if err != nil { + if strings.Contains(err.Error(), "record not found") { + response.Error(c, app_errors.ErrResourceNotFound) + } else { + response.Error(c, app_errors.ParseDBError(err)) + } + return + } + + response.Success(c, gin.H{ + "id": key.ID, + "enabled": key.Enabled, + "status": key.Status, + }) +} diff --git a/internal/keypool/provider.go b/internal/keypool/provider.go index 7af1f0e8e..336e2f8e4 100644 --- a/internal/keypool/provider.go +++ b/internal/keypool/provider.go @@ -79,6 +79,7 @@ func (p *KeyProvider) SelectKey(groupID uint) (*models.APIKey, error) { ID: uint(keyID), KeyValue: decryptedKeyValue, Status: keyDetails["status"], + Enabled: parseEnabled(keyDetails), FailureCount: failureCount, GroupID: groupID, CreatedAt: time.Unix(createdAt, 0), @@ -147,6 +148,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st failureCount, _ := strconv.ParseInt(keyDetails["failure_count"], 10, 64) isActive := keyDetails["status"] == models.KeyStatusActive + isEnabled := parseEnabled(keyDetails) if failureCount == 0 && isActive { return nil @@ -171,7 +173,7 @@ func (p *KeyProvider) handleSuccess(keyID uint, keyHashKey, activeKeysListKey st return fmt.Errorf("failed to update key details in store: %w", err) } - if !isActive { + if !isActive && isEnabled { logrus.WithField("keyID", keyID).Debug("Key has recovered and is being restored to active pool.") if err := p.store.LRem(activeKeysListKey, 0, keyID); err != nil { return fmt.Errorf("failed to LRem key before LPush on recovery: %w", err) @@ -265,7 +267,7 @@ func (p *KeyProvider) LoadKeysFromDB() error { } } - if key.Status == models.KeyStatusActive { + if isKeyAvailableForPool(key) { allActiveKeyIDs[key.GroupID] = append(allActiveKeyIDs[key.GroupID], key.ID) } } @@ -560,7 +562,7 @@ func (p *KeyProvider) addKeyToStore(key *models.APIKey) error { } // 2. If active, add to the active LIST - if key.Status == models.KeyStatusActive { + if isKeyAvailableForPool(key) { activeKeysListKey := fmt.Sprintf("group:%d:active_keys", key.GroupID) if err := p.store.LRem(activeKeysListKey, 0, key.ID); err != nil { return fmt.Errorf("failed to LRem key %d before LPush for group %d: %w", key.ID, key.GroupID, err) @@ -601,19 +603,65 @@ func (p *KeyProvider) addKeysToCacheBatch(groupID uint, keys []models.APIKey) er // 2. 收集所有密钥 ID activeKeysListKey := fmt.Sprintf("group:%d:active_keys", groupID) - activeKeyIDs := make([]any, len(keys)) + activeKeyIDs := make([]any, 0, len(keys)) for i := range keys { - activeKeyIDs[i] = keys[i].ID + if isKeyAvailableForPool(&keys[i]) { + activeKeyIDs = append(activeKeyIDs, keys[i].ID) + } } // 3. 批量 LPush 活跃密钥 - if err := p.store.LPush(activeKeysListKey, activeKeyIDs...); err != nil { - return fmt.Errorf("failed to batch LPush keys to group %d: %w", groupID, err) + if len(activeKeyIDs) > 0 { + if err := p.store.LPush(activeKeysListKey, activeKeyIDs...); err != nil { + return fmt.Errorf("failed to batch LPush keys to group %d: %w", groupID, err) + } } return nil } +// SetKeyEnabled updates a key's manual enabled switch and keeps the active pool in sync. +func (p *KeyProvider) SetKeyEnabled(keyID uint, enabled bool) (*models.APIKey, error) { + var updatedKey models.APIKey + + err := p.executeTransactionWithRetry(func(tx *gorm.DB) error { + var key models.APIKey + if err := tx.Set("gorm:query_option", "FOR UPDATE").First(&key, keyID).Error; err != nil { + return fmt.Errorf("failed to lock key %d for enabled update: %w", keyID, err) + } + + if err := tx.Model(&key).Update("enabled", enabled).Error; err != nil { + return fmt.Errorf("failed to update key enabled flag in DB: %w", err) + } + key.Enabled = enabled + updatedKey = key + + keyHashKey := fmt.Sprintf("key:%d", key.ID) + activeKeysListKey := fmt.Sprintf("group:%d:active_keys", key.GroupID) + + if err := p.store.HSet(keyHashKey, p.apiKeyToMap(&key)); err != nil { + return fmt.Errorf("failed to update key enabled flag in store: %w", err) + } + + if err := p.store.LRem(activeKeysListKey, 0, key.ID); err != nil { + return fmt.Errorf("failed to LRem key from active list: %w", err) + } + + if isKeyAvailableForPool(&key) { + if err := p.store.LPush(activeKeysListKey, key.ID); err != nil { + return fmt.Errorf("failed to LPush key to active list: %w", err) + } + } + + return nil + }) + if err != nil { + return nil, err + } + + return &updatedKey, nil +} + // removeKeyFromStore is a helper to remove a single key from the cache. func (p *KeyProvider) removeKeyFromStore(keyID, groupID uint) error { activeKeysListKey := fmt.Sprintf("group:%d:active_keys", groupID) @@ -634,12 +682,22 @@ func (p *KeyProvider) apiKeyToMap(key *models.APIKey) map[string]any { "id": fmt.Sprint(key.ID), "key_string": key.KeyValue, "status": key.Status, + "enabled": key.Enabled, "failure_count": key.FailureCount, "group_id": key.GroupID, "created_at": key.CreatedAt.Unix(), } } +func isKeyAvailableForPool(key *models.APIKey) bool { + return key.Enabled && key.Status == models.KeyStatusActive +} + +func parseEnabled(details map[string]string) bool { + enabled, ok := details["enabled"] + return !ok || enabled == "" || enabled == "true" || enabled == "1" +} + // pluckIDs extracts IDs from a slice of APIKey. func pluckIDs(keys []models.APIKey) []uint { ids := make([]uint, len(keys)) diff --git a/internal/models/types.go b/internal/models/types.go index 27089f856..a370b4234 100644 --- a/internal/models/types.go +++ b/internal/models/types.go @@ -119,6 +119,7 @@ type APIKey struct { KeyHash string `gorm:"type:varchar(128);index" json:"key_hash"` GroupID uint `gorm:"not null;index;index:idx_api_keys_group_last_used_id,priority:1" json:"group_id"` Status string `gorm:"type:varchar(50);not null;default:'active';index" json:"status"` + Enabled bool `gorm:"not null;default:true;index" json:"enabled"` Notes string `gorm:"type:varchar(255);default:''" json:"notes"` RequestCount int64 `gorm:"not null;default:0" json:"request_count"` FailureCount int64 `gorm:"not null;default:0" json:"failure_count"` diff --git a/internal/router/router.go b/internal/router/router.go index 22edd3dc6..3548f996c 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -142,6 +142,7 @@ func registerProtectedAPIRoutes(api *gin.RouterGroup, serverHandler *handler.Ser keys.POST("/validate-group", serverHandler.ValidateGroupKeys) keys.POST("/test-multiple", serverHandler.TestMultipleKeys) keys.PUT("/:id/notes", serverHandler.UpdateKeyNotes) + keys.PUT("/:id/enabled", serverHandler.UpdateKeyEnabled) } // Tasks diff --git a/internal/services/key_service.go b/internal/services/key_service.go index c3e7096ff..657b9ebc5 100644 --- a/internal/services/key_service.go +++ b/internal/services/key_service.go @@ -130,6 +130,7 @@ func (s *KeyService) processAndCreateKeys( KeyValue: encryptedKey, KeyHash: keyHash, Status: models.KeyStatusActive, + Enabled: true, }) } @@ -251,6 +252,11 @@ func (s *KeyService) ClearAllKeys(groupID uint) (int64, error) { return s.KeyProvider.RemoveAllKeys(groupID) } +// SetKeyEnabled updates whether a key is allowed to participate in request routing. +func (s *KeyService) SetKeyEnabled(keyID uint, enabled bool) (*models.APIKey, error) { + return s.KeyProvider.SetKeyEnabled(keyID, enabled) +} + // DeleteMultipleKeys handles the business logic of deleting keys from a text block. func (s *KeyService) DeleteMultipleKeys(groupID uint, keysText string) (*DeleteKeysResult, error) { keysToDelete := s.ParseKeysFromText(keysText) diff --git a/web/src/api/keys.ts b/web/src/api/keys.ts index 09cc35790..4a9305514 100644 --- a/web/src/api/keys.ts +++ b/web/src/api/keys.ts @@ -140,6 +140,19 @@ export const keysApi = { await http.put(`/keys/${keyId}/notes`, { notes }, { hideMessage: true }); }, + // 更新密钥启用状态 + async updateKeyEnabled( + keyId: number, + enabled: boolean + ): Promise> { + const res = await http.put( + `/keys/${keyId}/enabled`, + { enabled }, + { hideMessage: true } + ); + return res.data; + }, + // 测试密钥 async testKeys( group_id: number, diff --git a/web/src/components/keys/KeyTable.vue b/web/src/components/keys/KeyTable.vue index 7ac047580..8638d91eb 100644 --- a/web/src/components/keys/KeyTable.vue +++ b/web/src/components/keys/KeyTable.vue @@ -25,6 +25,7 @@ import { NSelect, NSpace, NSpin, + NSwitch, useDialog, type MessageReactive, } from "naive-ui"; @@ -89,6 +90,7 @@ const moreOptions = [ let testingMsg: MessageReactive | null = null; const isDeling = ref(false); const isRestoring = ref(false); +const togglingKeyIds = ref>(new Set()); const createDialogShow = ref(false); const deleteDialogShow = ref(false); @@ -323,6 +325,30 @@ async function saveKeyNotes() { } } +async function updateKeyEnabled(key: KeyRow, enabled: boolean) { + const previousEnabled = key.enabled; + const nextToggling = new Set(togglingKeyIds.value); + nextToggling.add(key.id); + togglingKeyIds.value = nextToggling; + key.enabled = enabled; + + try { + const updatedKey = await keysApi.updateKeyEnabled(key.id, enabled); + key.enabled = updatedKey.enabled; + window.$message.success(enabled ? t("keys.keyEnabled") : t("keys.keyDisabled")); + if (props.selectedGroup) { + triggerSyncOperationRefresh(props.selectedGroup.name, "KEY_ENABLED_UPDATE"); + } + } catch (error) { + key.enabled = previousEnabled; + console.error("Update key enabled failed", error); + } finally { + const remaining = new Set(togglingKeyIds.value); + remaining.delete(key.id); + togglingKeyIds.value = remaining; + } +} + async function restoreKey(key: KeyRow) { if (!props.selectedGroup?.id || !key.key_value || isRestoring.value) { return; @@ -692,23 +718,39 @@ function resetPage() { v-for="key in keys" :key="key.id" class="key-card" - :class="getStatusClass(key.status)" + :class="[getStatusClass(key.status), { 'is-disabled': !key.enabled }]" >
-
- - - {{ t("keys.validShort") }} - - - - {{ t("keys.invalidShort") }} - +
+
+ + + {{ t("keys.validShort") }} + + + + {{ t("keys.invalidShort") }} + + + {{ t("keys.routingDisabledShort") }} + +
+ +
+
diff --git a/web/src/locales/en-US.ts b/web/src/locales/en-US.ts index 3942c7844..6cf64caa8 100644 --- a/web/src/locales/en-US.ts +++ b/web/src/locales/en-US.ts @@ -162,6 +162,12 @@ export default { blacklistCount: "Blacklist Count", valid: "Valid", invalid: "Invalid", + keyEnabled: "Key enabled", + keyDisabled: "Key disabled", + keyRoutingSwitch: "Key routing switch", + enableKeyRouting: "Enable key routing", + disableKeyRouting: "Disable key routing", + routingDisabledShort: "Off", checking: "Checking", unchecked: "Unchecked", addToBlacklist: "Add to Blacklist", diff --git a/web/src/locales/ja-JP.ts b/web/src/locales/ja-JP.ts index 4d9321714..805347d39 100644 --- a/web/src/locales/ja-JP.ts +++ b/web/src/locales/ja-JP.ts @@ -161,6 +161,12 @@ export default { blacklistCount: "ブラックリスト回数", valid: "有効", invalid: "無効", + keyEnabled: "キーを有効化しました", + keyDisabled: "キーを無効化しました", + keyRoutingSwitch: "キールーティングスイッチ", + enableKeyRouting: "キールーティングを有効化", + disableKeyRouting: "キールーティングを無効化", + routingDisabledShort: "停止", checking: "チェック中", unchecked: "未チェック", addToBlacklist: "ブラックリストに追加", diff --git a/web/src/locales/zh-CN.ts b/web/src/locales/zh-CN.ts index ca1a64a52..f64d4f49c 100644 --- a/web/src/locales/zh-CN.ts +++ b/web/src/locales/zh-CN.ts @@ -160,6 +160,12 @@ export default { blacklistCount: "黑名单次数", valid: "有效", invalid: "无效", + keyEnabled: "密钥已启用", + keyDisabled: "密钥已禁用", + keyRoutingSwitch: "密钥路由开关", + enableKeyRouting: "启用密钥路由", + disableKeyRouting: "禁用密钥路由", + routingDisabledShort: "停用", checking: "检查中", unchecked: "未检查", addToBlacklist: "加入黑名单", diff --git a/web/src/types/models.ts b/web/src/types/models.ts index bfdc91163..a98bb9d97 100644 --- a/web/src/types/models.ts +++ b/web/src/types/models.ts @@ -21,6 +21,7 @@ export interface APIKey { key_value: string; notes?: string; status: KeyStatus; + enabled: boolean; request_count: number; failure_count: number; last_used_at?: string;