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 (
"context"
"fmt"
"strconv"
"strings"
"time"
)
@ -42,7 +43,7 @@ func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
return &FallbackChain{cooldown: cooldown}
}
func (c FallbackCandidate) cooldownKey() string {
func (c FallbackCandidate) cooldownKey(candidates []FallbackCandidate) string {
if strings.TrimSpace(c.CooldownKey) != "" {
return c.CooldownKey
}
@ -52,9 +53,49 @@ func (c FallbackCandidate) cooldownKey() string {
}
return ModelKey(c.Provider, c.Model)
}
if belongsToMultiKeySet(c, candidates) {
return ModelKey(c.Provider, c.Model)
}
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.
func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate {
return ResolveCandidatesWithLookup(cfg, defaultProvider, nil)
@ -126,7 +167,7 @@ func (fc *FallbackChain) Execute(
}
for i, candidate := range candidates {
cooldownKey := candidate.cooldownKey()
cooldownKey := candidate.cooldownKey(candidates)
// Check context before each attempt.
if ctx.Err() == context.Canceled {

View file

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