From 9678be7aa435c91bec7a15bd17db1256a0d079d4 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 21 Oct 2025 21:51:30 +0800 Subject: [PATCH] feat: add DisableCooling configuration to manage quota cooldown behavior --- cmd/server/main.go | 2 ++ internal/api/server.go | 10 ++++++++++ internal/config/config.go | 4 ++++ internal/watcher/watcher.go | 3 +++ sdk/cliproxy/auth/manager.go | 24 +++++++++++++++++++++--- 5 files changed, 40 insertions(+), 3 deletions(-) diff --git a/cmd/server/main.go b/cmd/server/main.go index ce6770b5..78259928 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -27,6 +27,7 @@ import ( "github.com/router-for-me/CLIProxyAPI/v6/internal/usage" "github.com/router-for-me/CLIProxyAPI/v6/internal/util" sdkAuth "github.com/router-for-me/CLIProxyAPI/v6/sdk/auth" + coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" log "github.com/sirupsen/logrus" ) @@ -377,6 +378,7 @@ func main() { } } usage.SetStatisticsEnabled(cfg.UsageStatisticsEnabled) + coreauth.SetQuotaCooldownDisabled(cfg.DisableCooling) if err = logging.ConfigureLogOutput(cfg.LoggingToFile); err != nil { log.Fatalf("failed to configure log output: %v", err) diff --git a/internal/api/server.go b/internal/api/server.go index 1df9c7a4..aae2b0e0 100644 --- a/internal/api/server.go +++ b/internal/api/server.go @@ -233,6 +233,7 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk s.oldConfigYaml, _ = yaml.Marshal(cfg) s.applyAccessConfig(nil, cfg) managementasset.SetCurrentConfig(cfg) + auth.SetQuotaCooldownDisabled(cfg.DisableCooling) // Initialize management handler s.mgmt = managementHandlers.NewHandler(cfg, configFilePath, authManager) if optionState.localPassword != "" { @@ -716,6 +717,15 @@ func (s *Server) UpdateClients(cfg *config.Config) { } } + if oldCfg == nil || oldCfg.DisableCooling != cfg.DisableCooling { + auth.SetQuotaCooldownDisabled(cfg.DisableCooling) + if oldCfg != nil { + log.Debugf("disable_cooling updated from %t to %t", oldCfg.DisableCooling, cfg.DisableCooling) + } else { + log.Debugf("disable_cooling toggled to %t", cfg.DisableCooling) + } + } + // Update log level dynamically when debug flag changes if oldCfg == nil || oldCfg.Debug != cfg.Debug { util.SetLogLevel(cfg) diff --git a/internal/config/config.go b/internal/config/config.go index 68208e5a..169eecc2 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -34,6 +34,9 @@ type Config struct { // UsageStatisticsEnabled toggles in-memory usage aggregation; when false, usage data is discarded. UsageStatisticsEnabled bool `yaml:"usage-statistics-enabled" json:"usage-statistics-enabled"` + // DisableCooling disables quota cooldown scheduling when true. + DisableCooling bool `yaml:"disable-cooling" json:"disable-cooling"` + // QuotaExceeded defines the behavior when a quota is exceeded. QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"` @@ -183,6 +186,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) { // Set defaults before unmarshal so that absent keys keep defaults. cfg.LoggingToFile = false cfg.UsageStatisticsEnabled = false + cfg.DisableCooling = false if err = yaml.Unmarshal(data, &cfg); err != nil { if optional { // In cloud deploy mode, if YAML parsing fails, return empty config instead of error. diff --git a/internal/watcher/watcher.go b/internal/watcher/watcher.go index 640a3d2f..85b48aae 100644 --- a/internal/watcher/watcher.go +++ b/internal/watcher/watcher.go @@ -1192,6 +1192,9 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string { if oldCfg.UsageStatisticsEnabled != newCfg.UsageStatisticsEnabled { changes = append(changes, fmt.Sprintf("usage-statistics-enabled: %t -> %t", oldCfg.UsageStatisticsEnabled, newCfg.UsageStatisticsEnabled)) } + if oldCfg.DisableCooling != newCfg.DisableCooling { + changes = append(changes, fmt.Sprintf("disable-cooling: %t -> %t", oldCfg.DisableCooling, newCfg.DisableCooling)) + } if oldCfg.RequestLog != newCfg.RequestLog { changes = append(changes, fmt.Sprintf("request-log: %t -> %t", oldCfg.RequestLog, newCfg.RequestLog)) } diff --git a/sdk/cliproxy/auth/manager.go b/sdk/cliproxy/auth/manager.go index 10832d31..c2e87d9d 100644 --- a/sdk/cliproxy/auth/manager.go +++ b/sdk/cliproxy/auth/manager.go @@ -8,6 +8,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/google/uuid" @@ -44,6 +45,13 @@ const ( quotaBackoffMax = 30 * time.Minute ) +var quotaCooldownDisabled atomic.Bool + +// SetQuotaCooldownDisabled toggles quota cooldown scheduling globally. +func SetQuotaCooldownDisabled(disable bool) { + quotaCooldownDisabled.Store(disable) +} + // Result captures execution outcome used to adjust auth state. type Result struct { // AuthID references the auth that produced this result. @@ -535,7 +543,10 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) { shouldSuspendModel = true case 429: cooldown, nextLevel := nextQuotaCooldown(state.Quota.BackoffLevel) - next := now.Add(cooldown) + var next time.Time + if cooldown > 0 { + next = now.Add(cooldown) + } state.NextRetryAfter = next state.Quota = QuotaState{ Exceeded: true, @@ -750,9 +761,13 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, now time.Time) { auth.Quota.Exceeded = true auth.Quota.Reason = "quota" cooldown, nextLevel := nextQuotaCooldown(auth.Quota.BackoffLevel) - auth.Quota.NextRecoverAt = now.Add(cooldown) + var next time.Time + if cooldown > 0 { + next = now.Add(cooldown) + } + auth.Quota.NextRecoverAt = next auth.Quota.BackoffLevel = nextLevel - auth.NextRetryAfter = auth.Quota.NextRecoverAt + auth.NextRetryAfter = next case 408, 500, 502, 503, 504: auth.StatusMessage = "transient upstream error" auth.NextRetryAfter = now.Add(1 * time.Minute) @@ -768,6 +783,9 @@ func nextQuotaCooldown(prevLevel int) (time.Duration, int) { if prevLevel < 0 { prevLevel = 0 } + if quotaCooldownDisabled.Load() { + return 0, prevLevel + } cooldown := quotaBackoffBase * time.Duration(1<