From d33ac5e1e9dad15e8de5be011a2e9eda921d538f Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Fri, 19 Jun 2026 04:26:04 +0800 Subject: [PATCH] feat(auth): add transient error cooldown configuration and adjust retry logic - Introduced `SetTransientErrorCooldownSeconds` to enable configurable cooldowns for transient errors (e.g., 408/500/502/503/504). - Updated retry scheduling logic to use the new `nextTransientErrorRetryAfter` function. - Modified config parsing to include `transient-error-cooldown-seconds` with support for disabling or defaulting to legacy behavior. - Expanded tests to validate transient cooldown logic with various configurations and edge cases. Closes: #3315 --- cmd/server/main.go | 1 + config.example.yaml | 4 + internal/api/server.go | 4 + internal/config/config.go | 5 + internal/config/parse.go | 1 + internal/watcher/diff/config_diff.go | 3 + internal/watcher/diff/config_diff_test.go | 100 +++++------ sdk/cliproxy/auth/conductor.go | 24 ++- sdk/cliproxy/auth/conductor_overrides_test.go | 157 ++++++++++++++++++ sdk/cliproxy/service.go | 1 + 10 files changed, 250 insertions(+), 50 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index e280b0db..a56ba307 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -505,6 +505,7 @@ func main() { redisqueue.SetUsageStatisticsEnabled(cfg.UsageStatisticsEnabled) redisqueue.SetRetentionSeconds(cfg.RedisUsageQueueRetentionSeconds) coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling) + coreauth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) if err = logging.ConfigureLogOutput(cfg); err != nil { log.Errorf("failed to configure log output: %v", err) diff --git a/config.example.yaml b/config.example.yaml index 551b4064..52d98f39 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -119,6 +119,10 @@ max-retry-interval: 30 # When true, disable auth/model cooldown scheduling globally (prevents blackout windows after failure states). disable-cooling: false +# Cooldown duration in seconds for transient upstream errors (408/500/502/503/504). +# Set to 0 to keep the legacy 60-second cooldown; set to -1 to disable transient error cooldowns. +transient-error-cooldown-seconds: 0 + # When true, globally disable Claude request cloaking (the Claude Code CLI disguise and # system prompt replacement), so the original system prompt is passed through to Claude as-is. # Individual credentials can still override this: a claude-api-key entry via its "cloak.mode", diff --git a/internal/api/server.go b/internal/api/server.go index 4572d3c1..beacc18a 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -328,6 +328,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk } managementasset.SetCurrentConfig(cfg) auth.SetQuotaCooldownDisabled(cfg.DisableCooling) + auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) applySignatureCacheConfig(nil, cfg) // Initialize management handler s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager) @@ -1596,6 +1597,9 @@ func (s *Server) UpdateClients(cfg *config.Config) { if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling { auth.SetQuotaCooldownDisabled(cfg.DisableCooling) } + if oldCfg == nil || oldCfg.TransientErrorCooldownSeconds != cfg.TransientErrorCooldownSeconds { + auth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) + } if oldCfg != nil && oldCfg.DisableImageGeneration != cfg.DisableImageGeneration { log.Infof("disable-image-generation updated: %v -> %v", oldCfg.DisableImageGeneration, cfg.DisableImageGeneration) diff --git a/internal/config/config.go b/internal/config/config.go index 9f8ba44e..2b5f18cd 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -80,6 +80,10 @@ type Config struct { // DisableCooling disables quota cooldown scheduling when true. DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"` + // TransientErrorCooldownSeconds controls cooldowns for transient upstream errors. + // 0 keeps the legacy default cooldown. Negative values disable these cooldowns. + TransientErrorCooldownSeconds int `yaml:"transient-error-cooldown-seconds" json:"transient-error-cooldown-seconds"` + // AuthAutoRefreshWorkers overrides the size of the core auth auto-refresh worker pool. // When <= 0, the default worker count is used. AuthAutoRefreshWorkers int `yaml:"auth-auto-refresh-workers" json:"auth-auto-refresh-workers"` @@ -684,6 +688,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { cfg.UsageStatisticsEnabled = false cfg.RedisUsageQueueRetentionSeconds = 60 cfg.DisableCooling = false + cfg.TransientErrorCooldownSeconds = 0 cfg.DisableImageGeneration = DisableImageGenerationOff cfg.Pprof.Enable = false cfg.Pprof.Addr = DefaultPprofAddr diff --git a/internal/config/parse.go b/internal/config/parse.go index b097976c..fb13447e 100644 --- a/internal/config/parse.go +++ b/internal/config/parse.go @@ -25,6 +25,7 @@ func ParseConfigBytes(data []byte) (*Config, error) { cfg.UsageStatisticsEnabled = false cfg.RedisUsageQueueRetentionSeconds = 60 cfg.DisableCooling = false + cfg.TransientErrorCooldownSeconds = 0 cfg.DisableImageGeneration = DisableImageGenerationOff cfg.Pprof.Enable = false cfg.Pprof.Addr = DefaultPprofAddr diff --git a/internal/watcher/diff/config_diff.go b/internal/watcher/diff/config_diff.go index 4b3799f5..3f508ebe 100644 --- a/internal/watcher/diff/config_diff.go +++ b/internal/watcher/diff/config_diff.go @@ -45,6 +45,9 @@ func BuildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldCfg.DisableCooling != newCfg.DisableCooling { changes = append(changes, fmt.Sprintf("disable-cooling: %t -> %t", oldCfg.DisableCooling, newCfg.DisableCooling)) } + if oldCfg.TransientErrorCooldownSeconds != newCfg.TransientErrorCooldownSeconds { + changes = append(changes, fmt.Sprintf("transient-error-cooldown-seconds: %d -> %d", oldCfg.TransientErrorCooldownSeconds, newCfg.TransientErrorCooldownSeconds)) + } if oldCfg.DisableClaudeCloakMode != newCfg.DisableClaudeCloakMode { changes = append(changes, fmt.Sprintf("disable-claude-cloak-mode: %t -> %t", oldCfg.DisableClaudeCloakMode, newCfg.DisableClaudeCloakMode)) } diff --git a/internal/watcher/diff/config_diff_test.go b/internal/watcher/diff/config_diff_test.go index e80bf017..13887db3 100644 --- a/internal/watcher/diff/config_diff_test.go +++ b/internal/watcher/diff/config_diff_test.go @@ -187,20 +187,21 @@ func TestBuildConfigChangeDetails_SecretsAndCounts(t *testing.T) { func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { oldCfg := &config.Config{ - Port: 1000, - AuthDir: "/old", - Debug: false, - LoggingToFile: false, - UsageStatisticsEnabled: false, - DisableCooling: false, - RequestRetry: 1, - MaxRetryCredentials: 1, - MaxRetryInterval: 1, - WebsocketAuth: false, - QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, - ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}}, - CodexKey: []config.CodexKey{{APIKey: "x1"}}, - RemoteManagement: config.RemoteManagement{DisableControlPanel: false, PanelGitHubRepository: "old/repo", SecretKey: "keep"}, + Port: 1000, + AuthDir: "/old", + Debug: false, + LoggingToFile: false, + UsageStatisticsEnabled: false, + DisableCooling: false, + TransientErrorCooldownSeconds: 0, + RequestRetry: 1, + MaxRetryCredentials: 1, + MaxRetryInterval: 1, + WebsocketAuth: false, + QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, + ClaudeKey: []config.ClaudeKey{{APIKey: "c1"}}, + CodexKey: []config.CodexKey{{APIKey: "x1"}}, + RemoteManagement: config.RemoteManagement{DisableControlPanel: false, PanelGitHubRepository: "old/repo", SecretKey: "keep"}, SDKConfig: sdkconfig.SDKConfig{ RequestLog: false, ProxyURL: "http://old-proxy", @@ -210,17 +211,18 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { }, } newCfg := &config.Config{ - Port: 2000, - AuthDir: "/new", - Debug: true, - LoggingToFile: true, - UsageStatisticsEnabled: true, - DisableCooling: true, - RequestRetry: 2, - MaxRetryCredentials: 3, - MaxRetryInterval: 3, - WebsocketAuth: true, - QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, + Port: 2000, + AuthDir: "/new", + Debug: true, + LoggingToFile: true, + UsageStatisticsEnabled: true, + DisableCooling: true, + TransientErrorCooldownSeconds: -1, + RequestRetry: 2, + MaxRetryCredentials: 3, + MaxRetryInterval: 3, + WebsocketAuth: true, + QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, ClaudeKey: []config.ClaudeKey{ {APIKey: "c1", BaseURL: "http://new", ProxyURL: "http://p", Headers: map[string]string{"H": "1"}, ExcludedModels: []string{"a"}}, {APIKey: "c2"}, @@ -250,6 +252,7 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { expectContains(t, details, "logging-to-file: false -> true") expectContains(t, details, "usage-statistics-enabled: false -> true") expectContains(t, details, "disable-cooling: false -> true") + expectContains(t, details, "transient-error-cooldown-seconds: 0 -> -1") expectContains(t, details, "disable-image-generation: false -> true") expectContains(t, details, "request-log: false -> true") expectContains(t, details, "request-retry: 1 -> 2") @@ -273,17 +276,18 @@ func TestBuildConfigChangeDetails_FlagsAndKeys(t *testing.T) { func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { oldCfg := &config.Config{ - Port: 1, - AuthDir: "/a", - Debug: false, - LoggingToFile: false, - UsageStatisticsEnabled: false, - DisableCooling: false, - RequestRetry: 1, - MaxRetryCredentials: 1, - MaxRetryInterval: 1, - WebsocketAuth: false, - QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, + Port: 1, + AuthDir: "/a", + Debug: false, + LoggingToFile: false, + UsageStatisticsEnabled: false, + DisableCooling: false, + TransientErrorCooldownSeconds: 0, + RequestRetry: 1, + MaxRetryCredentials: 1, + MaxRetryInterval: 1, + WebsocketAuth: false, + QuotaExceeded: config.QuotaExceeded{SwitchProject: false, SwitchPreviewModel: false, AntigravityCredits: false}, GeminiKey: []config.GeminiKey{ {APIKey: "g-old", BaseURL: "http://g-old", ProxyURL: "http://gp-old", Headers: map[string]string{"A": "1"}}, }, @@ -320,17 +324,18 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { }, } newCfg := &config.Config{ - Port: 2, - AuthDir: "/b", - Debug: true, - LoggingToFile: true, - UsageStatisticsEnabled: true, - DisableCooling: true, - RequestRetry: 2, - MaxRetryCredentials: 3, - MaxRetryInterval: 3, - WebsocketAuth: true, - QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, + Port: 2, + AuthDir: "/b", + Debug: true, + LoggingToFile: true, + UsageStatisticsEnabled: true, + DisableCooling: true, + TransientErrorCooldownSeconds: -1, + RequestRetry: 2, + MaxRetryCredentials: 3, + MaxRetryInterval: 3, + WebsocketAuth: true, + QuotaExceeded: config.QuotaExceeded{SwitchProject: true, SwitchPreviewModel: true, AntigravityCredits: true}, GeminiKey: []config.GeminiKey{ {APIKey: "g-new", BaseURL: "http://g-new", ProxyURL: "http://gp-new", Headers: map[string]string{"A": "2"}, ExcludedModels: []string{"x", "y"}}, }, @@ -380,6 +385,7 @@ func TestBuildConfigChangeDetails_AllBranches(t *testing.T) { expectContains(t, changes, "logging-to-file: false -> true") expectContains(t, changes, "usage-statistics-enabled: false -> true") expectContains(t, changes, "disable-cooling: false -> true") + expectContains(t, changes, "transient-error-cooldown-seconds: 0 -> -1") expectContains(t, changes, "disable-image-generation: false -> true") expectContains(t, changes, "request-retry: 1 -> 2") expectContains(t, changes, "max-retry-credentials: 1 -> 3") diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index cb3eb824..9d61f475 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -84,15 +84,23 @@ const ( refreshIneffectiveBackoff = 30 * time.Second quotaBackoffBase = time.Second quotaBackoffMax = 30 * time.Minute + transientErrorCooldown = time.Minute ) var quotaCooldownDisabled atomic.Bool +var transientErrorCooldownSeconds atomic.Int64 // SetQuotaCooldownDisabled toggles quota cooldown scheduling globally. func SetQuotaCooldownDisabled(disable bool) { quotaCooldownDisabled.Store(disable) } +// SetTransientErrorCooldownSeconds configures cooldowns for 408/500/502/503/504. +// 0 keeps the legacy default; negative values disable transient error cooldowns. +func SetTransientErrorCooldownSeconds(seconds int) { + transientErrorCooldownSeconds.Store(int64(seconds)) +} + func quotaCooldownDisabledForAuth(auth *Auth) bool { if auth != nil { if override, ok := auth.DisableCoolingOverride(); ok { @@ -102,6 +110,17 @@ func quotaCooldownDisabledForAuth(auth *Auth) bool { return quotaCooldownDisabled.Load() } +func nextTransientErrorRetryAfter(now time.Time) time.Time { + seconds := transientErrorCooldownSeconds.Load() + if seconds < 0 { + return time.Time{} + } + if seconds == 0 { + return now.Add(transientErrorCooldown) + } + return now.Add(time.Duration(seconds) * time.Second) +} + // Result captures execution outcome used to adjust auth state. type Result struct { // AuthID references the auth that produced this result. @@ -2909,8 +2928,7 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { if disableCooling { state.NextRetryAfter = time.Time{} } else { - next := now.Add(1 * time.Minute) - state.NextRetryAfter = next + state.NextRetryAfter = nextTransientErrorRetryAfter(now) } default: state.NextRetryAfter = time.Time{} @@ -3414,7 +3432,7 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati if disableCooling { auth.NextRetryAfter = time.Time{} } else { - auth.NextRetryAfter = now.Add(1 * time.Minute) + auth.NextRetryAfter = nextTransientErrorRetryAfter(now) } default: if auth.StatusMessage == "" { diff --git a/sdk/cliproxy/auth/conductor_overrides_test.go b/sdk/cliproxy/auth/conductor_overrides_test.go index 5acd331e..d4b10a32 100644 --- a/sdk/cliproxy/auth/conductor_overrides_test.go +++ b/sdk/cliproxy/auth/conductor_overrides_test.go @@ -522,6 +522,163 @@ func TestManager_MarkResult_RespectsAuthDisableCoolingOverride(t *testing.T) { } } +func TestManager_MarkResult_TransientErrorCooldownDefault(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + prevTransient := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(0) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + transientErrorCooldownSeconds.Store(prevTransient) + }) + + m := NewManager(nil, nil, nil) + + auth := &Auth{ + ID: "auth-transient-default", + Provider: "claude", + } + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + model := "test-model-transient-default" + m.MarkResult(context.Background(), Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadGateway, Message: "bad gateway"}, + }) + + updated, ok := m.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + state := updated.ModelStates[model] + if state == nil { + t.Fatalf("expected model state to be present") + } + if state.NextRetryAfter.IsZero() { + t.Fatal("expected transient error cooldown to keep the legacy default") + } + diff := time.Until(state.NextRetryAfter) + if diff < 55*time.Second || diff > 65*time.Second { + t.Fatalf("expected transient error cooldown to be ~60 seconds, got %v", diff) + } +} + +func TestManager_MarkResult_TransientErrorCooldownDisabled(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + prevTransient := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(-1) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + transientErrorCooldownSeconds.Store(prevTransient) + }) + + m := NewManager(nil, nil, nil) + + modelAuth := &Auth{ + ID: "auth-transient-model-disabled", + Provider: "claude", + } + if _, errRegisterModel := m.Register(context.Background(), modelAuth); errRegisterModel != nil { + t.Fatalf("register model auth: %v", errRegisterModel) + } + + model := "test-model-transient-disabled" + m.MarkResult(context.Background(), Result{ + AuthID: modelAuth.ID, + Provider: modelAuth.Provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusBadGateway, Message: "bad gateway"}, + }) + + updatedModelAuth, okModelAuth := m.GetByID(modelAuth.ID) + if !okModelAuth || updatedModelAuth == nil { + t.Fatalf("expected model auth to be present") + } + state := updatedModelAuth.ModelStates[model] + if state == nil { + t.Fatalf("expected model state to be present") + } + if !state.NextRetryAfter.IsZero() { + t.Fatalf("expected transient model cooldown to be disabled, got %v", state.NextRetryAfter) + } + + authLevelAuth := &Auth{ + ID: "auth-transient-auth-disabled", + Provider: "claude", + } + if _, errRegisterAuth := m.Register(context.Background(), authLevelAuth); errRegisterAuth != nil { + t.Fatalf("register auth-level auth: %v", errRegisterAuth) + } + + m.MarkResult(context.Background(), Result{ + AuthID: authLevelAuth.ID, + Provider: authLevelAuth.Provider, + Success: false, + Error: &Error{HTTPStatus: http.StatusServiceUnavailable, Message: "unavailable"}, + }) + + updatedAuthLevel, okAuthLevel := m.GetByID(authLevelAuth.ID) + if !okAuthLevel || updatedAuthLevel == nil { + t.Fatalf("expected auth-level auth to be present") + } + if !updatedAuthLevel.NextRetryAfter.IsZero() { + t.Fatalf("expected transient auth cooldown to be disabled, got %v", updatedAuthLevel.NextRetryAfter) + } +} + +func TestManager_MarkResult_TransientErrorCooldownDoesNotDisableAuthErrors(t *testing.T) { + prevQuota := quotaCooldownDisabled.Load() + quotaCooldownDisabled.Store(false) + prevTransient := transientErrorCooldownSeconds.Load() + SetTransientErrorCooldownSeconds(-1) + t.Cleanup(func() { + quotaCooldownDisabled.Store(prevQuota) + transientErrorCooldownSeconds.Store(prevTransient) + }) + + m := NewManager(nil, nil, nil) + + auth := &Auth{ + ID: "auth-transient-auth-error", + Provider: "claude", + } + if _, errRegister := m.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("register auth: %v", errRegister) + } + + model := "test-model-auth-error" + m.MarkResult(context.Background(), Result{ + AuthID: auth.ID, + Provider: auth.Provider, + Model: model, + Success: false, + Error: &Error{HTTPStatus: http.StatusForbidden, Message: "forbidden"}, + }) + + updated, ok := m.GetByID(auth.ID) + if !ok || updated == nil { + t.Fatalf("expected auth to be present") + } + state := updated.ModelStates[model] + if state == nil { + t.Fatalf("expected model state to be present") + } + if state.NextRetryAfter.IsZero() { + t.Fatal("expected auth error cooldown to remain enabled") + } + diff := time.Until(state.NextRetryAfter) + if diff < 29*time.Minute || diff > 31*time.Minute { + t.Fatalf("expected auth error cooldown to be ~30 minutes, got %v", diff) + } +} + func TestManager_MarkResult_RespectsAuthDisableCoolingOverride_On403(t *testing.T) { prev := quotaCooldownDisabled.Load() quotaCooldownDisabled.Store(false) diff --git a/sdk/cliproxy/service.go b/sdk/cliproxy/service.go index ab09f91f..589ecbdd 100644 --- a/sdk/cliproxy/service.go +++ b/sdk/cliproxy/service.go @@ -742,6 +742,7 @@ func (s *Service) applyRetryConfig(cfg *config.Config) { } maxInterval := time.Duration(cfg.MaxRetryInterval) * time.Second s.coreManager.SetRetryConfig(cfg.RequestRetry, maxInterval, cfg.MaxRetryCredentials) + coreauth.SetTransientErrorCooldownSeconds(cfg.TransientErrorCooldownSeconds) } func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName string, ok bool) {