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
+41 -10
View File
@@ -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"])
}
}
}
})
}
}