From 41a4dba67098b311c0661f7661e4acb00a1fffc1 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Tue, 9 Jun 2026 13:44:10 +0800 Subject: [PATCH] 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. --- internal/pluginhost/scheduler.go | 4 ++ sdk/cliproxy/auth/conductor.go | 14 +++++- sdk/cliproxy/auth/scheduler_test.go | 72 +++++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 2 deletions(-) diff --git a/internal/pluginhost/scheduler.go b/internal/pluginhost/scheduler.go index fa7489563..33781fb02 100644 --- a/internal/pluginhost/scheduler.go +++ b/internal/pluginhost/scheduler.go @@ -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 diff --git a/sdk/cliproxy/auth/conductor.go b/sdk/cliproxy/auth/conductor.go index db03a3516..bd9f466ba 100644 --- a/sdk/cliproxy/auth/conductor.go +++ b/sdk/cliproxy/auth/conductor.go @@ -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 { diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 48a7673eb..ae6ba86f0 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -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{}