Pr/key enabled switch upstream

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,
+	})
+}

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))

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"`

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

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)

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<Pick<APIKey, "id" | "enabled" | "status">> {
+    const res = await http.put(
+      `/keys/${keyId}/enabled`,
+      { enabled },
+      { hideMessage: true }
+    );
+    return res.data;
+  },
+
   // 测试密钥
   async testKeys(
     group_id: number,