feat(config): add per-auth disable_cooling override support

- Introduced `disable_cooling` metadata field for fine-grained control over cooldown scheduling.
- Updated `Auth` object to include `Metadata` with conditional logic for handling empty states.
- Added YAML configuration support for `disable_cooling` in API key definitions across providers.
- Enhanced unit tests to validate `disable_cooling` behavior in various scenarios.
This commit is contained in:
Luis Pater
2026-05-09 10:51:27 +08:00
parent c67096b687
commit 0f0fcd2304
5 changed files with 108 additions and 17 deletions
+4
View File
@@ -157,6 +157,7 @@ nonstream-keepalive-interval: 0
# gemini-api-key: # gemini-api-key:
# - api-key: "AIzaSy...01" # - api-key: "AIzaSy...01"
# prefix: "test" # optional: require calls like "test/gemini-3-pro-preview" to target this credential # prefix: "test" # optional: require calls like "test/gemini-3-pro-preview" to target this credential
# disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling
# base-url: "https://generativelanguage.googleapis.com" # base-url: "https://generativelanguage.googleapis.com"
# headers: # headers:
# X-Custom-Header: "custom-value" # X-Custom-Header: "custom-value"
@@ -176,6 +177,7 @@ nonstream-keepalive-interval: 0
# codex-api-key: # codex-api-key:
# - api-key: "sk-atSM..." # - api-key: "sk-atSM..."
# prefix: "test" # optional: require calls like "test/gpt-5-codex" to target this credential # prefix: "test" # optional: require calls like "test/gpt-5-codex" to target this credential
# disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling
# base-url: "https://www.example.com" # use the custom codex API endpoint # base-url: "https://www.example.com" # use the custom codex API endpoint
# headers: # headers:
# X-Custom-Header: "custom-value" # X-Custom-Header: "custom-value"
@@ -195,6 +197,7 @@ nonstream-keepalive-interval: 0
# - api-key: "sk-atSM..." # use the official claude API key, no need to set the base url # - api-key: "sk-atSM..." # use the official claude API key, no need to set the base url
# - api-key: "sk-atSM..." # - api-key: "sk-atSM..."
# prefix: "test" # optional: require calls like "test/claude-sonnet-latest" to target this credential # prefix: "test" # optional: require calls like "test/claude-sonnet-latest" to target this credential
# disable-cooling: false # optional: per-auth override for auth/model cooldown scheduling
# base-url: "https://www.example.com" # use the custom claude API endpoint # base-url: "https://www.example.com" # use the custom claude API endpoint
# headers: # headers:
# X-Custom-Header: "custom-value" # X-Custom-Header: "custom-value"
@@ -250,6 +253,7 @@ nonstream-keepalive-interval: 0
# disabled: false # optional: set to true to disable this provider without removing it # disabled: false # optional: set to true to disable this provider without removing it
# prefix: "test" # optional: require calls like "test/kimi-k2" to target this provider's credentials # prefix: "test" # optional: require calls like "test/kimi-k2" to target this provider's credentials
# base-url: "https://openrouter.ai/api/v1" # The base URL of the provider. # base-url: "https://openrouter.ai/api/v1" # The base URL of the provider.
# disable-cooling: false # optional: per-provider override for auth/model cooldown scheduling
# headers: # headers:
# X-Custom-Header: "custom-value" # X-Custom-Header: "custom-value"
# api-key-entries: # api-key-entries:
+12 -6
View File
@@ -226,12 +226,6 @@ type RoutingConfig struct {
// Supported values: "round-robin" (default), "fill-first". // Supported values: "round-robin" (default), "fill-first".
Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"` Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"`
// ClaudeCodeSessionAffinity enables session-sticky routing for Claude Code clients.
// When enabled, requests with the same session ID (extracted from metadata.user_id)
// are routed to the same auth credential when available.
// Deprecated: Use SessionAffinity instead for universal session support.
ClaudeCodeSessionAffinity bool `yaml:"claude-code-session-affinity,omitempty" json:"claude-code-session-affinity,omitempty"`
// SessionAffinity enables universal session-sticky routing for all clients. // SessionAffinity enables universal session-sticky routing for all clients.
// Session IDs are extracted from multiple sources: // Session IDs are extracted from multiple sources:
// metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex), // metadata.user_id (Claude Code session format), X-Session-ID, Session_id (Codex),
@@ -403,6 +397,9 @@ type ClaudeKey struct {
// ExcludedModels lists model IDs that should be excluded for this provider. // ExcludedModels lists model IDs that should be excluded for this provider.
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
// DisableCooling disables auth/model cooldown scheduling for this credential when true.
DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"`
// Cloak configures request cloaking for non-Claude-Code clients. // Cloak configures request cloaking for non-Claude-Code clients.
Cloak *CloakConfig `yaml:"cloak,omitempty" json:"cloak,omitempty"` Cloak *CloakConfig `yaml:"cloak,omitempty" json:"cloak,omitempty"`
@@ -458,6 +455,9 @@ type CodexKey struct {
// ExcludedModels lists model IDs that should be excluded for this provider. // ExcludedModels lists model IDs that should be excluded for this provider.
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
// DisableCooling disables auth/model cooldown scheduling for this credential when true.
DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"`
} }
func (k CodexKey) GetAPIKey() string { return k.APIKey } func (k CodexKey) GetAPIKey() string { return k.APIKey }
@@ -502,6 +502,9 @@ type GeminiKey struct {
// ExcludedModels lists model IDs that should be excluded for this provider. // ExcludedModels lists model IDs that should be excluded for this provider.
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"` ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
// DisableCooling disables auth/model cooldown scheduling for this credential when true.
DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"`
} }
func (k GeminiKey) GetAPIKey() string { return k.APIKey } func (k GeminiKey) GetAPIKey() string { return k.APIKey }
@@ -546,6 +549,9 @@ type OpenAICompatibility struct {
// Headers optionally adds extra HTTP headers for requests sent to this provider. // Headers optionally adds extra HTTP headers for requests sent to this provider.
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"` Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
// DisableCooling disables auth/model cooldown scheduling for this provider when true.
DisableCooling bool `yaml:"disable-cooling,omitempty" json:"disable-cooling,omitempty"`
} }
// OpenAICompatibilityAPIKey represents an API key configuration with optional proxy setting. // OpenAICompatibilityAPIKey represents an API key configuration with optional proxy setting.
+41
View File
@@ -60,6 +60,10 @@ func (s *ConfigSynthesizer) synthesizeGeminiKeys(ctx *SynthesisContext) []*corea
"source": fmt.Sprintf("config:gemini[%s]", token), "source": fmt.Sprintf("config:gemini[%s]", token),
"api_key": key, "api_key": key,
} }
metadata := map[string]any{}
if entry.DisableCooling {
metadata["disable_cooling"] = true
}
if entry.Priority != 0 { if entry.Priority != 0 {
attrs["priority"] = strconv.Itoa(entry.Priority) attrs["priority"] = strconv.Itoa(entry.Priority)
} }
@@ -78,10 +82,14 @@ func (s *ConfigSynthesizer) synthesizeGeminiKeys(ctx *SynthesisContext) []*corea
Status: coreauth.StatusActive, Status: coreauth.StatusActive,
ProxyURL: proxyURL, ProxyURL: proxyURL,
Attributes: attrs, Attributes: attrs,
Metadata: metadata,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
ApplyAuthExcludedModelsMeta(a, cfg, entry.ExcludedModels, "apikey") ApplyAuthExcludedModelsMeta(a, cfg, entry.ExcludedModels, "apikey")
if len(a.Metadata) == 0 {
a.Metadata = nil
}
out = append(out, a) out = append(out, a)
} }
return out return out
@@ -107,6 +115,10 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea
"source": fmt.Sprintf("config:claude[%s]", token), "source": fmt.Sprintf("config:claude[%s]", token),
"api_key": key, "api_key": key,
} }
metadata := map[string]any{}
if ck.DisableCooling {
metadata["disable_cooling"] = true
}
if ck.Priority != 0 { if ck.Priority != 0 {
attrs["priority"] = strconv.Itoa(ck.Priority) attrs["priority"] = strconv.Itoa(ck.Priority)
} }
@@ -126,10 +138,14 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea
Status: coreauth.StatusActive, Status: coreauth.StatusActive,
ProxyURL: proxyURL, ProxyURL: proxyURL,
Attributes: attrs, Attributes: attrs,
Metadata: metadata,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey") ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
if len(a.Metadata) == 0 {
a.Metadata = nil
}
out = append(out, a) out = append(out, a)
} }
return out return out
@@ -154,6 +170,10 @@ func (s *ConfigSynthesizer) synthesizeCodexKeys(ctx *SynthesisContext) []*coreau
"source": fmt.Sprintf("config:codex[%s]", token), "source": fmt.Sprintf("config:codex[%s]", token),
"api_key": key, "api_key": key,
} }
metadata := map[string]any{}
if ck.DisableCooling {
metadata["disable_cooling"] = true
}
if ck.Priority != 0 { if ck.Priority != 0 {
attrs["priority"] = strconv.Itoa(ck.Priority) attrs["priority"] = strconv.Itoa(ck.Priority)
} }
@@ -176,10 +196,14 @@ func (s *ConfigSynthesizer) synthesizeCodexKeys(ctx *SynthesisContext) []*coreau
Status: coreauth.StatusActive, Status: coreauth.StatusActive,
ProxyURL: proxyURL, ProxyURL: proxyURL,
Attributes: attrs, Attributes: attrs,
Metadata: metadata,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey") ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
if len(a.Metadata) == 0 {
a.Metadata = nil
}
out = append(out, a) out = append(out, a)
} }
return out return out
@@ -203,6 +227,7 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
providerName = "openai-compatibility" providerName = "openai-compatibility"
} }
base := strings.TrimSpace(compat.BaseURL) base := strings.TrimSpace(compat.BaseURL)
disableCooling := compat.DisableCooling
// Handle new APIKeyEntries format (preferred) // Handle new APIKeyEntries format (preferred)
createdEntries := 0 createdEntries := 0
@@ -218,6 +243,10 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
"compat_name": compat.Name, "compat_name": compat.Name,
"provider_key": providerName, "provider_key": providerName,
} }
metadata := map[string]any{}
if disableCooling {
metadata["disable_cooling"] = true
}
if compat.Priority != 0 { if compat.Priority != 0 {
attrs["priority"] = strconv.Itoa(compat.Priority) attrs["priority"] = strconv.Itoa(compat.Priority)
} }
@@ -236,9 +265,13 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
Status: coreauth.StatusActive, Status: coreauth.StatusActive,
ProxyURL: proxyURL, ProxyURL: proxyURL,
Attributes: attrs, Attributes: attrs,
Metadata: metadata,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
if len(a.Metadata) == 0 {
a.Metadata = nil
}
out = append(out, a) out = append(out, a)
createdEntries++ createdEntries++
} }
@@ -252,6 +285,10 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
"compat_name": compat.Name, "compat_name": compat.Name,
"provider_key": providerName, "provider_key": providerName,
} }
metadata := map[string]any{}
if disableCooling {
metadata["disable_cooling"] = true
}
if compat.Priority != 0 { if compat.Priority != 0 {
attrs["priority"] = strconv.Itoa(compat.Priority) attrs["priority"] = strconv.Itoa(compat.Priority)
} }
@@ -266,9 +303,13 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
Prefix: prefix, Prefix: prefix,
Status: coreauth.StatusActive, Status: coreauth.StatusActive,
Attributes: attrs, Attributes: attrs,
Metadata: metadata,
CreatedAt: now, CreatedAt: now,
UpdatedAt: now, UpdatedAt: now,
} }
if len(a.Metadata) == 0 {
a.Metadata = nil
}
out = append(out, a) out = append(out, a)
} }
} }
+41 -10
View File
@@ -68,11 +68,26 @@ func TestConfigSynthesizer_GeminiKeys(t *testing.T) {
if auths[0].Attributes["api_key"] != "test-key-123" { if auths[0].Attributes["api_key"] != "test-key-123" {
t.Errorf("expected api_key test-key-123, got %s", auths[0].Attributes["api_key"]) t.Errorf("expected api_key test-key-123, got %s", auths[0].Attributes["api_key"])
} }
if auths[0].Metadata != nil {
t.Errorf("expected metadata to be nil when disable_cooling not set, got %v", auths[0].Metadata)
}
if auths[0].Status != coreauth.StatusActive { if auths[0].Status != coreauth.StatusActive {
t.Errorf("expected status active, got %s", auths[0].Status) t.Errorf("expected status active, got %s", auths[0].Status)
} }
}, },
}, },
{
name: "gemini key disable cooling",
geminiKeys: []config.GeminiKey{
{APIKey: "test-key-123", Prefix: "team-a", DisableCooling: true},
},
wantLen: 1,
validate: func(t *testing.T, auths []*coreauth.Auth) {
if v, ok := auths[0].Metadata["disable_cooling"].(bool); !ok || !v {
t.Errorf("expected disable_cooling=true, got %v", auths[0].Metadata["disable_cooling"])
}
},
},
{ {
name: "gemini key with base url and proxy", name: "gemini key with base url and proxy",
geminiKeys: []config.GeminiKey{ geminiKeys: []config.GeminiKey{
@@ -160,9 +175,10 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) {
Config: &config.Config{ Config: &config.Config{
ClaudeKey: []config.ClaudeKey{ ClaudeKey: []config.ClaudeKey{
{ {
APIKey: "sk-ant-api-xxx", APIKey: "sk-ant-api-xxx",
Prefix: "main", Prefix: "main",
BaseURL: "https://api.anthropic.com", BaseURL: "https://api.anthropic.com",
DisableCooling: true,
Models: []config.ClaudeModel{ Models: []config.ClaudeModel{
{Name: "claude-3-opus"}, {Name: "claude-3-opus"},
{Name: "claude-3-sonnet"}, {Name: "claude-3-sonnet"},
@@ -197,6 +213,9 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) {
if _, ok := auths[0].Attributes["models_hash"]; !ok { if _, ok := auths[0].Attributes["models_hash"]; !ok {
t.Error("expected models_hash in attributes") t.Error("expected models_hash in attributes")
} }
if v, ok := auths[0].Metadata["disable_cooling"].(bool); !ok || !v {
t.Errorf("expected disable_cooling=true, got %v", auths[0].Metadata["disable_cooling"])
}
} }
func TestConfigSynthesizer_ClaudeKeys_SkipsEmptyAndHeaders(t *testing.T) { func TestConfigSynthesizer_ClaudeKeys_SkipsEmptyAndHeaders(t *testing.T) {
@@ -231,11 +250,12 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) {
Config: &config.Config{ Config: &config.Config{
CodexKey: []config.CodexKey{ CodexKey: []config.CodexKey{
{ {
APIKey: "codex-key-123", APIKey: "codex-key-123",
Prefix: "dev", Prefix: "dev",
BaseURL: "https://api.openai.com", BaseURL: "https://api.openai.com",
ProxyURL: "http://proxy.local", ProxyURL: "http://proxy.local",
Websockets: true, Websockets: true,
DisableCooling: true,
}, },
}, },
}, },
@@ -263,6 +283,9 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) {
if auths[0].Attributes["websockets"] != "true" { if auths[0].Attributes["websockets"] != "true" {
t.Errorf("expected websockets=true, got %s", auths[0].Attributes["websockets"]) t.Errorf("expected websockets=true, got %s", auths[0].Attributes["websockets"])
} }
if v, ok := auths[0].Metadata["disable_cooling"].(bool); !ok || !v {
t.Errorf("expected disable_cooling=true, got %v", auths[0].Metadata["disable_cooling"])
}
} }
func TestConfigSynthesizer_CodexKeys_SkipsEmptyAndHeaders(t *testing.T) { func TestConfigSynthesizer_CodexKeys_SkipsEmptyAndHeaders(t *testing.T) {
@@ -301,8 +324,9 @@ func TestConfigSynthesizer_OpenAICompat(t *testing.T) {
name: "with APIKeyEntries", name: "with APIKeyEntries",
compat: []config.OpenAICompatibility{ compat: []config.OpenAICompatibility{
{ {
Name: "CustomProvider", Name: "CustomProvider",
BaseURL: "https://custom.api.com", BaseURL: "https://custom.api.com",
DisableCooling: true,
APIKeyEntries: []config.OpenAICompatibilityAPIKey{ APIKeyEntries: []config.OpenAICompatibilityAPIKey{
{APIKey: "key-1"}, {APIKey: "key-1"},
{APIKey: "key-2"}, {APIKey: "key-2"},
@@ -365,6 +389,13 @@ func TestConfigSynthesizer_OpenAICompat(t *testing.T) {
if len(auths) != tt.wantLen { if len(auths) != tt.wantLen {
t.Fatalf("expected %d auths, got %d", tt.wantLen, len(auths)) t.Fatalf("expected %d auths, got %d", tt.wantLen, len(auths))
} }
if tt.name == "with APIKeyEntries" {
for i := range auths {
if v, ok := auths[i].Metadata["disable_cooling"].(bool); !ok || !v {
t.Fatalf("expected auth[%d].disable_cooling=true, got %v", i, auths[i].Metadata["disable_cooling"])
}
}
}
}) })
} }
} }
+10 -1
View File
@@ -355,19 +355,28 @@ func (a *Auth) ProxyInfo() string {
return "via proxy" return "via proxy"
} }
// DisableCoolingOverride returns the auth-file scoped disable_cooling override when present. // DisableCoolingOverride returns the auth scoped disable_cooling override when present.
// The value is read from metadata key "disable_cooling" (or legacy "disable-cooling"). // The value is read from metadata key "disable_cooling" (or legacy "disable-cooling").
//
// NOTE: This override is intentionally "true-only". When the metadata value is false, it is treated
// as "not set" so the global disable-cooling flag can still take effect.
func (a *Auth) DisableCoolingOverride() (bool, bool) { func (a *Auth) DisableCoolingOverride() (bool, bool) {
if a == nil || a.Metadata == nil { if a == nil || a.Metadata == nil {
return false, false return false, false
} }
if val, ok := a.Metadata["disable_cooling"]; ok { if val, ok := a.Metadata["disable_cooling"]; ok {
if parsed, okParse := parseBoolAny(val); okParse { if parsed, okParse := parseBoolAny(val); okParse {
if !parsed {
return false, false
}
return parsed, true return parsed, true
} }
} }
if val, ok := a.Metadata["disable-cooling"]; ok { if val, ok := a.Metadata["disable-cooling"]; ok {
if parsed, okParse := parseBoolAny(val); okParse { if parsed, okParse := parseBoolAny(val); okParse {
if !parsed {
return false, false
}
return parsed, true return parsed, true
} }
} }