feat(auth): enhance plugin scheduler with HasScheduler support and fast-path tests

- Added `pluginSchedulerState` interface with `HasScheduler` method for improved plugin scheduler state checks.
- Updated `Manager.hasPluginScheduler` to handle `HasScheduler` logic.
- Implemented and tested fast-path handling for inactive plugin schedulers, including mixed provider scenarios.
- Expanded unit test coverage to ensure correct behavior in various scheduler states.
This commit is contained in:
Luis Pater
2026-06-09 13:44:10 +08:00
parent 693ce1c55a
commit 41a4dba670
3 changed files with 88 additions and 2 deletions

View File

@@ -30,6 +30,10 @@ func (h *Host) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest)
return resp, true, nil
}
func (h *Host) HasScheduler() bool {
return h.schedulerRecord() != nil
}
func (h *Host) schedulerRecord() *capabilityRecord {
if h == nil {
return nil

View File

@@ -126,6 +126,10 @@ type PluginScheduler interface {
PickAuth(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error)
}
type pluginSchedulerState interface {
HasScheduler() bool
}
// StoppableSelector is an optional interface for selectors that hold resources.
// Selectors that implement this interface will have Stop called during shutdown.
type StoppableSelector interface {
@@ -245,9 +249,15 @@ func (m *Manager) hasPluginScheduler() bool {
return false
}
m.mu.RLock()
ok := m.pluginScheduler != nil
scheduler := m.pluginScheduler
m.mu.RUnlock()
return ok
if scheduler == nil {
return false
}
if state, ok := scheduler.(pluginSchedulerState); ok {
return state.HasScheduler()
}
return true
}
func isBuiltInSelector(selector Selector) bool {

View File

@@ -55,6 +55,14 @@ func (s *fakePluginScheduler) PickAuth(ctx context.Context, req pluginapi.Schedu
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
@@ -440,6 +448,38 @@ func TestManagerPluginSchedulerSkippedWhenHomeEnabled(t *testing.T) {
}
}
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{}
@@ -645,6 +685,38 @@ func TestManagerPluginSchedulerPickNextMixedSelectsProvider(t *testing.T) {
}
}
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{}