From adfaf3c1e78fb9838d8be909469b5f09ddb19938 Mon Sep 17 00:00:00 2001 From: tbphp Date: Sat, 5 Jul 2025 20:54:57 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20key=E6=A3=80=E6=B5=8B=E9=87=8D=E8=AF=95?= =?UTF-8?q?=E5=8A=9F=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/gpt-load/main.go | 2 +- internal/services/key_cron_service.go | 3 +- internal/services/key_validator_service.go | 78 +++++++++++----------- 3 files changed, 41 insertions(+), 42 deletions(-) diff --git a/cmd/gpt-load/main.go b/cmd/gpt-load/main.go index aca736c..3f4f27c 100644 --- a/cmd/gpt-load/main.go +++ b/cmd/gpt-load/main.go @@ -86,7 +86,7 @@ func main() { taskService := services.NewTaskService(storage) channelFactory := channel.NewFactory(settingsManager) - keyValidatorService := services.NewKeyValidatorService(database, channelFactory) + keyValidatorService := services.NewKeyValidatorService(database, channelFactory, settingsManager) keyManualValidationService := services.NewKeyManualValidationService(database, keyValidatorService, taskService, settingsManager) keyCronService := services.NewKeyCronService(database, keyValidatorService, settingsManager) diff --git a/internal/services/key_cron_service.go b/internal/services/key_cron_service.go index 46a48dc..f69ce9d 100644 --- a/internal/services/key_cron_service.go +++ b/internal/services/key_cron_service.go @@ -56,7 +56,7 @@ func (s *KeyCronService) run() { // Dynamically get the interval for the next run intervalMinutes := s.SettingsManager.GetInt("key_validation_interval_minutes", 60) if intervalMinutes <= 0 { - intervalMinutes = 60 // Fallback to a safe default + intervalMinutes = 60 } nextRunTimer := time.NewTimer(time.Duration(intervalMinutes) * time.Minute) @@ -84,6 +84,7 @@ func (s *KeyCronService) validateAllGroups(ctx context.Context) { // Get effective settings for the group effectiveSettings := s.SettingsManager.GetEffectiveConfig(g.Config) interval := time.Duration(effectiveSettings.KeyValidationIntervalMinutes) * time.Minute + logrus.Infof("KeyCronService: Validating group %s with interval %s", g.Name, interval) // Check if it's time to validate this group if g.LastValidatedAt == nil || time.Since(*g.LastValidatedAt) > interval { diff --git a/internal/services/key_validator_service.go b/internal/services/key_validator_service.go index 6cc11d2..8908489 100644 --- a/internal/services/key_validator_service.go +++ b/internal/services/key_validator_service.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "gpt-load/internal/channel" + "gpt-load/internal/config" "gpt-load/internal/models" "github.com/sirupsen/logrus" @@ -19,21 +20,22 @@ type KeyTestResult struct { // KeyValidatorService provides methods to validate API keys. type KeyValidatorService struct { - DB *gorm.DB - channelFactory *channel.Factory + DB *gorm.DB + channelFactory *channel.Factory + SettingsManager *config.SystemSettingsManager } // NewKeyValidatorService creates a new KeyValidatorService. -func NewKeyValidatorService(db *gorm.DB, factory *channel.Factory) *KeyValidatorService { +func NewKeyValidatorService(db *gorm.DB, factory *channel.Factory, settingsManager *config.SystemSettingsManager) *KeyValidatorService { return &KeyValidatorService{ - DB: db, - channelFactory: factory, + DB: db, + channelFactory: factory, + SettingsManager: settingsManager, } } // ValidateSingleKey performs a validation check on a single API key. func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models.APIKey, group *models.Group) (bool, error) { - // 添加超时保护 if ctx.Err() != nil { return false, fmt.Errorf("context cancelled or timed out: %w", ctx.Err()) } @@ -49,56 +51,52 @@ func (s *KeyValidatorService) ValidateSingleKey(ctx context.Context, key *models return false, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err) } - // 记录验证开始 - logrus.WithFields(logrus.Fields{ - "key_id": key.ID, - "group_id": group.ID, - "group_name": group.Name, - }).Debug("Starting key validation") - - isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue) - if validationErr != nil { - logrus.WithFields(logrus.Fields{ - "key_id": key.ID, - "group_id": group.ID, - "group_name": group.Name, - "error": validationErr, - }).Debug("Key validation failed") - return false, validationErr + effectiveSettings := s.SettingsManager.GetEffectiveConfig(group.Config) + retries := effectiveSettings.BlacklistThreshold + if retries <= 0 { + retries = 1 } - // 记录验证结果 - logrus.WithFields(logrus.Fields{ - "key_id": key.ID, - "group_id": group.ID, - "group_name": group.Name, - "is_valid": isValid, - }).Debug("Key validation completed") + var lastErr error + for range retries { + isValid, validationErr := ch.ValidateKey(ctx, key.KeyValue) + if validationErr == nil && isValid { + logrus.WithFields(logrus.Fields{ + "key_id": key.ID, + "is_valid": isValid, + }).Debug("Key validation successful") + return true, nil + } - return isValid, nil + lastErr = validationErr + } + + logrus.WithFields(logrus.Fields{ + "key_id": key.ID, + "group_id": group.ID, + "max_retries": retries, + }).Debug("Key validation failed after all retries") + + return false, lastErr } // TestMultipleKeys performs a synchronous validation for a list of key values within a specific group. func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *models.Group, keyValues []string) ([]KeyTestResult, error) { results := make([]KeyTestResult, len(keyValues)) - ch, err := s.channelFactory.GetChannel(group) - if err != nil { - return nil, fmt.Errorf("failed to get channel for group %s: %w", group.Name, err) - } // Find which of the provided keys actually exist in the database for this group var existingKeys []models.APIKey if err := s.DB.Where("group_id = ? AND key_value IN ?", group.ID, keyValues).Find(&existingKeys).Error; err != nil { return nil, fmt.Errorf("failed to query keys from DB: %w", err) } - existingKeyMap := make(map[string]bool) + existingKeyMap := make(map[string]models.APIKey) for _, k := range existingKeys { - existingKeyMap[k.KeyValue] = true + existingKeyMap[k.KeyValue] = k } for i, kv := range keyValues { - // Pre-check: ensure the key belongs to the group to prevent unnecessary API calls - if !existingKeyMap[kv] { + apiKey, exists := existingKeyMap[kv] + if !exists { results[i] = KeyTestResult{ KeyValue: kv, IsValid: false, @@ -107,11 +105,11 @@ func (s *KeyValidatorService) TestMultipleKeys(ctx context.Context, group *model continue } - isValid, validationErr := ch.ValidateKey(ctx, kv) + isValid, validationErr := s.ValidateSingleKey(ctx, &apiKey, group) results[i] = KeyTestResult{ KeyValue: kv, IsValid: isValid, - Error: "", // Explicitly set error to empty string on success + Error: "", } if validationErr != nil { results[i].Error = validationErr.Error()