From 00c4377a21b8c6d5d495a499f1ab479e4460503f Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Sun, 16 Aug 2026 21:55:50 +0800
Subject: [PATCH] fix(auth): share session affinity across model variant
 suffixes

- Canonicalize model IDs when building session-affinity cache/fallback keys so variant suffixes (for example thinking modes) map to the same binding.
- Normalize model values from result metadata/on-result release paths to release and rebind bindings consistently across canonical model keys.

Closes: #5016
---
 sdk/cliproxy/auth/selector.go      |  9 ++--
 sdk/cliproxy/auth/selector_test.go | 69 ++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go
index accde458..7a053b78 100644
--- a/sdk/cliproxy/auth/selector.go
+++ b/sdk/cliproxy/auth/selector.go
@@ -684,10 +684,11 @@ func (s *SessionAffinitySelector) Pick(ctx context.Context, provider, model stri
 	}
 	fallbackAuths := highestPriorityAuths(available)
 
-	cacheKey := provider + "::" + primaryID + "::" + model
+	modelKey := canonicalModelKey(model)
+	cacheKey := provider + "::" + primaryID + "::" + modelKey
 	fallbackKey := ""
 	if fallbackID != "" && fallbackID != primaryID {
-		fallbackKey = provider + "::" + fallbackID + "::" + model
+		fallbackKey = provider + "::" + fallbackID + "::" + modelKey
 	}
 	bind := func(authID string) {
 		if fallbackKey != "" {
@@ -783,9 +784,9 @@ func (s *SessionAffinitySelector) OnResult(res Result) {
 	if raw, ok := res.Options.Metadata[cliproxyexecutor.SessionAffinityProviderMetadataKey].(string); ok && raw != "" {
 		ns = raw
 	}
-	nsModel := res.Model
+	nsModel := canonicalModelKey(res.Model)
 	if raw, ok := res.Options.Metadata[cliproxyexecutor.SessionAffinityModelMetadataKey].(string); ok && raw != "" {
-		nsModel = raw
+		nsModel = canonicalModelKey(raw)
 	}
 
 	cacheKey := ns + "::" + primaryID + "::" + nsModel
diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go
index 19677b24..c8e1e5dc 100644
--- a/sdk/cliproxy/auth/selector_test.go
+++ b/sdk/cliproxy/auth/selector_test.go
@@ -779,6 +779,75 @@ func TestSessionAffinitySelector_SameSessionSameAuth(t *testing.T) {
 	}
 }
 
+func TestSessionAffinitySelector_ThinkingSuffixVariantsPreserveBindingAndRelease(t *testing.T) {
+	t.Parallel()
+
+	fallback := &RoundRobinSelector{}
+	selector := NewSessionAffinitySelector(fallback)
+	defer selector.Stop()
+
+	auths := []*Auth{
+		{ID: "auth-a"},
+		{ID: "auth-b"},
+		{ID: "auth-c"},
+	}
+
+	payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`)
+	opts := cliproxyexecutor.Options{OriginalRequest: payload}
+
+	first, errFirst := selector.Pick(context.Background(), "anthropic", "claude-sonnet-4-5", opts, auths)
+	if errFirst != nil {
+		t.Fatalf("first Pick() error = %v", errFirst)
+	}
+	if first == nil {
+		t.Fatalf("first Pick() returned nil")
+	}
+
+	// Suffix variant claude-sonnet-4-5(high) should reuse the exact same auth binding
+	second, errSecond := selector.Pick(context.Background(), "anthropic", "claude-sonnet-4-5(high)", opts, auths)
+	if errSecond != nil {
+		t.Fatalf("second Pick() error = %v", errSecond)
+	}
+	if second.ID != first.ID {
+		t.Fatalf("second Pick() auth.ID = %q, want %q (thinking suffix variant should keep session stickiness)", second.ID, first.ID)
+	}
+
+	// Third request with claude-sonnet-4-5(medium) should also reuse the same auth
+	third, errThird := selector.Pick(context.Background(), "anthropic", "claude-sonnet-4-5(medium)", opts, auths)
+	if errThird != nil {
+		t.Fatalf("third Pick() error = %v", errThird)
+	}
+	if third.ID != first.ID {
+		t.Fatalf("third Pick() auth.ID = %q, want %q (thinking suffix variant should keep session stickiness)", third.ID, first.ID)
+	}
+
+	// Failure on a thinking-suffix variant (with explicit metadata) should properly release the session binding
+	optsWithMetadata := cliproxyexecutor.Options{
+		OriginalRequest: payload,
+		Metadata: map[string]any{
+			cliproxyexecutor.SessionAffinityProviderMetadataKey: "anthropic",
+			cliproxyexecutor.SessionAffinityModelMetadataKey:    "claude-sonnet-4-5(high)",
+		},
+	}
+	selector.OnResult(Result{
+		Provider: "anthropic",
+		Model:    "claude-sonnet-4-5(high)",
+		AuthID:   first.ID,
+		Success:  false,
+		Error:    &Error{Code: "rate_limited", Message: "rate limited"},
+		Options:  optsWithMetadata,
+	})
+
+	// After release, next pick should reselect using fallback selector
+	next, errNext := selector.Pick(context.Background(), "anthropic", "claude-sonnet-4-5", opts, auths)
+	if errNext != nil {
+		t.Fatalf("next Pick() error = %v", errNext)
+	}
+	if next.ID == first.ID {
+		t.Fatalf("next Pick() auth.ID = %q, should have reselected a different auth after failure release", next.ID)
+	}
+}
+
 func TestSessionAffinitySelector_WeightedBindingRebindsAfterWeightBecomesZero(t *testing.T) {
 	t.Parallel()
 
-- 
2.51.2