Merge pull request #1386 from shenshuoyaoyouguang/sync-auth-changes

fix(auth): normalize model key for thinking suffix in selectors

sdk/cliproxy/auth/conductor.go

@@ -1299,7 +1299,7 @@ func updateAggregatedAvailability(auth *Auth, now time.Time) {
 			stateUnavailable = true
 		} else if state.Unavailable {
 			if state.NextRetryAfter.IsZero() {
-				stateUnavailable = true
+				stateUnavailable = false
 			} else if state.NextRetryAfter.After(now) {
 				stateUnavailable = true
 				if earliestRetry.IsZero() || state.NextRetryAfter.Before(earliestRetry) {

sdk/cliproxy/auth/conductor_availability_test.go

@@ -0,0 +1,62 @@
+package auth
+
+import (
+	"testing"
+	"time"
+)
+
+func TestUpdateAggregatedAvailability_UnavailableWithoutNextRetryDoesNotBlockAuth(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	model := "test-model"
+	auth := &Auth{
+		ID: "a",
+		ModelStates: map[string]*ModelState{
+			model: {
+				Status:      StatusError,
+				Unavailable: true,
+			},
+		},
+	}
+
+	updateAggregatedAvailability(auth, now)
+
+	if auth.Unavailable {
+		t.Fatalf("auth.Unavailable = true, want false")
+	}
+	if !auth.NextRetryAfter.IsZero() {
+		t.Fatalf("auth.NextRetryAfter = %v, want zero", auth.NextRetryAfter)
+	}
+}
+
+func TestUpdateAggregatedAvailability_FutureNextRetryBlocksAuth(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	model := "test-model"
+	next := now.Add(5 * time.Minute)
+	auth := &Auth{
+		ID: "a",
+		ModelStates: map[string]*ModelState{
+			model: {
+				Status:         StatusError,
+				Unavailable:    true,
+				NextRetryAfter: next,
+			},
+		},
+	}
+
+	updateAggregatedAvailability(auth, now)
+
+	if !auth.Unavailable {
+		t.Fatalf("auth.Unavailable = false, want true")
+	}
+	if auth.NextRetryAfter.IsZero() {
+		t.Fatalf("auth.NextRetryAfter = zero, want %v", next)
+	}
+	if auth.NextRetryAfter.Sub(next) > time.Second || next.Sub(auth.NextRetryAfter) > time.Second {
+		t.Fatalf("auth.NextRetryAfter = %v, want %v", auth.NextRetryAfter, next)
+	}
+}
+

sdk/cliproxy/auth/selector.go

@@ -12,6 +12,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
 	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
 )
 
@@ -19,6 +20,7 @@ import (
 type RoundRobinSelector struct {
 	mu      sync.Mutex
 	cursors map[string]int
+	maxKeys int
 }
 
 // FillFirstSelector selects the first available credential (deterministic ordering).
@@ -119,6 +121,19 @@ func authPriority(auth *Auth) int {
 	return parsed
 }
 
+func canonicalModelKey(model string) string {
+	model = strings.TrimSpace(model)
+	if model == "" {
+		return ""
+	}
+	parsed := thinking.ParseSuffix(model)
+	modelName := strings.TrimSpace(parsed.ModelName)
+	if modelName == "" {
+		return model
+	}
+	return modelName
+}
+
 func collectAvailableByPriority(auths []*Auth, model string, now time.Time) (available map[int][]*Auth, cooldownCount int, earliest time.Time) {
 	available = make(map[int][]*Auth)
 	for i := 0; i < len(auths); i++ {
@@ -185,11 +200,18 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
 	if err != nil {
 		return nil, err
 	}
-	key := provider + ":" + model
+	key := provider + ":" + canonicalModelKey(model)
 	s.mu.Lock()
 	if s.cursors == nil {
 		s.cursors = make(map[string]int)
 	}
+	limit := s.maxKeys
+	if limit <= 0 {
+		limit = 4096
+	}
+	if _, ok := s.cursors[key]; !ok && len(s.cursors) >= limit {
+		s.cursors = make(map[string]int)
+	}
 	index := s.cursors[key]
 
 	if index >= 2_147_483_640 {
@@ -223,7 +245,14 @@ func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, block
 	}
 	if model != "" {
 		if len(auth.ModelStates) > 0 {
-			if state, ok := auth.ModelStates[model]; ok && state != nil {
+			state, ok := auth.ModelStates[model]
+			if (!ok || state == nil) && model != "" {
+				baseModel := canonicalModelKey(model)
+				if baseModel != "" && baseModel != model {
+					state, ok = auth.ModelStates[baseModel]
+				}
+			}
+			if ok && state != nil {
 				if state.Status == StatusDisabled {
 					return true, blockReasonDisabled, time.Time{}
 				}

sdk/cliproxy/auth/selector_test.go

@@ -2,7 +2,9 @@ package auth
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
+	"net/http"
 	"sync"
 	"testing"
 	"time"
@@ -175,3 +177,228 @@ func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
 	default:
 	}
 }
+
+func TestSelectorPick_AllCooldownReturnsModelCooldownError(t *testing.T) {
+	t.Parallel()
+
+	model := "test-model"
+	now := time.Now()
+	next := now.Add(60 * time.Second)
+	auths := []*Auth{
+		{
+			ID: "a",
+			ModelStates: map[string]*ModelState{
+				model: {
+					Status:         StatusActive,
+					Unavailable:    true,
+					NextRetryAfter: next,
+					Quota: QuotaState{
+						Exceeded:      true,
+						NextRecoverAt: next,
+					},
+				},
+			},
+		},
+		{
+			ID: "b",
+			ModelStates: map[string]*ModelState{
+				model: {
+					Status:         StatusActive,
+					Unavailable:    true,
+					NextRetryAfter: next,
+					Quota: QuotaState{
+						Exceeded:      true,
+						NextRecoverAt: next,
+					},
+				},
+			},
+		},
+	}
+
+	t.Run("mixed provider redacts provider field", func(t *testing.T) {
+		t.Parallel()
+
+		selector := &FillFirstSelector{}
+		_, err := selector.Pick(context.Background(), "mixed", model, cliproxyexecutor.Options{}, auths)
+		if err == nil {
+			t.Fatalf("Pick() error = nil")
+		}
+
+		var mce *modelCooldownError
+		if !errors.As(err, &mce) {
+			t.Fatalf("Pick() error = %T, want *modelCooldownError", err)
+		}
+		if mce.StatusCode() != http.StatusTooManyRequests {
+			t.Fatalf("StatusCode() = %d, want %d", mce.StatusCode(), http.StatusTooManyRequests)
+		}
+
+		headers := mce.Headers()
+		if got := headers.Get("Retry-After"); got == "" {
+			t.Fatalf("Headers().Get(Retry-After) = empty")
+		}
+
+		var payload map[string]any
+		if err := json.Unmarshal([]byte(mce.Error()), &payload); err != nil {
+			t.Fatalf("json.Unmarshal(Error()) error = %v", err)
+		}
+		rawErr, ok := payload["error"].(map[string]any)
+		if !ok {
+			t.Fatalf("Error() payload missing error object: %v", payload)
+		}
+		if got, _ := rawErr["code"].(string); got != "model_cooldown" {
+			t.Fatalf("Error().error.code = %q, want %q", got, "model_cooldown")
+		}
+		if _, ok := rawErr["provider"]; ok {
+			t.Fatalf("Error().error.provider exists for mixed provider: %v", rawErr["provider"])
+		}
+	})
+
+	t.Run("non-mixed provider includes provider field", func(t *testing.T) {
+		t.Parallel()
+
+		selector := &FillFirstSelector{}
+		_, err := selector.Pick(context.Background(), "gemini", model, cliproxyexecutor.Options{}, auths)
+		if err == nil {
+			t.Fatalf("Pick() error = nil")
+		}
+
+		var mce *modelCooldownError
+		if !errors.As(err, &mce) {
+			t.Fatalf("Pick() error = %T, want *modelCooldownError", err)
+		}
+
+		var payload map[string]any
+		if err := json.Unmarshal([]byte(mce.Error()), &payload); err != nil {
+			t.Fatalf("json.Unmarshal(Error()) error = %v", err)
+		}
+		rawErr, ok := payload["error"].(map[string]any)
+		if !ok {
+			t.Fatalf("Error() payload missing error object: %v", payload)
+		}
+		if got, _ := rawErr["provider"].(string); got != "gemini" {
+			t.Fatalf("Error().error.provider = %q, want %q", got, "gemini")
+		}
+	})
+}
+
+func TestIsAuthBlockedForModel_UnavailableWithoutNextRetryIsNotBlocked(t *testing.T) {
+	t.Parallel()
+
+	now := time.Now()
+	model := "test-model"
+	auth := &Auth{
+		ID: "a",
+		ModelStates: map[string]*ModelState{
+			model: {
+				Status:      StatusActive,
+				Unavailable: true,
+				Quota: QuotaState{
+					Exceeded: true,
+				},
+			},
+		},
+	}
+
+	blocked, reason, next := isAuthBlockedForModel(auth, model, now)
+	if blocked {
+		t.Fatalf("blocked = true, want false")
+	}
+	if reason != blockReasonNone {
+		t.Fatalf("reason = %v, want %v", reason, blockReasonNone)
+	}
+	if !next.IsZero() {
+		t.Fatalf("next = %v, want zero", next)
+	}
+}
+
+func TestFillFirstSelectorPick_ThinkingSuffixFallsBackToBaseModelState(t *testing.T) {
+	t.Parallel()
+
+	selector := &FillFirstSelector{}
+	now := time.Now()
+
+	baseModel := "test-model"
+	requestedModel := "test-model(high)"
+
+	high := &Auth{
+		ID:         "high",
+		Attributes: map[string]string{"priority": "10"},
+		ModelStates: map[string]*ModelState{
+			baseModel: {
+				Status:         StatusActive,
+				Unavailable:    true,
+				NextRetryAfter: now.Add(30 * time.Minute),
+				Quota: QuotaState{
+					Exceeded: true,
+				},
+			},
+		},
+	}
+	low := &Auth{
+		ID:         "low",
+		Attributes: map[string]string{"priority": "0"},
+	}
+
+	got, err := selector.Pick(context.Background(), "mixed", requestedModel, cliproxyexecutor.Options{}, []*Auth{high, low})
+	if err != nil {
+		t.Fatalf("Pick() error = %v", err)
+	}
+	if got == nil {
+		t.Fatalf("Pick() auth = nil")
+	}
+	if got.ID != "low" {
+		t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "low")
+	}
+}
+
+func TestRoundRobinSelectorPick_ThinkingSuffixSharesCursor(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{}
+	auths := []*Auth{
+		{ID: "b"},
+		{ID: "a"},
+	}
+
+	first, err := selector.Pick(context.Background(), "gemini", "test-model(high)", cliproxyexecutor.Options{}, auths)
+	if err != nil {
+		t.Fatalf("Pick() first error = %v", err)
+	}
+	second, err := selector.Pick(context.Background(), "gemini", "test-model(low)", cliproxyexecutor.Options{}, auths)
+	if err != nil {
+		t.Fatalf("Pick() second error = %v", err)
+	}
+	if first == nil || second == nil {
+		t.Fatalf("Pick() returned nil auth")
+	}
+	if first.ID != "a" {
+		t.Fatalf("Pick() first auth.ID = %q, want %q", first.ID, "a")
+	}
+	if second.ID != "b" {
+		t.Fatalf("Pick() second auth.ID = %q, want %q", second.ID, "b")
+	}
+}
+
+func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) {
+	t.Parallel()
+
+	selector := &RoundRobinSelector{maxKeys: 2}
+	auths := []*Auth{{ID: "a"}}
+
+	_, _ = selector.Pick(context.Background(), "gemini", "m1", cliproxyexecutor.Options{}, auths)
+	_, _ = selector.Pick(context.Background(), "gemini", "m2", cliproxyexecutor.Options{}, auths)
+	_, _ = selector.Pick(context.Background(), "gemini", "m3", cliproxyexecutor.Options{}, auths)
+
+	selector.mu.Lock()
+	defer selector.mu.Unlock()
+
+	if selector.cursors == nil {
+		t.Fatalf("selector.cursors = nil")
+	}
+	if len(selector.cursors) != 1 {
+		t.Fatalf("len(selector.cursors) = %d, want %d", len(selector.cursors), 1)
+	}
+	if _, ok := selector.cursors["gemini:m3"]; !ok {
+		t.Fatalf("selector.cursors missing key %q", "gemini:m3")
+	}
+}