fix(provider): preserve multikey cooldown isolation

This commit is contained in:
Alix-007 2026-03-19 11:08:39 +08:00
parent 7577f952cd
commit 01d907c9c5
2 changed files with 50 additions and 10 deletions

View file

@ -3,6 +3,7 @@ package providers
import ( import (
"context" "context"
"fmt" "fmt"
"strconv"
"strings" "strings"
"time" "time"
) )
@ -42,7 +43,7 @@ func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
return &FallbackChain{cooldown: cooldown} return &FallbackChain{cooldown: cooldown}
} }
func (c FallbackCandidate) cooldownKey() string { func (c FallbackCandidate) cooldownKey(candidates []FallbackCandidate) string {
if strings.TrimSpace(c.CooldownKey) != "" { if strings.TrimSpace(c.CooldownKey) != "" {
return c.CooldownKey return c.CooldownKey
} }
@ -52,9 +53,49 @@ func (c FallbackCandidate) cooldownKey() string {
} }
return ModelKey(c.Provider, c.Model) return ModelKey(c.Provider, c.Model)
} }
if belongsToMultiKeySet(c, candidates) {
return ModelKey(c.Provider, c.Model)
}
return c.Provider return c.Provider
} }
func belongsToMultiKeySet(candidate FallbackCandidate, candidates []FallbackCandidate) bool {
provider := NormalizeProvider(candidate.Provider)
model := strings.ToLower(strings.TrimSpace(candidate.Model))
if provider == "" || model == "" {
return false
}
baseModel, isReplica := multiKeyBaseModel(model)
if isReplica {
return true
}
for _, other := range candidates {
if NormalizeProvider(other.Provider) != provider {
continue
}
otherBase, otherIsReplica := multiKeyBaseModel(other.Model)
if otherIsReplica && otherBase == baseModel {
return true
}
}
return false
}
func multiKeyBaseModel(model string) (string, bool) {
normalized := strings.ToLower(strings.TrimSpace(model))
idx := strings.LastIndex(normalized, "__key_")
if idx <= 0 {
return normalized, false
}
if _, err := strconv.Atoi(normalized[idx+len("__key_"):]); err != nil {
return normalized, false
}
return normalized[:idx], true
}
// ResolveCandidates parses model config into a deduplicated candidate list. // ResolveCandidates parses model config into a deduplicated candidate list.
func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate { func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate {
return ResolveCandidatesWithLookup(cfg, defaultProvider, nil) return ResolveCandidatesWithLookup(cfg, defaultProvider, nil)
@ -126,7 +167,7 @@ func (fc *FallbackChain) Execute(
} }
for i, candidate := range candidates { for i, candidate := range candidates {
cooldownKey := candidate.cooldownKey() cooldownKey := candidate.cooldownKey(candidates)
// Check context before each attempt. // Check context before each attempt.
if ctx.Err() == context.Canceled { if ctx.Err() == context.Canceled {

View file

@ -165,8 +165,8 @@ func TestFallback_CooldownSkip(t *testing.T) {
ct, _ := newTestTracker(now) ct, _ := newTestTracker(now)
fc := NewFallbackChain(ct) fc := NewFallbackChain(ct)
// Put openai/gpt-4 in cooldown (using ModelKey now) // Put openai in cooldown
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit) ct.MarkFailure("openai", FailoverRateLimit)
candidates := []FallbackCandidate{ candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"), makeCandidate("openai", "gpt-4"),
@ -258,9 +258,9 @@ func TestFallback_AllInCooldown(t *testing.T) {
ct := NewCooldownTracker() ct := NewCooldownTracker()
fc := NewFallbackChain(ct) fc := NewFallbackChain(ct)
// Put all models in cooldown (using ModelKey now) // Put all providers in cooldown
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit) ct.MarkFailure("openai", FailoverRateLimit)
ct.MarkFailure(ModelKey("anthropic", "claude"), FailoverBilling) ct.MarkFailure("anthropic", FailoverBilling)
candidates := []FallbackCandidate{ candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"), makeCandidate("openai", "gpt-4"),
@ -336,13 +336,12 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
fc := NewFallbackChain(ct) fc := NewFallbackChain(ct)
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")} candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
modelKey := ModelKey("openai", "gpt-4")
attempt := 0 attempt := 0
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
attempt++ attempt++
if attempt == 1 { if attempt == 1 {
ct.MarkFailure(modelKey, FailoverRateLimit) // simulate failure tracked elsewhere ct.MarkFailure("openai", FailoverRateLimit) // simulate failure tracked elsewhere
} }
return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil
} }
@ -351,7 +350,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
if !ct.IsAvailable(modelKey) { if !ct.IsAvailable("openai") {
t.Error("success should reset cooldown") t.Error("success should reset cooldown")
} }
} }