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:
@@ -60,6 +60,10 @@ func (s *ConfigSynthesizer) synthesizeGeminiKeys(ctx *SynthesisContext) []*corea
|
||||
"source": fmt.Sprintf("config:gemini[%s]", token),
|
||||
"api_key": key,
|
||||
}
|
||||
metadata := map[string]any{}
|
||||
if entry.DisableCooling {
|
||||
metadata["disable_cooling"] = true
|
||||
}
|
||||
if entry.Priority != 0 {
|
||||
attrs["priority"] = strconv.Itoa(entry.Priority)
|
||||
}
|
||||
@@ -78,10 +82,14 @@ func (s *ConfigSynthesizer) synthesizeGeminiKeys(ctx *SynthesisContext) []*corea
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
Metadata: metadata,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
ApplyAuthExcludedModelsMeta(a, cfg, entry.ExcludedModels, "apikey")
|
||||
if len(a.Metadata) == 0 {
|
||||
a.Metadata = nil
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
@@ -107,6 +115,10 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea
|
||||
"source": fmt.Sprintf("config:claude[%s]", token),
|
||||
"api_key": key,
|
||||
}
|
||||
metadata := map[string]any{}
|
||||
if ck.DisableCooling {
|
||||
metadata["disable_cooling"] = true
|
||||
}
|
||||
if ck.Priority != 0 {
|
||||
attrs["priority"] = strconv.Itoa(ck.Priority)
|
||||
}
|
||||
@@ -126,10 +138,14 @@ func (s *ConfigSynthesizer) synthesizeClaudeKeys(ctx *SynthesisContext) []*corea
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
Metadata: metadata,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
|
||||
if len(a.Metadata) == 0 {
|
||||
a.Metadata = nil
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
@@ -154,6 +170,10 @@ func (s *ConfigSynthesizer) synthesizeCodexKeys(ctx *SynthesisContext) []*coreau
|
||||
"source": fmt.Sprintf("config:codex[%s]", token),
|
||||
"api_key": key,
|
||||
}
|
||||
metadata := map[string]any{}
|
||||
if ck.DisableCooling {
|
||||
metadata["disable_cooling"] = true
|
||||
}
|
||||
if ck.Priority != 0 {
|
||||
attrs["priority"] = strconv.Itoa(ck.Priority)
|
||||
}
|
||||
@@ -176,10 +196,14 @@ func (s *ConfigSynthesizer) synthesizeCodexKeys(ctx *SynthesisContext) []*coreau
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
Metadata: metadata,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
ApplyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
|
||||
if len(a.Metadata) == 0 {
|
||||
a.Metadata = nil
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
@@ -203,6 +227,7 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
|
||||
providerName = "openai-compatibility"
|
||||
}
|
||||
base := strings.TrimSpace(compat.BaseURL)
|
||||
disableCooling := compat.DisableCooling
|
||||
|
||||
// Handle new APIKeyEntries format (preferred)
|
||||
createdEntries := 0
|
||||
@@ -218,6 +243,10 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
|
||||
"compat_name": compat.Name,
|
||||
"provider_key": providerName,
|
||||
}
|
||||
metadata := map[string]any{}
|
||||
if disableCooling {
|
||||
metadata["disable_cooling"] = true
|
||||
}
|
||||
if compat.Priority != 0 {
|
||||
attrs["priority"] = strconv.Itoa(compat.Priority)
|
||||
}
|
||||
@@ -236,9 +265,13 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
Metadata: metadata,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if len(a.Metadata) == 0 {
|
||||
a.Metadata = nil
|
||||
}
|
||||
out = append(out, a)
|
||||
createdEntries++
|
||||
}
|
||||
@@ -252,6 +285,10 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
|
||||
"compat_name": compat.Name,
|
||||
"provider_key": providerName,
|
||||
}
|
||||
metadata := map[string]any{}
|
||||
if disableCooling {
|
||||
metadata["disable_cooling"] = true
|
||||
}
|
||||
if compat.Priority != 0 {
|
||||
attrs["priority"] = strconv.Itoa(compat.Priority)
|
||||
}
|
||||
@@ -266,9 +303,13 @@ func (s *ConfigSynthesizer) synthesizeOpenAICompat(ctx *SynthesisContext) []*cor
|
||||
Prefix: prefix,
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: attrs,
|
||||
Metadata: metadata,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if len(a.Metadata) == 0 {
|
||||
a.Metadata = nil
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,11 +68,26 @@ func TestConfigSynthesizer_GeminiKeys(t *testing.T) {
|
||||
if auths[0].Attributes["api_key"] != "test-key-123" {
|
||||
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 {
|
||||
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",
|
||||
geminiKeys: []config.GeminiKey{
|
||||
@@ -160,9 +175,10 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) {
|
||||
Config: &config.Config{
|
||||
ClaudeKey: []config.ClaudeKey{
|
||||
{
|
||||
APIKey: "sk-ant-api-xxx",
|
||||
Prefix: "main",
|
||||
BaseURL: "https://api.anthropic.com",
|
||||
APIKey: "sk-ant-api-xxx",
|
||||
Prefix: "main",
|
||||
BaseURL: "https://api.anthropic.com",
|
||||
DisableCooling: true,
|
||||
Models: []config.ClaudeModel{
|
||||
{Name: "claude-3-opus"},
|
||||
{Name: "claude-3-sonnet"},
|
||||
@@ -197,6 +213,9 @@ func TestConfigSynthesizer_ClaudeKeys(t *testing.T) {
|
||||
if _, ok := auths[0].Attributes["models_hash"]; !ok {
|
||||
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) {
|
||||
@@ -231,11 +250,12 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) {
|
||||
Config: &config.Config{
|
||||
CodexKey: []config.CodexKey{
|
||||
{
|
||||
APIKey: "codex-key-123",
|
||||
Prefix: "dev",
|
||||
BaseURL: "https://api.openai.com",
|
||||
ProxyURL: "http://proxy.local",
|
||||
Websockets: true,
|
||||
APIKey: "codex-key-123",
|
||||
Prefix: "dev",
|
||||
BaseURL: "https://api.openai.com",
|
||||
ProxyURL: "http://proxy.local",
|
||||
Websockets: true,
|
||||
DisableCooling: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
@@ -263,6 +283,9 @@ func TestConfigSynthesizer_CodexKeys(t *testing.T) {
|
||||
if auths[0].Attributes["websockets"] != "true" {
|
||||
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) {
|
||||
@@ -301,8 +324,9 @@ func TestConfigSynthesizer_OpenAICompat(t *testing.T) {
|
||||
name: "with APIKeyEntries",
|
||||
compat: []config.OpenAICompatibility{
|
||||
{
|
||||
Name: "CustomProvider",
|
||||
BaseURL: "https://custom.api.com",
|
||||
Name: "CustomProvider",
|
||||
BaseURL: "https://custom.api.com",
|
||||
DisableCooling: true,
|
||||
APIKeyEntries: []config.OpenAICompatibilityAPIKey{
|
||||
{APIKey: "key-1"},
|
||||
{APIKey: "key-2"},
|
||||
@@ -365,6 +389,13 @@ func TestConfigSynthesizer_OpenAICompat(t *testing.T) {
|
||||
if len(auths) != tt.wantLen {
|
||||
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"])
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user