package auth import ( "context" "errors" "net/http" "testing" "time" internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config" "github.com/router-for-me/CLIProxyAPI/v7/internal/registry" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor" "github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi" ) type schedulerTestExecutor struct{} func (schedulerTestExecutor) Identifier() string { return "test" } func (schedulerTestExecutor) Execute(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, nil } func (schedulerTestExecutor) ExecuteStream(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) { return nil, nil } func (schedulerTestExecutor) Refresh(ctx context.Context, auth *Auth) (*Auth, error) { return auth, nil } func (schedulerTestExecutor) CountTokens(ctx context.Context, auth *Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) { return cliproxyexecutor.Response{}, nil } func (schedulerTestExecutor) HttpRequest(ctx context.Context, auth *Auth, req *http.Request) (*http.Response, error) { return nil, nil } type fakePluginScheduler struct { resp pluginapi.SchedulerPickResponse handled bool err error calls int requests []pluginapi.SchedulerPickRequest pick func(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) } func (s *fakePluginScheduler) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { s.calls++ s.requests = append(s.requests, req) if s.pick != nil { return s.pick(ctx, req) } return s.resp, s.handled, s.err } type inactivePluginScheduler struct { fakePluginScheduler } func (s *inactivePluginScheduler) HasScheduler() bool { return false } type trackingSelector struct { calls int lastAuthID []string } func (s *trackingSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) { s.calls++ s.lastAuthID = s.lastAuthID[:0] for _, auth := range auths { s.lastAuthID = append(s.lastAuthID, auth.ID) } if len(auths) == 0 { return nil, nil } return auths[len(auths)-1], nil } func newSchedulerForTest(selector Selector, auths ...*Auth) *authScheduler { scheduler := newAuthScheduler(selector) scheduler.rebuild(auths) return scheduler } func registerSchedulerModels(t *testing.T, provider string, model string, authIDs ...string) { t.Helper() reg := registry.GetGlobalRegistry() for _, authID := range authIDs { reg.RegisterClient(authID, provider, []*registry.ModelInfo{{ID: model}}) } t.Cleanup(func() { for _, authID := range authIDs { reg.UnregisterClient(authID) } }) } func TestSchedulerPick_RoundRobinHighestPriority(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "low", Provider: "gemini", Attributes: map[string]string{"priority": "0"}}, &Auth{ID: "high-b", Provider: "gemini", Attributes: map[string]string{"priority": "10"}}, &Auth{ID: "high-a", Provider: "gemini", Attributes: map[string]string{"priority": "10"}}, ) want := []string{"high-a", "high-b", "high-a"} for index, wantID := range want { got, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickSingle() #%d auth = nil", index) } if got.ID != wantID { t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, wantID) } } } func TestSchedulerPick_FillFirstSticksToFirstReady(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &FillFirstSelector{}, &Auth{ID: "b", Provider: "gemini"}, &Auth{ID: "a", Provider: "gemini"}, &Auth{ID: "c", Provider: "gemini"}, ) for index := 0; index < 3; index++ { got, errPick := scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickSingle() #%d auth = nil", index) } if got.ID != "a" { t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, "a") } } } func TestSchedulerPick_PromotesExpiredCooldownBeforePick(t *testing.T) { t.Parallel() model := "gemini-2.5-pro" registerSchedulerModels(t, "gemini", model, "cooldown-expired") scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ ID: "cooldown-expired", Provider: "gemini", ModelStates: map[string]*ModelState{ model: { Status: StatusError, Unavailable: true, NextRetryAfter: time.Now().Add(-1 * time.Second), }, }, }, ) got, errPick := scheduler.pickSingle(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() error = %v", errPick) } if got == nil { t.Fatalf("pickSingle() auth = nil") } if got.ID != "cooldown-expired" { t.Fatalf("pickSingle() auth.ID = %q, want %q", got.ID, "cooldown-expired") } } func TestSchedulerPick_CodexWebsocketPrefersWebsocketEnabledSubset(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "codex-http", Provider: "codex"}, &Auth{ID: "codex-ws-a", Provider: "codex", Attributes: map[string]string{"websockets": "true"}}, &Auth{ID: "codex-ws-b", Provider: "codex", Attributes: map[string]string{"websockets": "true"}}, ) ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) want := []string{"codex-ws-a", "codex-ws-b", "codex-ws-a"} for index, wantID := range want { got, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickSingle() #%d auth = nil", index) } if got.ID != wantID { t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, wantID) } } } func TestSchedulerPick_XAIWebsocketPrefersWebsocketEnabledSubset(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "xai-http", Provider: "xai"}, &Auth{ID: "xai-ws-a", Provider: "xai", Attributes: map[string]string{"websockets": "true"}}, &Auth{ID: "xai-ws-b", Provider: "xai", Attributes: map[string]string{"websockets": "true"}}, ) ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) want := []string{"xai-ws-a", "xai-ws-b", "xai-ws-a"} for index, wantID := range want { got, errPick := scheduler.pickSingle(ctx, "xai", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickSingle() #%d auth = nil", index) } if got.ID != wantID { t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, wantID) } } } func TestSchedulerPick_CodexWebsocketPrefersWebsocketEnabledAcrossPriorities(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "codex-http", Provider: "codex", Attributes: map[string]string{"priority": "10"}}, &Auth{ID: "codex-ws-a", Provider: "codex", Attributes: map[string]string{"priority": "0", "websockets": "true"}}, &Auth{ID: "codex-ws-b", Provider: "codex", Attributes: map[string]string{"priority": "0", "websockets": "true"}}, ) ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background()) want := []string{"codex-ws-a", "codex-ws-b", "codex-ws-a"} for index, wantID := range want { got, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickSingle() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickSingle() #%d auth = nil", index) } if got.ID != wantID { t.Fatalf("pickSingle() #%d auth.ID = %q, want %q", index, got.ID, wantID) } } } func TestSchedulerPick_MixedProvidersUsesWeightedProviderRotationOverReadyCandidates(t *testing.T) { t.Parallel() scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "gemini-a", Provider: "gemini"}, &Auth{ID: "gemini-b", Provider: "gemini"}, &Auth{ID: "claude-a", Provider: "claude"}, ) wantProviders := []string{"gemini", "gemini", "claude", "gemini"} wantIDs := []string{"gemini-a", "gemini-b", "claude-a", "gemini-a"} for index := range wantProviders { got, provider, errPick := scheduler.pickMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickMixed() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickMixed() #%d auth = nil", index) } if provider != wantProviders[index] { t.Fatalf("pickMixed() #%d provider = %q, want %q", index, provider, wantProviders[index]) } if got.ID != wantIDs[index] { t.Fatalf("pickMixed() #%d auth.ID = %q, want %q", index, got.ID, wantIDs[index]) } } } func TestSchedulerPick_MixedProvidersPrefersHighestPriorityTier(t *testing.T) { t.Parallel() model := "gpt-default" registerSchedulerModels(t, "provider-low", model, "low") registerSchedulerModels(t, "provider-high-a", model, "high-a") registerSchedulerModels(t, "provider-high-b", model, "high-b") scheduler := newSchedulerForTest( &RoundRobinSelector{}, &Auth{ID: "low", Provider: "provider-low", Attributes: map[string]string{"priority": "4"}}, &Auth{ID: "high-a", Provider: "provider-high-a", Attributes: map[string]string{"priority": "7"}}, &Auth{ID: "high-b", Provider: "provider-high-b", Attributes: map[string]string{"priority": "7"}}, ) providers := []string{"provider-low", "provider-high-a", "provider-high-b"} wantProviders := []string{"provider-high-a", "provider-high-b", "provider-high-a", "provider-high-b"} wantIDs := []string{"high-a", "high-b", "high-a", "high-b"} for index := range wantProviders { got, provider, errPick := scheduler.pickMixed(context.Background(), providers, model, cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickMixed() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickMixed() #%d auth = nil", index) } if provider != wantProviders[index] { t.Fatalf("pickMixed() #%d provider = %q, want %q", index, provider, wantProviders[index]) } if got.ID != wantIDs[index] { t.Fatalf("pickMixed() #%d auth.ID = %q, want %q", index, got.ID, wantIDs[index]) } } } func TestManager_PickNextMixed_UsesWeightedProviderRotationBeforeCredentialRotation(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } wantProviders := []string{"gemini", "gemini", "claude", "gemini"} wantIDs := []string{"gemini-a", "gemini-b", "claude-a", "gemini-a"} for index := range wantProviders { got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, map[string]struct{}{}) if errPick != nil { t.Fatalf("pickNextMixed() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickNextMixed() #%d auth = nil", index) } if provider != wantProviders[index] { t.Fatalf("pickNextMixed() #%d provider = %q, want %q", index, provider, wantProviders[index]) } if got.ID != wantIDs[index] { t.Fatalf("pickNextMixed() #%d auth.ID = %q, want %q", index, got.ID, wantIDs[index]) } } } func TestManager_PickNextMixed_DisallowFreeAuthSkipsCodexFreePlan(t *testing.T) { t.Parallel() model := "gpt-5.4-mini" registerSchedulerModels(t, "codex", model, "codex-a-free", "codex-b-plus") manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["codex"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-a-free", Provider: "codex", Attributes: map[string]string{"plan_type": "free"}}); errRegister != nil { t.Fatalf("Register(codex-a-free) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "codex-b-plus", Provider: "codex", Attributes: map[string]string{"plan_type": "plus"}}); errRegister != nil { t.Fatalf("Register(codex-b-plus) error = %v", errRegister) } opts := cliproxyexecutor.Options{ Metadata: map[string]any{cliproxyexecutor.DisallowFreeAuthMetadataKey: true}, } got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"codex"}, model, opts, map[string]struct{}{}) if errPick != nil { t.Fatalf("pickNextMixed() error = %v", errPick) } if got == nil { t.Fatalf("pickNextMixed() auth = nil") } if provider != "codex" { t.Fatalf("pickNextMixed() provider = %q, want %q", provider, "codex") } if got.ID != "codex-b-plus" { t.Fatalf("pickNextMixed() auth.ID = %q, want %q", got.ID, "codex-b-plus") } } func TestManagerPluginSchedulerSelectsAuthID(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } scheduler := &fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-b"}, handled: true, } manager.SetPluginScheduler(scheduler) got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{Stream: true}, nil) if errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } if got == nil { t.Fatalf("pickNext() auth = nil") } if got.ID != "auth-b" { t.Fatalf("pickNext() auth.ID = %q, want %q", got.ID, "auth-b") } if scheduler.calls != 1 { t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 1) } if len(scheduler.requests) != 1 { t.Fatalf("len(scheduler.requests) = %d, want %d", len(scheduler.requests), 1) } if !scheduler.requests[0].Stream { t.Fatalf("scheduler request Stream = false, want true") } } func TestManagerPluginSchedulerSkippedWhenHomeEnabled(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.SetConfig(&internalconfig.Config{Home: internalconfig.HomeConfig{Enabled: true}}) scheduler := &fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"}, handled: true, } manager.SetPluginScheduler(scheduler) _, _, _ = manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if scheduler.calls != 0 { t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0) } } func TestManagerInactivePluginSchedulerKeepsFastPath(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } scheduler := &inactivePluginScheduler{} manager.SetPluginScheduler(scheduler) gotA, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() first error = %v", errPick) } gotB, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() second error = %v", errPick) } if gotA == nil || gotB == nil { t.Fatalf("pickNext() auths = %v, %v; want non-nil", gotA, gotB) } if gotA.ID != "auth-a" || gotB.ID != "auth-b" { t.Fatalf("fast path picks = %q, %q; want auth-a, auth-b", gotA.ID, gotB.ID) } if scheduler.calls != 0 { t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0) } } func TestManagerPluginSchedulerCalledOutsideManagerLock(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } scheduler := &fakePluginScheduler{ handled: true, pick: func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { if !manager.mu.TryLock() { t.Fatalf("plugin scheduler called while manager lock is held") } manager.mu.Unlock() return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"}, true, nil }, } manager.SetPluginScheduler(scheduler) got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } if got == nil { t.Fatalf("pickNext() auth = nil") } if got.ID != "auth-a" { t.Fatalf("pickNext() auth.ID = %q, want auth-a", got.ID) } if scheduler.calls != 1 { t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 1) } } func TestManagerPluginSchedulerErrorStopsPick(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } scheduler := &fakePluginScheduler{ handled: true, err: errors.New("tenant denied"), } manager.SetPluginScheduler(scheduler) got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick == nil { t.Fatalf("pickNext() error = nil, want tenant denied") } if errPick.Error() != "tenant denied" { t.Fatalf("pickNext() error = %v, want tenant denied", errPick) } if got != nil { t.Fatalf("pickNext() auth = %v, want nil", got) } } func TestManagerPluginSchedulerFallsBackWhenUnhandledOrUnknown(t *testing.T) { for _, tc := range []struct { name string resp pluginapi.SchedulerPickResponse handled bool }{ { name: "unhandled", resp: pluginapi.SchedulerPickResponse{Handled: false}, handled: false, }, { name: "unknown auth id", resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "missing"}, handled: true, }, } { t.Run(tc.name, func(t *testing.T) { manager := NewManager(nil, &FillFirstSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } scheduler := &fakePluginScheduler{resp: tc.resp, handled: tc.handled} manager.SetPluginScheduler(scheduler) got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } if got == nil { t.Fatalf("pickNext() auth = nil") } if got.ID != "auth-a" { t.Fatalf("pickNext() auth.ID = %q, want %q", got.ID, "auth-a") } }) } } func TestManagerPluginSchedulerDelegatesBuiltin(t *testing.T) { t.Run("round-robin", func(t *testing.T) { manager := NewManager(nil, &FillFirstSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } manager.SetPluginScheduler(&fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinRoundRobin}, handled: true, }) gotA, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() first error = %v", errPick) } gotB, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() second error = %v", errPick) } if gotA == nil || gotB == nil { t.Fatalf("pickNext() auths = %v, %v; want non-nil", gotA, gotB) } if gotA.ID != "auth-a" || gotB.ID != "auth-b" { t.Fatalf("round-robin picks = %q, %q; want auth-a, auth-b", gotA.ID, gotB.ID) } }) t.Run("round-robin model cursors", func(t *testing.T) { reg := registry.GetGlobalRegistry() models := []*registry.ModelInfo{{ID: "model-a"}, {ID: "model-b"}} for _, authID := range []string{"auth-a", "auth-b"} { reg.RegisterClient(authID, "gemini", models) t.Cleanup(func() { reg.UnregisterClient(authID) }) } manager := NewManager(nil, &FillFirstSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } manager.SetPluginScheduler(&fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinRoundRobin}, handled: true, }) gotModelA, _, errPick := manager.pickNext(context.Background(), "gemini", "model-a", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext(model-a) error = %v", errPick) } gotModelB, _, errPick := manager.pickNext(context.Background(), "gemini", "model-b", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext(model-b) error = %v", errPick) } if gotModelA == nil || gotModelB == nil { t.Fatalf("pickNext() auths = %v, %v; want non-nil", gotModelA, gotModelB) } if gotModelA.ID != "auth-a" || gotModelB.ID != "auth-a" { t.Fatalf("model-scoped round-robin picks = %q, %q; want auth-a, auth-a", gotModelA.ID, gotModelB.ID) } }) t.Run("fill-first", func(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } manager.SetPluginScheduler(&fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinFillFirst}, handled: true, }) got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } if got == nil { t.Fatalf("pickNext() auth = nil") } if got.ID != "auth-a" { t.Fatalf("fill-first pick = %q, want auth-a", got.ID) } }) } func TestManagerPluginSchedulerDelegateRoundRobinUsesNativeMixedRotation(t *testing.T) { manager := NewManager(nil, &FillFirstSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } manager.SetPluginScheduler(&fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, DelegateBuiltin: pluginapi.SchedulerBuiltinRoundRobin}, handled: true, }) wantProviders := []string{"gemini", "gemini", "claude", "gemini"} wantIDs := []string{"gemini-a", "gemini-b", "claude-a", "gemini-a"} for index := range wantProviders { got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNextMixed() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickNextMixed() #%d auth = nil", index) } if provider != wantProviders[index] { t.Fatalf("pickNextMixed() #%d provider = %q, want %q", index, provider, wantProviders[index]) } if got.ID != wantIDs[index] { t.Fatalf("pickNextMixed() #%d auth.ID = %q, want %q", index, got.ID, wantIDs[index]) } } } func TestManagerPluginSchedulerPickNextMixedSelectsProvider(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } scheduler := &fakePluginScheduler{ resp: pluginapi.SchedulerPickResponse{Handled: true, AuthID: "claude-a"}, handled: true, } manager.SetPluginScheduler(scheduler) got, executor, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNextMixed() error = %v", errPick) } if got == nil { t.Fatalf("pickNextMixed() auth = nil") } if got.ID != "claude-a" { t.Fatalf("pickNextMixed() auth.ID = %q, want claude-a", got.ID) } if provider != "claude" { t.Fatalf("pickNextMixed() provider = %q, want claude", provider) } if executor == nil { t.Fatalf("pickNextMixed() executor = nil") } if len(scheduler.requests) != 1 { t.Fatalf("len(scheduler.requests) = %d, want %d", len(scheduler.requests), 1) } req := scheduler.requests[0] if req.Provider != "" { t.Fatalf("scheduler request Provider = %q, want empty for mixed provider pick", req.Provider) } if len(req.Providers) != 2 || req.Providers[0] != "gemini" || req.Providers[1] != "claude" { t.Fatalf("scheduler request Providers = %#v, want [gemini claude]", req.Providers) } } func TestManagerInactivePluginSchedulerKeepsMixedFastPath(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } scheduler := &inactivePluginScheduler{} manager.SetPluginScheduler(scheduler) got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNextMixed() error = %v", errPick) } if got == nil { t.Fatalf("pickNextMixed() auth = nil") } if provider != "gemini" { t.Fatalf("pickNextMixed() provider = %q, want gemini", provider) } if got.ID != "gemini-a" { t.Fatalf("pickNextMixed() auth.ID = %q, want gemini-a", got.ID) } if scheduler.calls != 0 { t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0) } } func TestManagerPluginSchedulerCandidatesAreSafeCopies(t *testing.T) { manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} auth := &Auth{ ID: "auth-a", Provider: "gemini", Status: StatusActive, Attributes: map[string]string{ "access_token": "token-value", "api_key": "api-key-value", "cookie": "cookie-value", "priority": "7", "team": "alpha", }, Metadata: map[string]any{"tenant": "one"}, } if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } scheduler := &fakePluginScheduler{ handled: true, pick: func(ctx context.Context, req pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error) { if len(req.Candidates) != 1 { t.Fatalf("len(req.Candidates) = %d, want %d", len(req.Candidates), 1) } candidate := req.Candidates[0] if candidate.ID != "auth-a" || candidate.Provider != "gemini" || candidate.Priority != 7 || candidate.Status != string(StatusActive) { t.Fatalf("scheduler candidate = %#v, want sanitized auth-a metadata", candidate) } for _, key := range []string{"access_token", "api_key", "cookie"} { if _, ok := candidate.Attributes[key]; ok { t.Fatalf("scheduler candidate Attributes contains sensitive key %q", key) } } if candidate.Attributes["priority"] != "7" { t.Fatalf("scheduler candidate priority attribute = %q, want 7", candidate.Attributes["priority"]) } if len(candidate.Metadata) != 0 { t.Fatalf("scheduler candidate Metadata = %#v, want empty", candidate.Metadata) } candidate.Attributes["team"] = "mutated" req.Candidates[0] = candidate return pluginapi.SchedulerPickResponse{Handled: true, AuthID: "auth-a"}, true, nil }, } manager.SetPluginScheduler(scheduler) if _, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil); errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } manager.mu.RLock() gotAttr := manager.auths["auth-a"].Attributes["team"] gotAPIKey := manager.auths["auth-a"].Attributes["api_key"] manager.mu.RUnlock() if gotAttr != "alpha" { t.Fatalf("manager auth attribute team = %q, want alpha", gotAttr) } if gotAPIKey != "api-key-value" { t.Fatalf("manager auth attribute api_key = %q, want api-key-value", gotAPIKey) } } func TestManagerCustomSelector_FallsBackToLegacyPath(t *testing.T) { t.Parallel() selector := &trackingSelector{} manager := NewManager(nil, selector, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.auths["auth-a"] = &Auth{ID: "auth-a", Provider: "gemini"} manager.auths["auth-b"] = &Auth{ID: "auth-b", Provider: "gemini"} got, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, map[string]struct{}{}) if errPick != nil { t.Fatalf("pickNext() error = %v", errPick) } if got == nil { t.Fatalf("pickNext() auth = nil") } if selector.calls != 1 { t.Fatalf("selector.calls = %d, want %d", selector.calls, 1) } if len(selector.lastAuthID) != 2 { t.Fatalf("len(selector.lastAuthID) = %d, want %d", len(selector.lastAuthID), 2) } if got.ID != selector.lastAuthID[len(selector.lastAuthID)-1] { t.Fatalf("pickNext() auth.ID = %q, want selector-picked %q", got.ID, selector.lastAuthID[len(selector.lastAuthID)-1]) } } func TestManager_InitializesSchedulerForBuiltInSelector(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) if manager.scheduler == nil { t.Fatalf("manager.scheduler = nil") } if manager.scheduler.strategy != schedulerStrategyRoundRobin { t.Fatalf("manager.scheduler.strategy = %v, want %v", manager.scheduler.strategy, schedulerStrategyRoundRobin) } manager.SetSelector(&FillFirstSelector{}) if manager.scheduler.strategy != schedulerStrategyFillFirst { t.Fatalf("manager.scheduler.strategy = %v, want %v", manager.scheduler.strategy, schedulerStrategyFillFirst) } } func TestManager_SchedulerTracksRegisterAndUpdate(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } got, errPick := manager.scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("scheduler.pickSingle() error = %v", errPick) } if got == nil || got.ID != "auth-a" { t.Fatalf("scheduler.pickSingle() auth = %v, want auth-a", got) } if _, errUpdate := manager.Update(context.Background(), &Auth{ID: "auth-a", Provider: "gemini", Disabled: true}); errUpdate != nil { t.Fatalf("Update(auth-a) error = %v", errUpdate) } got, errPick = manager.scheduler.pickSingle(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("scheduler.pickSingle() after update error = %v", errPick) } if got == nil || got.ID != "auth-b" { t.Fatalf("scheduler.pickSingle() after update auth = %v, want auth-b", got) } } func TestManager_PickNextMixed_UsesSchedulerRotation(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["gemini"] = schedulerTestExecutor{} manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-b) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } wantProviders := []string{"gemini", "gemini", "claude", "gemini"} wantIDs := []string{"gemini-a", "gemini-b", "claude-a", "gemini-a"} for index := range wantProviders { got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNextMixed() #%d error = %v", index, errPick) } if got == nil { t.Fatalf("pickNextMixed() #%d auth = nil", index) } if provider != wantProviders[index] { t.Fatalf("pickNextMixed() #%d provider = %q, want %q", index, provider, wantProviders[index]) } if got.ID != wantIDs[index] { t.Fatalf("pickNextMixed() #%d auth.ID = %q, want %q", index, got.ID, wantIDs[index]) } } } func TestManager_PickNextMixed_SkipsProvidersWithoutExecutors(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) manager.executors["claude"] = schedulerTestExecutor{} if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(gemini-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil { t.Fatalf("Register(claude-a) error = %v", errRegister) } got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("pickNextMixed() error = %v", errPick) } if got == nil { t.Fatalf("pickNextMixed() auth = nil") } if provider != "claude" { t.Fatalf("pickNextMixed() provider = %q, want %q", provider, "claude") } if got.ID != "claude-a" { t.Fatalf("pickNextMixed() auth.ID = %q, want %q", got.ID, "claude-a") } } func TestManager_SchedulerTracksMarkResultCooldownAndRecovery(t *testing.T) { t.Parallel() manager := NewManager(nil, &RoundRobinSelector{}, nil) reg := registry.GetGlobalRegistry() reg.RegisterClient("auth-a", "gemini", []*registry.ModelInfo{{ID: "test-model"}}) reg.RegisterClient("auth-b", "gemini", []*registry.ModelInfo{{ID: "test-model"}}) t.Cleanup(func() { reg.UnregisterClient("auth-a") reg.UnregisterClient("auth-b") }) if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-a) error = %v", errRegister) } if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil { t.Fatalf("Register(auth-b) error = %v", errRegister) } manager.MarkResult(context.Background(), Result{ AuthID: "auth-a", Provider: "gemini", Model: "test-model", Success: false, Error: &Error{HTTPStatus: 429, Message: "quota"}, }) got, errPick := manager.scheduler.pickSingle(context.Background(), "gemini", "test-model", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("scheduler.pickSingle() after cooldown error = %v", errPick) } if got == nil || got.ID != "auth-b" { t.Fatalf("scheduler.pickSingle() after cooldown auth = %v, want auth-b", got) } manager.MarkResult(context.Background(), Result{ AuthID: "auth-a", Provider: "gemini", Model: "test-model", Success: true, }) seen := make(map[string]struct{}, 2) for index := 0; index < 2; index++ { got, errPick = manager.scheduler.pickSingle(context.Background(), "gemini", "test-model", cliproxyexecutor.Options{}, nil) if errPick != nil { t.Fatalf("scheduler.pickSingle() after recovery #%d error = %v", index, errPick) } if got == nil { t.Fatalf("scheduler.pickSingle() after recovery #%d auth = nil", index) } seen[got.ID] = struct{}{} } if len(seen) != 2 { t.Fatalf("len(seen) = %d, want %d", len(seen), 2) } }