From 01d907c9c5420aa7670aa6cd0f1a26432d497161 Mon Sep 17 00:00:00 2001
From: Alix-007 <267018309+Alix-007@users.noreply.github.com>
Date: Thu, 19 Mar 2026 11:08:39 +0800
Subject: [PATCH] fix(provider): preserve multikey cooldown isolation

---
 pkg/providers/fallback.go      | 45 ++++++++++++++++++++++++++++++++--
 pkg/providers/fallback_test.go | 15 ++++++------
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go
index 8b3b94d94..319d72063 100644
--- a/pkg/providers/fallback.go
+++ b/pkg/providers/fallback.go
@@ -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 {
diff --git a/pkg/providers/fallback_test.go b/pkg/providers/fallback_test.go
index 875b0745d..fb4c26591 100644
--- a/pkg/providers/fallback_test.go
+++ b/pkg/providers/fallback_test.go
@@ -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")
 	}
 }