package auth import ( "context" "encoding/json" "errors" "fmt" "math" "net/http" "strings" "sync" "testing" "time" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" cliproxysession "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/session" sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator" ) func TestFillFirstSelectorPick_Deterministic(t *testing.T) { t.Parallel() selector := &FillFirstSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } if got == nil { t.Fatalf("Pick() auth = nil") } if got.ID != "a" { t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a") } } func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) { t.Parallel() selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } want := []string{"a", "b", "c", "a", "b"} for i, id := range want { got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() #%d error = %v", i, err) } if got == nil { t.Fatalf("Pick() #%d auth = nil", i) } if got.ID != id { t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id) } } } func TestWeightedRoundRobinSelectorPick_DistributesAndSkipsNonPositiveWeights(t *testing.T) { t.Parallel() selector := &WeightedRoundRobinSelector{} auths := []*Auth{ {ID: "a", Attributes: map[string]string{AttributeWeight: "5"}}, {ID: "b", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "c", Attributes: map[string]string{AttributeWeight: "2"}}, {ID: "disabled-by-weight", Attributes: map[string]string{AttributeWeight: "0"}}, } counts := make(map[string]int) for index := 0; index < 100; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil { t.Fatalf("Pick() #%d error = %v", index, errPick) } counts[got.ID]++ } want := map[string]int{"a": 50, "b": 30, "c": 20} for authID, wantCount := range want { if counts[authID] != wantCount { t.Fatalf("auth %q picks = %d, want %d", authID, counts[authID], wantCount) } } if counts["disabled-by-weight"] != 0 { t.Fatalf("non-positive weight auth picks = %d, want 0", counts["disabled-by-weight"]) } } func TestWeightedRoundRobinSelectorPick_ResetsCreditsWhenWeightsChange(t *testing.T) { t.Parallel() selector := &WeightedRoundRobinSelector{} authA := &Auth{ID: "a", Attributes: map[string]string{AttributeWeight: "1000000"}} authB := &Auth{ID: "b", Attributes: map[string]string{AttributeWeight: "1"}} auths := []*Auth{authA, authB} for index := 0; index < 1000; index++ { if _, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths); errPick != nil { t.Fatalf("warmup Pick() #%d error = %v", index, errPick) } } authA.Attributes[AttributeWeight] = "1" counts := make(map[string]int) for index := 0; index < 20; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil { t.Fatalf("Pick() after weight change #%d error = %v", index, errPick) } counts[got.ID]++ } if counts["a"] != 10 || counts["b"] != 10 { t.Fatalf("picks after weight change = %#v, want a:b=10:10", counts) } } func TestWeightedRoundRobinSelectorPick_RebalancesWhenHighestWeightUnavailable(t *testing.T) { t.Parallel() selector := &WeightedRoundRobinSelector{} auths := []*Auth{ {ID: "a", Disabled: true, Attributes: map[string]string{AttributeWeight: "5"}}, {ID: "b", Attributes: map[string]string{AttributeWeight: "3"}}, {ID: "c", Attributes: map[string]string{AttributeWeight: "2"}}, } counts := make(map[string]int) for index := 0; index < 100; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil { t.Fatalf("Pick() #%d error = %v", index, errPick) } counts[got.ID]++ } if counts["a"] != 0 || counts["b"] != 60 || counts["c"] != 40 { t.Fatalf("weighted failover counts = %#v, want b:c=60:40 with a skipped", counts) } } func TestWeightedRoundRobinSelectorPick_SkipsUnavailableAndQuotaExceededWithoutRecovery(t *testing.T) { t.Parallel() model := "test-model" selector := &WeightedRoundRobinSelector{} auths := []*Auth{ { ID: "model-unavailable", ModelStates: map[string]*ModelState{ model: {Unavailable: true}, }, }, {ID: "quota-exceeded", Quota: QuotaState{Exceeded: true}}, {ID: "available"}, } gotModel, errModel := selector.Pick(context.Background(), "gemini", model, cliproxyexecutor.Options{}, auths) if errModel != nil || gotModel == nil || gotModel.ID != "available" { t.Fatalf("model Pick() = %#v, %v; want available", gotModel, errModel) } for index := 0; index < 4; index++ { gotAuth, errAuth := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if errAuth != nil || gotAuth == nil { t.Fatalf("auth Pick() #%d = %#v, %v; want available auth", index, gotAuth, errAuth) } if gotAuth.ID == "quota-exceeded" { t.Fatalf("auth Pick() #%d selected quota-exceeded credential", index) } } } func TestAuthWeight_MetadataFallbackAndAttributePrecedence(t *testing.T) { t.Parallel() if got := authWeight(&Auth{Metadata: map[string]any{AttributeWeight: float64(7)}}); got != 7 { t.Fatalf("authWeight(metadata) = %d, want 7", got) } if got := authWeight(&Auth{ Attributes: map[string]string{AttributeWeight: "3"}, Metadata: map[string]any{AttributeWeight: float64(7)}, }); got != 3 { t.Fatalf("authWeight(attribute and metadata) = %d, want attribute weight 3", got) } } func TestAuthWeight_InvalidAndOverflowValuesAreExcluded(t *testing.T) { t.Parallel() for _, raw := range []string{"1.5", "1000001", "9223372036854775807", "9223372036854775808"} { auth := &Auth{Attributes: map[string]string{AttributeWeight: raw}} if got := authWeight(auth); got != 0 { t.Fatalf("authWeight(%q) = %d, want 0", raw, got) } } if got := authWeight(&Auth{Metadata: map[string]any{AttributeWeight: 1.5}}); got != 0 { t.Fatalf("authWeight(invalid metadata) = %d, want 0", got) } if got := authWeight(&Auth{Attributes: map[string]string{AttributeWeight: "-1"}}); got != 0 { t.Fatalf("authWeight(-1) = %d, want 0", got) } } func TestPickSmoothWeightedAuth_SaturatesCorruptState(t *testing.T) { t.Parallel() current := map[string]int64{"a": math.MaxInt64, "b": math.MinInt64} picked := pickSmoothWeightedAuth([]*Auth{{ID: "a"}, {ID: "b"}}, current) if picked == nil { t.Fatal("pickSmoothWeightedAuth() returned nil") } if current["a"] != math.MaxInt64-2 || current["b"] != math.MinInt64+1 { t.Fatalf("current state = %#v, want saturated arithmetic", current) } } func TestWeightedRoundRobinSelectorPick_RecoveredAuthReturnsWithoutAccumulatedCredit(t *testing.T) { t.Parallel() selector := &WeightedRoundRobinSelector{} authA := &Auth{ID: "a", Attributes: map[string]string{AttributeWeight: "5"}} authB := &Auth{ID: "b", Attributes: map[string]string{AttributeWeight: "1"}} auths := []*Auth{authA, authB} for index := 0; index < 6; index++ { if _, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths); errPick != nil { t.Fatalf("warmup Pick() #%d error = %v", index, errPick) } } authA.Unavailable = true authA.NextRetryAfter = time.Now().Add(time.Hour) for index := 0; index < 6; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil || got == nil || got.ID != "b" { t.Fatalf("unavailable Pick() #%d = %#v, %v; want b", index, got, errPick) } } authA.Unavailable = false authA.NextRetryAfter = time.Time{} counts := make(map[string]int) for index := 0; index < 6; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil { t.Fatalf("recovered Pick() #%d error = %v", index, errPick) } counts[got.ID]++ } if counts["a"] != 5 || counts["b"] != 1 { t.Fatalf("recovered picks = %#v, want a:b=5:1", counts) } } func TestWeightedRoundRobinSelectorPick_DefaultWeightIsOne(t *testing.T) { t.Parallel() selector := &WeightedRoundRobinSelector{} auths := []*Auth{{ID: "a"}, {ID: "b"}, {ID: "c"}} counts := make(map[string]int) for index := 0; index < 30; index++ { got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, auths) if errPick != nil { t.Fatalf("Pick() #%d error = %v", index, errPick) } counts[got.ID]++ } for _, authID := range []string{"a", "b", "c"} { if counts[authID] != 10 { t.Fatalf("auth %q picks = %d, want 10", authID, counts[authID]) } } } func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) { t.Parallel() selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "c", Attributes: map[string]string{"priority": "0"}}, {ID: "a", Attributes: map[string]string{"priority": "10"}}, {ID: "b", Attributes: map[string]string{"priority": "10"}}, } want := []string{"a", "b", "a", "b"} for i, id := range want { got, err := selector.Pick(context.Background(), "mixed", "", cliproxyexecutor.Options{}, auths) if err != nil { t.Fatalf("Pick() #%d error = %v", i, err) } if got == nil { t.Fatalf("Pick() #%d auth = nil", i) } if got.ID != id { t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id) } if got.ID == "c" { t.Fatalf("Pick() #%d unexpectedly selected lower priority auth", i) } } } func TestFillFirstSelectorPick_PriorityFallbackCooldown(t *testing.T) { t.Parallel() selector := &FillFirstSelector{} now := time.Now() model := "test-model" high := &Auth{ ID: "high", Attributes: map[string]string{"priority": "10"}, ModelStates: map[string]*ModelState{ model: { 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", model, 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_Concurrent(t *testing.T) { selector := &RoundRobinSelector{} auths := []*Auth{ {ID: "b"}, {ID: "a"}, {ID: "c"}, } start := make(chan struct{}) var wg sync.WaitGroup errCh := make(chan error, 1) goroutines := 32 iterations := 100 for i := 0; i < goroutines; i++ { wg.Add(1) go func() { defer wg.Done() <-start for j := 0; j < iterations; j++ { got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths) if err != nil { select { case errCh <- err: default: } return } if got == nil { select { case errCh <- errors.New("Pick() returned nil auth"): default: } return } if got.ID == "" { select { case errCh <- errors.New("Pick() returned auth with empty ID"): default: } return } } }() } close(start) wg.Wait() select { case err := <-errCh: t.Fatalf("concurrent Pick() error = %v", err) 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_UnavailableWithoutNextRetryIsBlocked(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 = false, want true") } if reason != blockReasonOther { t.Fatalf("reason = %v, want %v", reason, blockReasonOther) } if !next.IsZero() { t.Fatalf("next = %v, want zero", next) } } func TestIsAuthBlockedForModel_AuthQuotaExceededWithoutRecoveryIsBlocked(t *testing.T) { t.Parallel() auth := &Auth{ID: "a", Quota: QuotaState{Exceeded: true}} for _, model := range []string{"", "test-model"} { blocked, reason, next := isAuthBlockedForModel(auth, model, time.Now()) if !blocked || reason != blockReasonOther || !next.IsZero() { t.Fatalf("isAuthBlockedForModel(%q) = %v, %v, %v; want true, other, zero", model, blocked, reason, next) } } } func TestIsAuthBlockedForModel_ExpiredRecoveryIsAvailable(t *testing.T) { t.Parallel() now := time.Now() auth := &Auth{ ID: "a", Unavailable: true, NextRetryAfter: now.Add(-time.Minute), Quota: QuotaState{ Exceeded: true, NextRecoverAt: now.Add(-time.Second), }, } blocked, reason, next := isAuthBlockedForModel(auth, "", now) if blocked || reason != blockReasonNone || !next.IsZero() { t.Fatalf("isAuthBlockedForModel() = %v, %v, %v; want false, none, zero", blocked, reason, 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 TestIsAuthBlockedForModel_ThinkingSuffixStatesBlockCanonicalModel(t *testing.T) { t.Parallel() now := time.Now() laterRetry := now.Add(2 * time.Hour) auth := &Auth{ ID: "a", ModelStates: map[string]*ModelState{ "test-model(high)": { Status: StatusError, Unavailable: true, NextRetryAfter: now.Add(time.Hour), Quota: QuotaState{ Exceeded: true, NextRecoverAt: now.Add(time.Hour), }, }, "test-model(low)": { Status: StatusError, Unavailable: true, NextRetryAfter: laterRetry, Quota: QuotaState{ Exceeded: true, NextRecoverAt: laterRetry, }, }, }, } for _, model := range []string{"test-model", "test-model(medium)", "test-model(low)"} { blocked, reason, next := isAuthBlockedForModel(auth, model, now) if !blocked || reason != blockReasonCooldown || !next.Equal(laterRetry) { t.Fatalf("isAuthBlockedForModel(%q) = %v, %v, %v; want true, cooldown, %v", model, blocked, reason, next, laterRetry) } } } 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") } } func TestExtractSessionID(t *testing.T) { t.Parallel() tests := []struct { name string payload string want string }{ { name: "valid_claude_code_format", payload: `{"metadata":{"user_id":"user_3f221fe75652cf9a89a31647f16274bb8036a9b85ac4dc226a4df0efec8dc04d_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`, want: "claude:ac980658-63bd-4fb3-97ba-8da64cb1e344", }, { name: "json_user_id_with_session_id", payload: `{"metadata":{"user_id":"{\"device_id\":\"be82c3aee1e0c2d74535bacc85f9f559228f02dd8a17298cf522b71e6c375714\",\"account_uuid\":\"\",\"session_id\":\"e26d4046-0f88-4b09-bb5b-f863ab5fb24e\"}"}}`, want: "claude:e26d4046-0f88-4b09-bb5b-f863ab5fb24e", }, { name: "json_user_id_without_session_id", payload: `{"metadata":{"user_id":"{\"device_id\":\"abc123\"}"}}`, want: `user:{"device_id":"abc123"}`, }, { name: "no_session_but_user_id", payload: `{"metadata":{"user_id":"user_abc123"}}`, want: "user:user_abc123", }, { name: "conversation_id", payload: `{"conversation_id":"conv-12345"}`, want: "conv:conv-12345", }, { name: "no_metadata", payload: `{"model":"claude-3"}`, want: "", }, { name: "empty_payload", payload: ``, want: "", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := extractSessionID([]byte(tt.payload)) if got != tt.want { t.Errorf("extractSessionID() = %q, want %q", got, tt.want) } }) } } func TestSessionAffinitySelector_SameSessionSameAuth(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelector(fallback) auths := []*Auth{ {ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}, } // Use valid UUID format for session ID payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} // Same session should always pick the same auth first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } if first == nil { t.Fatalf("Pick() returned nil") } // Verify consistency: same session, same auths -> same result for i := 0; i < 10; i++ { got, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { t.Fatalf("Pick() #%d error = %v", i, err) } if got.ID != first.ID { t.Fatalf("Pick() #%d auth.ID = %q, want %q (same session should pick same auth)", i, got.ID, first.ID) } } } 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() selector := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) defer selector.Stop() authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "1"}} authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "1"}} auths := []*Auth{authA, authB} opts := cliproxyexecutor.Options{OriginalRequest: []byte(`{"metadata":{"user_id":"user_xxx_account__session_weight-change"}}`)} first, errFirst := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if errFirst != nil { t.Fatalf("first Pick() error = %v", errFirst) } if first.ID != authA.ID { t.Fatalf("first Pick() auth.ID = %q, want %q", first.ID, authA.ID) } authA.Attributes[AttributeWeight] = "0" second, errSecond := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if errSecond != nil { t.Fatalf("Pick() after weight update error = %v", errSecond) } if second.ID != authB.ID { t.Fatalf("Pick() after weight update auth.ID = %q, want %q", second.ID, authB.ID) } authA.Attributes[AttributeWeight] = "10" third, errThird := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if errThird != nil { t.Fatalf("Pick() after rebind error = %v", errThird) } if third.ID != authB.ID { t.Fatalf("Pick() after rebind auth.ID = %q, want sticky auth %q", third.ID, authB.ID) } } func TestSessionAffinitySelector_WeightedNewSessionsResetAfterWeightChange(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelector(&WeightedRoundRobinSelector{}) defer selector.Stop() authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "1000000"}} authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "1"}} auths := []*Auth{authA, authB} pickSession := func(index int) *Auth { t.Helper() opts := cliproxyexecutor.Options{OriginalRequest: []byte(fmt.Sprintf(`{"session_id":"session-%d"}`, index))} picked, errPick := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if errPick != nil { t.Fatalf("Pick(session-%d) error = %v", index, errPick) } return picked } for index := 0; index < 1000; index++ { pickSession(index) } authA.Attributes[AttributeWeight] = "1" counts := make(map[string]int) for index := 1000; index < 1020; index++ { counts[pickSession(index).ID]++ } if counts[authA.ID] != 10 || counts[authB.ID] != 10 { t.Fatalf("new session picks after weight change = %#v, want 10 each", counts) } } func TestSessionAffinitySelector_NoSessionFallback(t *testing.T) { t.Parallel() fallback := &FillFirstSelector{} selector := NewSessionAffinitySelector(fallback) auths := []*Auth{ {ID: "auth-b"}, {ID: "auth-a"}, {ID: "auth-c"}, } // No session in payload, should fallback to FillFirstSelector (picks "auth-a" after sorting) payload := []byte(`{"model":"claude-3"}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} got, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } if got.ID != "auth-a" { t.Fatalf("Pick() auth.ID = %q, want %q (should fallback to FillFirst)", got.ID, "auth-a") } } func TestSessionAffinitySelector_DifferentSessionsDifferentAuths(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelector(fallback) auths := []*Auth{ {ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}, } // Use valid UUID format for session IDs session1 := []byte(`{"metadata":{"user_id":"user_xxx_account__session_11111111-1111-1111-1111-111111111111"}}`) session2 := []byte(`{"metadata":{"user_id":"user_xxx_account__session_22222222-2222-2222-2222-222222222222"}}`) opts1 := cliproxyexecutor.Options{OriginalRequest: session1} opts2 := cliproxyexecutor.Options{OriginalRequest: session2} auth1, _ := selector.Pick(context.Background(), "claude", "claude-3", opts1, auths) auth2, _ := selector.Pick(context.Background(), "claude", "claude-3", opts2, auths) // Different sessions may or may not pick different auths (depends on hash collision) // But each session should be consistent for i := 0; i < 5; i++ { got1, _ := selector.Pick(context.Background(), "claude", "claude-3", opts1, auths) got2, _ := selector.Pick(context.Background(), "claude", "claude-3", opts2, auths) if got1.ID != auth1.ID { t.Fatalf("session1 Pick() #%d inconsistent: got %q, want %q", i, got1.ID, auth1.ID) } if got2.ID != auth2.ID { t.Fatalf("session2 Pick() #%d inconsistent: got %q, want %q", i, got2.ID, auth2.ID) } } } func TestSessionAffinitySelector_FailoverWhenAuthUnavailable(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{ {ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}, } payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_failover-test-uuid"}}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} // First pick establishes binding first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } // Remove the bound auth from available list (simulating rate limit) availableWithoutFirst := make([]*Auth, 0, len(auths)-1) for _, a := range auths { if a.ID != first.ID { availableWithoutFirst = append(availableWithoutFirst, a) } } // With failover enabled, should pick a new auth second, err := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) if err != nil { t.Fatalf("Pick() after failover error = %v", err) } if second.ID == first.ID { t.Fatalf("Pick() after failover returned same auth %q, expected different", first.ID) } // Subsequent picks should consistently return the new binding for i := 0; i < 5; i++ { got, _ := selector.Pick(context.Background(), "claude", "claude-3", opts, availableWithoutFirst) if got.ID != second.ID { t.Fatalf("Pick() #%d after failover inconsistent: got %q, want %q", i, got.ID, second.ID) } } } func TestExtractSessionID_ClaudeCodePriorityOverHeader(t *testing.T) { t.Parallel() // Claude Code metadata.user_id remains higher priority than a generic X-Session-ID header. headers := make(http.Header) headers.Set("X-Session-ID", "header-session-id") payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`) got := ExtractSessionID(headers, payload, nil) want := "claude:ac980658-63bd-4fb3-97ba-8da64cb1e344" if got != want { t.Errorf("ExtractSessionID() = %q, want %q (Claude Code should have highest priority over header)", got, want) } } func TestExtractSessionID_ClaudeCodePriorityOverIdempotencyKey(t *testing.T) { t.Parallel() // Claude Code metadata.user_id should have highest priority, even when idempotency_key is present metadata := map[string]any{"idempotency_key": "idem-12345"} payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_ac980658-63bd-4fb3-97ba-8da64cb1e344"}}`) got := ExtractSessionID(nil, payload, metadata) want := "claude:ac980658-63bd-4fb3-97ba-8da64cb1e344" if got != want { t.Errorf("ExtractSessionID() = %q, want %q (Claude Code should have highest priority over idempotency_key)", got, want) } } func TestExtractSessionID_Headers(t *testing.T) { t.Parallel() headers := make(http.Header) headers.Set("X-Session-ID", "my-explicit-session") got := ExtractSessionID(headers, nil, nil) want := "header:my-explicit-session" if got != want { t.Errorf("ExtractSessionID() with header = %q, want %q", got, want) } } func TestExtractSessionID_CodexSessionIDHeader(t *testing.T) { t.Parallel() headers := make(http.Header) headers.Set("Session_id", "codex-session-123") got := ExtractSessionID(headers, nil, nil) want := "codex:codex-session-123" if got != want { t.Errorf("ExtractSessionID() with Session_id = %q, want %q", got, want) } } func TestExtractSessionID_ClientRequestIDHeader(t *testing.T) { t.Parallel() headers := make(http.Header) headers.Set("X-Client-Request-Id", "pi-session-123") got := ExtractSessionID(headers, nil, nil) want := "clientreq:pi-session-123" if got != want { t.Errorf("ExtractSessionID() with X-Client-Request-Id = %q, want %q", got, want) } } func TestExtractSessionID_CodexSessionIDPriorityOverClientRequestID(t *testing.T) { t.Parallel() headers := make(http.Header) headers.Set("X-Client-Request-Id", "pi-session-123") headers.Set("Session_id", "codex-session-456") got := ExtractSessionID(headers, nil, nil) want := "codex:codex-session-456" if got != want { t.Errorf("ExtractSessionID() = %q, want %q (Session_id should take priority over X-Client-Request-Id)", got, want) } } // TestExtractSessionID_IdempotencyKey verifies that idempotency_key is intentionally // ignored for session affinity (it's auto-generated per-request, causing cache misses). func TestExtractSessionID_IdempotencyKey(t *testing.T) { t.Parallel() metadata := map[string]any{"idempotency_key": "idem-12345"} got := ExtractSessionID(nil, nil, metadata) // idempotency_key is disabled - should return empty (no payload to hash) if got != "" { t.Errorf("ExtractSessionID() with idempotency_key = %q, want empty (idempotency_key is disabled)", got) } } func TestExtractSessionID_DerivedSessionAndExplicitPriority(t *testing.T) { t.Parallel() metadata := map[string]any{cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:derived-root"} payload := []byte(`{"messages":[{"role":"user","content":"hello"}]}`) if got := ExtractSessionID(nil, payload, metadata); got != "derived:ctx:v1:derived-root" { t.Fatalf("ExtractSessionID() = %q, want derived identity", got) } executionMetadata := map[string]any{ cliproxyexecutor.ExecutionSessionMetadataKey: "execution-session", cliproxyexecutor.DerivedSessionIDMetadataKey: "ctx:v1:derived-root", } if got := ExtractSessionID(nil, payload, executionMetadata); got != "execution:execution-session" { t.Fatalf("ExtractSessionID() = %q, want explicit execution session", got) } explicitPayload := []byte(`{"session_id":"explicit-session","prompt_cache_key":"explicit-cache","messages":[{"role":"user","content":"hello"}]}`) if got := ExtractSessionID(nil, explicitPayload, metadata); got != "session:explicit-session" { t.Fatalf("ExtractSessionID() = %q, want explicit body session", got) } userPayload := []byte(`{"metadata":{"user_id":"explicit-user"},"conversation_id":"explicit-conversation","messages":[{"role":"user","content":"hello"}]}`) if got := ExtractSessionID(nil, userPayload, metadata); got != "user:explicit-user" { t.Fatalf("ExtractSessionID() = %q, want explicit metadata.user_id", got) } lowercaseHeaders := http.Header{"x-session-id": []string{" lowercase-session "}} if got := ExtractSessionID(lowercaseHeaders, payload, metadata); got != "header:lowercase-session" { t.Fatalf("ExtractSessionID() = %q, want case-insensitive trimmed header session", got) } headers := make(http.Header) headers.Set("X-Session-ID", "header-session") if got := ExtractSessionID(headers, explicitPayload, metadata); got != "header:header-session" { t.Fatalf("ExtractSessionID() = %q, want explicit header session", got) } } func TestExtractSessionID_MessageHashFallback(t *testing.T) { t.Parallel() // First request (user only) generates short hash firstRequestPayload := []byte(`{"messages":[{"role":"user","content":"Hello world"}]}`) shortHash := ExtractSessionID(nil, firstRequestPayload, nil) if shortHash == "" { t.Error("ExtractSessionID() first request should return short hash") } if !strings.HasPrefix(shortHash, "msg:") { t.Errorf("ExtractSessionID() = %q, want prefix 'msg:'", shortHash) } // Multi-turn with assistant generates full hash (different from short hash) multiTurnPayload := []byte(`{"messages":[ {"role":"user","content":"Hello world"}, {"role":"assistant","content":"Hi! How can I help?"}, {"role":"user","content":"Tell me a joke"} ]}`) fullHash := ExtractSessionID(nil, multiTurnPayload, nil) if fullHash == "" { t.Error("ExtractSessionID() multi-turn should return full hash") } if fullHash == shortHash { t.Error("Full hash should differ from short hash (includes assistant)") } // Same multi-turn payload should produce same hash fullHash2 := ExtractSessionID(nil, multiTurnPayload, nil) if fullHash != fullHash2 { t.Errorf("ExtractSessionID() not stable: got %q then %q", fullHash, fullHash2) } } func TestExtractSessionID_ClaudeAPITopLevelSystem(t *testing.T) { t.Parallel() // Claude API: system prompt in top-level "system" field (array format) arraySystem := []byte(`{ "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], "system": [{"type": "text", "text": "You are Claude Code"}] }`) got1 := ExtractSessionID(nil, arraySystem, nil) if got1 == "" || !strings.HasPrefix(got1, "msg:") { t.Errorf("ExtractSessionID() with array system = %q, want msg:* prefix", got1) } // Claude API: system prompt in top-level "system" field (string format) stringSystem := []byte(`{ "messages": [{"role": "user", "content": "Hello"}], "system": "You are Claude Code" }`) got2 := ExtractSessionID(nil, stringSystem, nil) if got2 == "" || !strings.HasPrefix(got2, "msg:") { t.Errorf("ExtractSessionID() with string system = %q, want msg:* prefix", got2) } // Multi-turn with top-level system should produce stable hash multiTurn := []byte(`{ "messages": [ {"role": "user", "content": "Hello"}, {"role": "assistant", "content": "Hi!"}, {"role": "user", "content": "Help me"} ], "system": "You are Claude Code" }`) got3 := ExtractSessionID(nil, multiTurn, nil) if got3 == "" { t.Error("ExtractSessionID() multi-turn with top-level system should return hash") } if got3 == got2 { t.Error("Multi-turn hash should differ from first-turn hash (includes assistant)") } } func TestExtractSessionID_GeminiFormat(t *testing.T) { t.Parallel() // Gemini format with systemInstruction and contents payload := []byte(`{ "systemInstruction": {"parts": [{"text": "You are a helpful assistant."}]}, "contents": [ {"role": "user", "parts": [{"text": "Hello Gemini"}]}, {"role": "model", "parts": [{"text": "Hi there!"}]} ] }`) got := ExtractSessionID(nil, payload, nil) if got == "" { t.Error("ExtractSessionID() with Gemini format should return hash-based session ID") } if !strings.HasPrefix(got, "msg:") { t.Errorf("ExtractSessionID() = %q, want prefix 'msg:'", got) } // Same payload should produce same hash got2 := ExtractSessionID(nil, payload, nil) if got != got2 { t.Errorf("ExtractSessionID() not stable: got %q then %q", got, got2) } // Different user message should produce different hash differentPayload := []byte(`{ "systemInstruction": {"parts": [{"text": "You are a helpful assistant."}]}, "contents": [ {"role": "user", "parts": [{"text": "Hello different"}]}, {"role": "model", "parts": [{"text": "Hi there!"}]} ] }`) got3 := ExtractSessionID(nil, differentPayload, nil) if got == got3 { t.Errorf("ExtractSessionID() should produce different hash for different user message") } } func TestExtractSessionID_OpenAIResponsesAPI(t *testing.T) { t.Parallel() firstTurn := []byte(`{ "instructions": "You are Codex, based on GPT-5.", "input": [ {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "system instructions"}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} ] }`) got1 := ExtractSessionID(nil, firstTurn, nil) if got1 == "" { t.Error("ExtractSessionID() should return hash for OpenAI Responses API format") } if !strings.HasPrefix(got1, "msg:") { t.Errorf("ExtractSessionID() = %q, want prefix 'msg:'", got1) } secondTurn := []byte(`{ "instructions": "You are Codex, based on GPT-5.", "input": [ {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "system instructions"}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "thinking..."}], "encrypted_content": "xxx"}, {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hello!"}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "what can you do"}]} ] }`) got2 := ExtractSessionID(nil, secondTurn, nil) if got2 == "" { t.Error("ExtractSessionID() should return hash for second turn") } if got1 == got2 { t.Log("First turn and second turn have different hashes (expected: second includes assistant)") } thirdTurn := []byte(`{ "instructions": "You are Codex, based on GPT-5.", "input": [ {"type": "message", "role": "developer", "content": [{"type": "input_text", "text": "system instructions"}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, {"type": "reasoning", "summary": [{"type": "summary_text", "text": "thinking..."}], "encrypted_content": "xxx"}, {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Hello!"}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "what can you do"}]}, {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "I can help with..."}]}, {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "thanks"}]} ] }`) got3 := ExtractSessionID(nil, thirdTurn, nil) if got2 != got3 { t.Errorf("Second and third turn should have same hash (same first assistant): got %q vs %q", got2, got3) } } func TestSessionAffinitySelector_ThreeScenarios(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}} testCases := []struct { name string scenario string payload []byte }{ { name: "OpenAI_Scenario1_NewRequest", scenario: "new", payload: []byte(`{"messages":[{"role":"system","content":"You are helpful"},{"role":"user","content":"Hello"}]}`), }, { name: "OpenAI_Scenario2_SecondTurn", scenario: "second", payload: []byte(`{"messages":[{"role":"system","content":"You are helpful"},{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"},{"role":"user","content":"Help me"}]}`), }, { name: "OpenAI_Scenario3_ManyTurns", scenario: "many", payload: []byte(`{"messages":[{"role":"system","content":"You are helpful"},{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there!"},{"role":"user","content":"Help me"},{"role":"assistant","content":"Sure!"},{"role":"user","content":"Thanks"}]}`), }, { name: "Gemini_Scenario1_NewRequest", scenario: "new", payload: []byte(`{"systemInstruction":{"parts":[{"text":"You are helpful"}]},"contents":[{"role":"user","parts":[{"text":"Hello Gemini"}]}]}`), }, { name: "Gemini_Scenario2_SecondTurn", scenario: "second", payload: []byte(`{"systemInstruction":{"parts":[{"text":"You are helpful"}]},"contents":[{"role":"user","parts":[{"text":"Hello Gemini"}]},{"role":"model","parts":[{"text":"Hi!"}]},{"role":"user","parts":[{"text":"Help"}]}]}`), }, { name: "Gemini_Scenario3_ManyTurns", scenario: "many", payload: []byte(`{"systemInstruction":{"parts":[{"text":"You are helpful"}]},"contents":[{"role":"user","parts":[{"text":"Hello Gemini"}]},{"role":"model","parts":[{"text":"Hi!"}]},{"role":"user","parts":[{"text":"Help"}]},{"role":"model","parts":[{"text":"Sure!"}]},{"role":"user","parts":[{"text":"Thanks"}]}]}`), }, { name: "Claude_Scenario1_NewRequest", scenario: "new", payload: []byte(`{"messages":[{"role":"user","content":"Hello Claude"}]}`), }, { name: "Claude_Scenario2_SecondTurn", scenario: "second", payload: []byte(`{"messages":[{"role":"user","content":"Hello Claude"},{"role":"assistant","content":"Hello!"},{"role":"user","content":"Help me"}]}`), }, { name: "Claude_Scenario3_ManyTurns", scenario: "many", payload: []byte(`{"messages":[{"role":"user","content":"Hello Claude"},{"role":"assistant","content":"Hello!"},{"role":"user","content":"Help"},{"role":"assistant","content":"Sure!"},{"role":"user","content":"Thanks"}]}`), }, } for _, tc := range testCases { tc := tc t.Run(tc.name, func(t *testing.T) { opts := cliproxyexecutor.Options{OriginalRequest: tc.payload} picked, err := selector.Pick(context.Background(), "provider", "model", opts, auths) if err != nil { t.Fatalf("Pick() error = %v", err) } if picked == nil { t.Fatal("Pick() returned nil") } t.Logf("%s: picked %s", tc.name, picked.ID) }) } t.Run("Scenario2And3_SameAuth", func(t *testing.T) { openaiS2 := []byte(`{"messages":[{"role":"system","content":"Stable test"},{"role":"user","content":"First msg"},{"role":"assistant","content":"Response"},{"role":"user","content":"Second"}]}`) openaiS3 := []byte(`{"messages":[{"role":"system","content":"Stable test"},{"role":"user","content":"First msg"},{"role":"assistant","content":"Response"},{"role":"user","content":"Second"},{"role":"assistant","content":"More"},{"role":"user","content":"Third"}]}`) opts2 := cliproxyexecutor.Options{OriginalRequest: openaiS2} opts3 := cliproxyexecutor.Options{OriginalRequest: openaiS3} picked2, _ := selector.Pick(context.Background(), "test", "model", opts2, auths) picked3, _ := selector.Pick(context.Background(), "test", "model", opts3, auths) if picked2.ID != picked3.ID { t.Errorf("Scenario2 and Scenario3 should pick same auth: got %s vs %s", picked2.ID, picked3.ID) } }) t.Run("Scenario1To2_InheritBinding", func(t *testing.T) { s1 := []byte(`{"messages":[{"role":"system","content":"Inherit test"},{"role":"user","content":"Initial"}]}`) s2 := []byte(`{"messages":[{"role":"system","content":"Inherit test"},{"role":"user","content":"Initial"},{"role":"assistant","content":"Reply"},{"role":"user","content":"Continue"}]}`) opts1 := cliproxyexecutor.Options{OriginalRequest: s1} opts2 := cliproxyexecutor.Options{OriginalRequest: s2} picked1, _ := selector.Pick(context.Background(), "inherit", "model", opts1, auths) picked2, _ := selector.Pick(context.Background(), "inherit", "model", opts2, auths) if picked1.ID != picked2.ID { t.Errorf("Scenario2 should inherit Scenario1 binding: got %s vs %s", picked1.ID, picked2.ID) } }) } func TestSessionAffinitySelectorBodyIdentifierTransitionsPreserveBinding(t *testing.T) { t.Parallel() bothPayload := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) primaryID, fallbackID := extractSessionIDs(nil, bothPayload, nil) if primaryID != "pck:shared-cache-bucket" || fallbackID != "conv:conversation-session" { t.Fatalf("extractSessionIDs() = (%q, %q), want prompt-cache primary with conversation fallback", primaryID, fallbackID) } for _, tt := range []struct { name string firstPayload []byte }{ {name: "prompt cache first", firstPayload: []byte(`{"prompt_cache_key":"shared-cache-bucket"}`)}, {name: "conversation first", firstPayload: []byte(`{"conversation":{"id":"conversation-session"}}`)}, } { t.Run(tt.name, func(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} provider := "responses-transition-" + tt.name first, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: tt.firstPayload}, auths) if err != nil { t.Fatalf("first Pick() error = %v", err) } second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: bothPayload}, auths) if err != nil { t.Fatalf("combined-identifier Pick() error = %v", err) } if second.ID != first.ID { t.Fatalf("combined identifiers changed auth from %q to %q", first.ID, second.ID) } }) } } func TestSessionAffinitySelectorCombinedIdentifiersBindConversationFallback(t *testing.T) { t.Parallel() selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} provider := "responses-combined-to-conversation" combined := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) conversationOnly := []byte(`{"conversation":{"id":"conversation-session"}}`) first, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: combined}, auths) if err != nil { t.Fatalf("combined-identifier Pick() error = %v", err) } second, err := selector.Pick(context.Background(), provider, "gpt-test", cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) if err != nil { t.Fatalf("conversation-only Pick() error = %v", err) } if second.ID != first.ID { t.Fatalf("dropping prompt_cache_key changed auth from %q to %q", first.ID, second.ID) } } func TestSessionAffinitySelectorPrimaryTrafficKeepsConversationAliasAlive(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} provider := "responses-active-primary-alias" model := "gpt-test" combined := []byte(`{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`) promptOnly := []byte(`{"prompt_cache_key":"shared-cache-bucket"}`) conversationOnly := []byte(`{"conversation":{"id":"conversation-session"}}`) first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combined}, auths) if err != nil { t.Fatalf("combined Pick() error = %v", err) } conversationKey := provider + "::conv:conversation-session::" + model selector.cache.mu.Lock() conversationEntry := selector.cache.entries[conversationKey] conversationEntry.expiresAt = time.Now().Add(-time.Second) selector.cache.entries[conversationKey] = conversationEntry selector.cache.mu.Unlock() primary, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: promptOnly}, auths) if err != nil { t.Fatalf("prompt-only Pick() error = %v", err) } if primary.ID != first.ID { t.Fatalf("prompt-only auth = %q, want %q", primary.ID, first.ID) } fallback, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) if err != nil { t.Fatalf("conversation-only Pick() error = %v", err) } if fallback.ID != first.ID { t.Fatalf("conversation alias expired during active primary traffic: got %q, want %q", fallback.ID, first.ID) } } func TestSessionAffinitySelectorSharedPromptKeyPreservesConversationAliases(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} provider := "responses-shared-prompt-key" model := "gpt-test" combinedA := []byte(`{"conversation":{"id":"conversation-a"},"prompt_cache_key":"shared-cache-bucket"}`) combinedB := []byte(`{"conversation":{"id":"conversation-b"},"prompt_cache_key":"shared-cache-bucket"}`) conversationA := []byte(`{"conversation":{"id":"conversation-a"}}`) conversationB := []byte(`{"conversation":{"id":"conversation-b"}}`) first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combinedA}, auths) if err != nil { t.Fatalf("conversation A combined Pick() error = %v", err) } second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combinedB}, auths) if err != nil { t.Fatalf("conversation B combined Pick() error = %v", err) } if second.ID != first.ID { t.Fatalf("shared prompt key changed auth from %q to %q", first.ID, second.ID) } for name, payload := range map[string][]byte{"conversation A": conversationA, "conversation B": conversationB} { picked, errPick := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: payload}, auths) if errPick != nil { t.Fatalf("%s Pick() error = %v", name, errPick) } if picked.ID != first.ID { t.Fatalf("%s alias selected %q, want %q", name, picked.ID, first.ID) } } } func TestSessionAffinitySelectorConversationIDContainingPromptMarkerRemainsStable(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} provider := "responses-opaque-conversation" model := "gpt-test" combined := []byte(`{"conversation":{"id":"a::pck:b"},"prompt_cache_key":"shared-cache-bucket"}`) conversationOnly := []byte(`{"conversation":{"id":"a::pck:b"}}`) first, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: combined}, auths) if err != nil { t.Fatalf("combined Pick() error = %v", err) } second, err := selector.Pick(context.Background(), provider, model, cliproxyexecutor.Options{OriginalRequest: conversationOnly}, auths) if err != nil { t.Fatalf("conversation-only Pick() error = %v", err) } if second.ID != first.ID { t.Fatalf("opaque conversation alias selected %q, want %q", second.ID, first.ID) } } func TestSessionCacheSharedPromptKeyCapsStableAliasesByRecency(t *testing.T) { cache := NewSessionCache(time.Minute) defer cache.Stop() const promptKey = "openai::pck:shared-cache-bucket::gpt-test" for index := 0; index < 128; index++ { conversation := fmt.Sprintf("openai::conv:conversation-%03d::gpt-test", index) cache.SetAliases("auth-a", promptKey, conversation) } cache.mu.RLock() defer cache.mu.RUnlock() if len(cache.entries) > 65 { t.Fatalf("cache entries = %d, want one prompt key plus at most 64 stable aliases", len(cache.entries)) } if _, ok := cache.entries["openai::conv:conversation-127::gpt-test"]; !ok { t.Fatal("newest conversation alias was not retained") } if _, ok := cache.entries["openai::conv:conversation-000::gpt-test"]; ok { t.Fatal("oldest conversation alias was retained after stable-alias cap") } } func TestSessionCacheRotatingPrimaryEvictsObsoleteAliases(t *testing.T) { cache := NewSessionCache(time.Minute) defer cache.Stop() const fallback = "openai::conv:conversation-session::gpt-test" for index := 0; index < 16; index++ { primary := fmt.Sprintf("openai::pck:cache-%02d::gpt-test", index) cache.SetAliases("auth-a", primary, fallback) } latest := "openai::pck:cache-15::gpt-test" oldest := "openai::pck:cache-00::gpt-test" cache.mu.RLock() defer cache.mu.RUnlock() if len(cache.entries) != 2 { t.Fatalf("cache entries = %d, want only latest primary and fallback", len(cache.entries)) } if _, ok := cache.entries[latest]; !ok { t.Fatalf("latest primary %q was not retained", latest) } if _, ok := cache.entries[fallback]; !ok { t.Fatalf("fallback %q was not retained", fallback) } if _, ok := cache.entries[oldest]; ok { t.Fatalf("obsolete primary %q was retained", oldest) } if aliases := cache.entries[fallback].aliases; len(aliases) != 2 { t.Fatalf("fallback alias group = %#v, want exactly two active identifiers", aliases) } } func TestSessionAffinitySelector_MultiModelSession(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() // auth-a supports only model-a, auth-b supports only model-b authA := &Auth{ID: "auth-a"} authB := &Auth{ID: "auth-b"} // Same session ID for all requests payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_multi-model-test"}}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} // Request model-a with only auth-a available for that model authsForModelA := []*Auth{authA} pickedA, err := selector.Pick(context.Background(), "provider", "model-a", opts, authsForModelA) if err != nil { t.Fatalf("Pick() for model-a error = %v", err) } if pickedA.ID != "auth-a" { t.Fatalf("Pick() for model-a = %q, want auth-a", pickedA.ID) } // Request model-b with only auth-b available for that model authsForModelB := []*Auth{authB} pickedB, err := selector.Pick(context.Background(), "provider", "model-b", opts, authsForModelB) if err != nil { t.Fatalf("Pick() for model-b error = %v", err) } if pickedB.ID != "auth-b" { t.Fatalf("Pick() for model-b = %q, want auth-b", pickedB.ID) } // Switch back to model-a - should still get auth-a (separate binding per model) pickedA2, err := selector.Pick(context.Background(), "provider", "model-a", opts, authsForModelA) if err != nil { t.Fatalf("Pick() for model-a (2nd) error = %v", err) } if pickedA2.ID != "auth-a" { t.Fatalf("Pick() for model-a (2nd) = %q, want auth-a", pickedA2.ID) } // Verify bindings are stable for multiple calls for i := 0; i < 5; i++ { gotA, _ := selector.Pick(context.Background(), "provider", "model-a", opts, authsForModelA) gotB, _ := selector.Pick(context.Background(), "provider", "model-b", opts, authsForModelB) if gotA.ID != "auth-a" { t.Fatalf("Pick() #%d for model-a = %q, want auth-a", i, gotA.ID) } if gotB.ID != "auth-b" { t.Fatalf("Pick() #%d for model-b = %q, want auth-b", i, gotB.ID) } } } func TestExtractSessionID_MultimodalContent(t *testing.T) { t.Parallel() // First request generates short hash firstRequestPayload := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"Hello world"},{"type":"image","source":{"data":"..."}}]}]}`) shortHash := ExtractSessionID(nil, firstRequestPayload, nil) if shortHash == "" { t.Error("ExtractSessionID() first request should return short hash") } if !strings.HasPrefix(shortHash, "msg:") { t.Errorf("ExtractSessionID() = %q, want prefix 'msg:'", shortHash) } // Multi-turn generates full hash multiTurnPayload := []byte(`{"messages":[ {"role":"user","content":[{"type":"text","text":"Hello world"},{"type":"image","source":{"data":"..."}}]}, {"role":"assistant","content":"I see an image!"}, {"role":"user","content":"What is it?"} ]}`) fullHash := ExtractSessionID(nil, multiTurnPayload, nil) if fullHash == "" { t.Error("ExtractSessionID() multimodal multi-turn should return full hash") } if fullHash == shortHash { t.Error("Full hash should differ from short hash") } // Different user content produces different hash differentPayload := []byte(`{"messages":[ {"role":"user","content":[{"type":"text","text":"Different content"}]}, {"role":"assistant","content":"I see something different!"} ]}`) differentHash := ExtractSessionID(nil, differentPayload, nil) if fullHash == differentHash { t.Errorf("ExtractSessionID() should produce different hash for different content") } } func TestSessionAffinitySelector_CrossProviderIsolation(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() authClaude := &Auth{ID: "auth-claude"} authGemini := &Auth{ID: "auth-gemini"} // Same session ID for both providers payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_cross-provider-test"}}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} // Request via claude provider pickedClaude, err := selector.Pick(context.Background(), "claude", "claude-3", opts, []*Auth{authClaude}) if err != nil { t.Fatalf("Pick() for claude error = %v", err) } if pickedClaude.ID != "auth-claude" { t.Fatalf("Pick() for claude = %q, want auth-claude", pickedClaude.ID) } // Same session but via gemini provider should get different auth pickedGemini, err := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", opts, []*Auth{authGemini}) if err != nil { t.Fatalf("Pick() for gemini error = %v", err) } if pickedGemini.ID != "auth-gemini" { t.Fatalf("Pick() for gemini = %q, want auth-gemini", pickedGemini.ID) } // Verify both bindings remain stable for i := 0; i < 5; i++ { gotC, _ := selector.Pick(context.Background(), "claude", "claude-3", opts, []*Auth{authClaude}) gotG, _ := selector.Pick(context.Background(), "gemini", "gemini-2.5-pro", opts, []*Auth{authGemini}) if gotC.ID != "auth-claude" { t.Fatalf("Pick() #%d for claude = %q, want auth-claude", i, gotC.ID) } if gotG.ID != "auth-gemini" { t.Fatalf("Pick() #%d for gemini = %q, want auth-gemini", i, gotG.ID) } } } func TestSessionCache_GetAndRefresh(t *testing.T) { t.Parallel() cache := NewSessionCache(100 * time.Millisecond) defer cache.Stop() cache.Set("session1", "auth1") // Verify initial value got, ok := cache.GetAndRefresh("session1") if !ok || got != "auth1" { t.Fatalf("GetAndRefresh() = %q, %v, want auth1, true", got, ok) } // Wait half TTL and access again (should refresh) time.Sleep(60 * time.Millisecond) got, ok = cache.GetAndRefresh("session1") if !ok || got != "auth1" { t.Fatalf("GetAndRefresh() after 60ms = %q, %v, want auth1, true", got, ok) } // Wait another 60ms (total 120ms from original, but TTL refreshed at 60ms) // Entry should still be valid because TTL was refreshed time.Sleep(60 * time.Millisecond) got, ok = cache.GetAndRefresh("session1") if !ok || got != "auth1" { t.Fatalf("GetAndRefresh() after refresh = %q, %v, want auth1, true (TTL should have been refreshed)", got, ok) } // Now wait full TTL without access time.Sleep(110 * time.Millisecond) got, ok = cache.GetAndRefresh("session1") if ok { t.Fatalf("GetAndRefresh() after expiry = %q, %v, want '', false", got, ok) } } func TestSessionAffinitySelector_RoundRobinDistribution(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{ {ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}, } sessionCount := 12 counts := make(map[string]int) for i := 0; i < sessionCount; i++ { payload := []byte(fmt.Sprintf(`{"metadata":{"user_id":"user_xxx_account__session_%08d-0000-0000-0000-000000000000"}}`, i)) opts := cliproxyexecutor.Options{OriginalRequest: payload} got, err := selector.Pick(context.Background(), "provider", "model", opts, auths) if err != nil { t.Fatalf("Pick() session %d error = %v", i, err) } counts[got.ID]++ } expected := sessionCount / len(auths) for _, auth := range auths { got := counts[auth.ID] if got != expected { t.Errorf("auth %s got %d sessions, want %d (round-robin should distribute evenly)", auth.ID, got, expected) } } } func TestSessionAffinitySelector_Concurrent(t *testing.T) { t.Parallel() fallback := &RoundRobinSelector{} selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: fallback, TTL: time.Minute, }) defer selector.Stop() auths := []*Auth{ {ID: "auth-a"}, {ID: "auth-b"}, {ID: "auth-c"}, } payload := []byte(`{"metadata":{"user_id":"user_xxx_account__session_concurrent-test"}}`) opts := cliproxyexecutor.Options{OriginalRequest: payload} // First pick to establish binding first, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { t.Fatalf("Initial Pick() error = %v", err) } expectedID := first.ID start := make(chan struct{}) var wg sync.WaitGroup errCh := make(chan error, 1) goroutines := 32 iterations := 50 for i := 0; i < goroutines; i++ { wg.Add(1) go func() { defer wg.Done() <-start for j := 0; j < iterations; j++ { got, err := selector.Pick(context.Background(), "claude", "claude-3", opts, auths) if err != nil { select { case errCh <- err: default: } return } if got.ID != expectedID { select { case errCh <- fmt.Errorf("concurrent Pick() returned %q, want %q", got.ID, expectedID): default: } return } } }() } close(start) wg.Wait() select { case err := <-errCh: t.Fatalf("concurrent Pick() error = %v", err) default: } } func TestExtractSessionIDNativeSignals(t *testing.T) { t.Parallel() tests := []struct { name string headers http.Header payload string want string }{ { name: "claude code header", headers: http.Header{"X-Claude-Code-Session-Id": []string{"claude-session"}}, want: "claude:claude-session", }, { name: "lowercase claude code header", headers: http.Header{"x-claude-code-session-id": []string{"lowercase-session"}}, want: "claude:lowercase-session", }, { name: "codex hyphen header", headers: http.Header{"Session-Id": []string{"codex-session"}}, want: "codex:codex-session", }, { name: "codex underscore header", headers: http.Header{"Session_id": []string{"legacy-codex-session"}}, want: "codex:legacy-codex-session", }, { name: "open code session affinity", headers: http.Header{"X-Session-Affinity": []string{"ses_opencode"}}, want: "affinity:ses_opencode", }, { name: "prompt cache key", payload: `{"prompt_cache_key":"prompt-session"}`, want: "pck:prompt-session", }, { name: "responses conversation object", payload: `{"conversation":{"id":"conv-object"}}`, want: "conv:conv-object", }, { name: "responses conversation string", payload: `{"conversation":"conv-string"}`, want: "conv:conv-string", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) } }) } } func TestExtractSessionIDNativeSignalPriority(t *testing.T) { t.Parallel() tests := []struct { name string headers http.Header payload string want string }{ { name: "claude header beats metadata", headers: http.Header{ "X-Claude-Code-Session-Id": []string{"header-session"}, }, payload: `{"metadata":{"user_id":"user_hash_account__session_22222222-2222-4222-8222-222222222222"}}`, want: "claude:header-session", }, { name: "claude metadata beats codex header", headers: http.Header{ "Session-Id": []string{"codex-session"}, }, payload: `{"metadata":{"user_id":"user_hash_account__session_22222222-2222-4222-8222-222222222222"}}`, want: "claude:22222222-2222-4222-8222-222222222222", }, { name: "codex header beats x session id and prompt key", headers: http.Header{ "Session-Id": []string{"codex-session"}, "X-Session-Id": []string{"generic-session"}, }, payload: `{"prompt_cache_key":"prompt-session"}`, want: "codex:codex-session", }, { name: "x session id beats affinity", headers: http.Header{ "X-Session-Id": []string{"generic-session"}, "X-Session-Affinity": []string{"affinity-session"}, }, want: "header:generic-session", }, { name: "prompt cache key beats conversation id", payload: `{"conversation":{"id":"conversation-session"},"prompt_cache_key":"shared-cache-bucket"}`, want: "pck:shared-cache-bucket", }, { name: "client request id beats body fallbacks", headers: http.Header{ "X-Client-Request-Id": []string{"client-session"}, }, payload: `{"prompt_cache_key":"prompt-session","conversation":{"id":"conversation-session"}}`, want: "clientreq:client-session", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) } }) } } func TestExtractSessionIDRejectsInvalidExplicitSignals(t *testing.T) { t.Parallel() tooLong := strings.Repeat("a", 257) tests := []struct { name string headers http.Header payload string want string }{ { name: "whitespace", headers: http.Header{"X-Claude-Code-Session-Id": []string{" "}}, want: "", }, { name: "newline", headers: http.Header{"X-Session-Id": []string{"bad\nsession"}}, want: "", }, { name: "control character", headers: http.Header{"Session-Id": []string{"bad\x00session"}}, want: "", }, { name: "too long", headers: http.Header{"X-Client-Request-Id": []string{tooLong}}, want: "", }, { name: "invalid stronger signal falls through", headers: http.Header{ "X-Claude-Code-Session-Id": []string{"bad\nsession"}, "Session-Id": []string{"valid-codex"}, }, want: "codex:valid-codex", }, { name: "invalid prompt key falls through to conversation", payload: `{"prompt_cache_key":" ","conversation":{"id":"valid-conversation"}}`, want: "conv:valid-conversation", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() if got := ExtractSessionID(tt.headers, []byte(tt.payload), nil); got != tt.want { t.Fatalf("ExtractSessionID() = %q, want %q", got, tt.want) } }) } } func TestExtractSessionIDClaudeMetadataParsesBeforeBoundingSessionID(t *testing.T) { t.Parallel() const sessionID = "11111111-1111-4111-8111-111111111111" metadata := map[string]string{ "device_id": strings.Repeat("d", 64), "account_uuid": "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa", "session_id": sessionID, "organization_uuid": "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", "email": "user@example.com", } for _, tt := range []struct { name string encode func(any) ([]byte, error) }{ {name: "rich compact json", encode: json.Marshal}, {name: "pretty printed json", encode: func(v any) ([]byte, error) { return json.MarshalIndent(v, "", " ") }}, } { t.Run(tt.name, func(t *testing.T) { userID, errMarshal := tt.encode(metadata) if errMarshal != nil { t.Fatalf("marshal metadata: %v", errMarshal) } payload, errPayload := json.Marshal(map[string]any{ "metadata": map[string]string{"user_id": string(userID)}, }) if errPayload != nil { t.Fatalf("marshal payload: %v", errPayload) } if got := ExtractSessionID(nil, payload, nil); got != "claude:"+sessionID { t.Fatalf("ExtractSessionID() = %q, want %q", got, "claude:"+sessionID) } }) } } func TestSessionAffinitySelectorUsesRequestPayloadWhenOriginalRequestMissing(t *testing.T) { selector := NewSessionAffinitySelectorWithConfig(SessionAffinityConfig{ Fallback: &RoundRobinSelector{}, TTL: time.Minute, }) defer selector.Stop() request := cliproxyexecutor.Request{ Model: "gpt-test", Payload: []byte(`{"conversation":{"id":"request-only-conversation"},"input":"hello"}`), } _, opts := cliproxysession.Enrich(request, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatOpenAIResponse, }) auths := []*Auth{{ID: "auth-a"}, {ID: "auth-b"}} first, errFirst := selector.Pick(context.Background(), "openai", request.Model, opts, auths) if errFirst != nil { t.Fatalf("first Pick() error = %v", errFirst) } second, errSecond := selector.Pick(context.Background(), "openai", request.Model, opts, auths) if errSecond != nil { t.Fatalf("second Pick() error = %v", errSecond) } if second.ID != first.ID { t.Fatalf("request-only conversation changed auth from %q to %q", first.ID, second.ID) } } func TestSessionCache_StopConcurrent(t *testing.T) { t.Parallel() for iter := 0; iter < 100; iter++ { cache := NewSessionCache(time.Minute) var wg sync.WaitGroup for i := 0; i < 20; i++ { wg.Add(1) go func() { defer wg.Done() cache.Stop() }() } wg.Wait() } } type mockStoppableSelector struct { stopped bool } func (m *mockStoppableSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { return nil, nil } func (m *mockStoppableSelector) Stop() { m.stopped = true } func TestManagerSetSelectorStopsReplacedStoppableSelector(t *testing.T) { t.Parallel() mockSelector := &mockStoppableSelector{} manager := NewManager(nil, mockSelector, nil) manager.SetSelector(&RoundRobinSelector{}) if !mockSelector.stopped { t.Fatal("expected previous StoppableSelector to be stopped when replaced via SetSelector") } } type zeroSizeSelectorA struct { stopped *bool } func (z zeroSizeSelectorA) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { return nil, nil } func (z zeroSizeSelectorA) Stop() { if z.stopped != nil { *z.stopped = true } } type zeroSizeSelectorB struct{} func (z zeroSizeSelectorB) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { return nil, nil } func TestManagerSetSelectorDifferentZeroSizedSelectors(t *testing.T) { t.Parallel() stoppedA := false selA := zeroSizeSelectorA{stopped: &stoppedA} selB := zeroSizeSelectorB{} manager := NewManager(nil, selA, nil) manager.SetSelector(selB) if !stoppedA { t.Fatal("expected zeroSizeSelectorA to be stopped when replaced by zeroSizeSelectorB") } if manager.Selector() != selB { t.Fatalf("expected manager selector to be selB, got %#v", manager.Selector()) } } type uncomparableSelector struct { fn func() } func (u uncomparableSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { return nil, nil } func TestManagerSetSelectorUncomparableTypes(t *testing.T) { t.Parallel() manager := NewManager(nil, nil, nil) sel1 := uncomparableSelector{fn: func() {}} sel2 := uncomparableSelector{fn: func() {}} // Setting uncomparable types must not panic manager.SetSelector(sel1) manager.SetSelector(sel2) manager.SetSelector(nil) } func TestManagerSetSelectorSameInstanceDoesNotStop(t *testing.T) { t.Parallel() mockSelector := &mockStoppableSelector{} manager := NewManager(nil, mockSelector, nil) // Setting the same instance should be a no-op and not call Stop manager.SetSelector(mockSelector) if mockSelector.stopped { t.Fatal("setting the same selector instance unexpectedly called Stop") } } func TestManagerSetSelectorConcurrent(t *testing.T) { t.Parallel() manager := NewManager(nil, nil, nil) var wg sync.WaitGroup for i := 0; i < 20; i++ { wg.Add(1) go func() { defer wg.Done() for j := 0; j < 10; j++ { sel := &mockStoppableSelector{} manager.SetSelector(sel) } }() } wg.Wait() }