From c3b8c574e99d9ab1dc2a46022a5a4e8b7890cb77 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 3 Jan 2026 06:16:33 +0000 Subject: [PATCH 1/3] Initial plan From be46241ae9a5f483ba6f86909ba70a694fdc7bb7 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 3 Jan 2026 06:23:10 +0000 Subject: [PATCH 2/3] Add race detection test exposing quota group index data race Co-authored-by: nghyane <59473462+nghyane@users.noreply.github.com> --- internal/provider/quota_group_race_test.go | 160 +++++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 internal/provider/quota_group_race_test.go diff --git a/internal/provider/quota_group_race_test.go b/internal/provider/quota_group_race_test.go new file mode 100644 index 0000000..4c696cb --- /dev/null +++ b/internal/provider/quota_group_race_test.go @@ -0,0 +1,160 @@ +package provider + +import ( + "sync" + "testing" + "time" +) + +// TestQuotaGroupIndexRaceCondition tests for race conditions when multiple +// goroutines attempt to create/access the quota group index simultaneously. +func TestQuotaGroupIndexRaceCondition(t *testing.T) { + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + Runtime: nil, // Start with nil runtime + } + + // Register the Antigravity quota resolver + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + var wg sync.WaitGroup + iterations := 100 + + // Simulate concurrent quota group operations + for i := 0; i < iterations; i++ { + wg.Add(1) + go func(iteration int) { + defer wg.Done() + + // Alternate between set and get operations + if iteration%2 == 0 { + idx := getOrCreateQuotaGroupIndex(auth) + if idx == nil { + t.Error("getOrCreateQuotaGroupIndex returned nil") + return + } + idx.setGroupBlocked("claude", "claude-sonnet", time.Now().Add(time.Minute), time.Now().Add(time.Minute)) + } else { + idx := getQuotaGroupIndex(auth) + if idx != nil { + blocked, _ := idx.isGroupBlocked("claude", time.Now()) + _ = blocked // Just check if it works + } + } + }(i) + } + + wg.Wait() + + // Verify the index is in a valid state + idx := getQuotaGroupIndex(auth) + if idx == nil { + t.Fatal("Index should exist after concurrent operations") + } + + // Should be able to check blocked state without panic + _, _ = idx.isGroupBlocked("claude", time.Now()) +} + +// TestConcurrentQuotaPropagation tests concurrent quota propagation and clearing. +func TestConcurrentQuotaPropagation(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + "claude-opus-4": {Status: StatusActive}, + "claude-haiku-4": {Status: StatusActive}, + "gemini-2.5-pro": {Status: StatusActive}, + "gemini-2.5-flash": {Status: StatusActive}, + }, + } + + var wg sync.WaitGroup + now := time.Now() + + // Concurrent propagation + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: now.Add(time.Minute), + BackoffLevel: 1, + } + _ = propagateQuotaToGroup(auth, "claude-sonnet-4", quota, now.Add(time.Minute), now) + }() + } + + // Concurrent clearing + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + time.Sleep(10 * time.Millisecond) // Slight delay to let propagation happen first + _ = clearQuotaGroupOnSuccess(auth, "claude-opus-4", now) + }() + } + + wg.Wait() + + // Verify auth is in a consistent state + if auth.ModelStates == nil { + t.Fatal("ModelStates should not be nil") + } + + for model, state := range auth.ModelStates { + if state == nil { + t.Errorf("Model %s has nil state", model) + } + } +} + +// TestQuotaGroupIndexConsistency tests that quota index and ModelStates stay in sync. +func TestQuotaGroupIndexConsistency(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + }, + } + + now := time.Now() + retryTime := now.Add(time.Minute) + + // Set quota exceeded via propagation + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: retryTime, + BackoffLevel: 1, + } + propagateQuotaToGroup(auth, "claude-sonnet-4", quota, retryTime, now) + + // Check that isAuthBlockedForModel correctly detects the block for an uninitialized model + // in the same quota group + blocked, _, next := isAuthBlockedForModel(auth, "claude-opus-4", now) + if !blocked { + t.Error("claude-opus-4 should be blocked due to quota group") + } + if next.IsZero() || !next.After(now) { + t.Errorf("Expected future retry time, got %v", next) + } + + // Clear the quota + clearQuotaGroupOnSuccess(auth, "claude-sonnet-4", now) + + // Verify the uninitialized model is no longer blocked + blocked, _, _ = isAuthBlockedForModel(auth, "claude-opus-4", now) + if blocked { + t.Error("claude-opus-4 should not be blocked after quota clear") + } +} From 56e8078cb8e1d11c240c9353a8bce51cfcbb40bd Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 3 Jan 2026 06:27:53 +0000 Subject: [PATCH 3/3] Fix race condition in quota group routing by cloning auth in pickNext Co-authored-by: nghyane <59473462+nghyane@users.noreply.github.com> --- internal/provider/manager.go | 5 +- internal/provider/quota_group.go | 11 +- internal/provider/quota_group_race_test.go | 141 +++++++++++++++++---- 3 files changed, 128 insertions(+), 29 deletions(-) diff --git a/internal/provider/manager.go b/internal/provider/manager.go index 63b890b..d25017d 100644 --- a/internal/provider/manager.go +++ b/internal/provider/manager.go @@ -665,7 +665,10 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts Opt if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(candidate.ID, modelKey) { continue } - candidates = append(candidates, candidate) + // Create a shallow clone for the selector to avoid race conditions. + // The selector needs to read auth.Runtime concurrently with MarkResult updates. + // By cloning here (under RLock), we ensure the selector gets a consistent snapshot. + candidates = append(candidates, candidate.Clone()) } if len(candidates) == 0 { m.mu.RUnlock() diff --git a/internal/provider/quota_group.go b/internal/provider/quota_group.go index 9e141e9..8001abd 100644 --- a/internal/provider/quota_group.go +++ b/internal/provider/quota_group.go @@ -129,8 +129,15 @@ func (rd *AuthRuntimeData) GetProviderData() any { } // getOrCreateQuotaGroupIndex returns the quota group index from auth.Runtime, -// creating it if necessary. This function is thread-safe and preserves -// any existing provider-specific runtime data. +// creating it if necessary. +// +// CRITICAL: This function MUST only be called while holding the Manager's lock (m.mu). +// It is NOT safe to call this concurrently on the same auth object without external +// synchronization. The Manager ensures this by holding m.mu during MarkResult and other +// state-modifying operations. +// +// The returned quotaGroupIndex has its own internal mutex for thread-safe read/write +// operations after initialization. func getOrCreateQuotaGroupIndex(auth *Auth) *quotaGroupIndex { if auth == nil { return nil diff --git a/internal/provider/quota_group_race_test.go b/internal/provider/quota_group_race_test.go index 4c696cb..c3bc970 100644 --- a/internal/provider/quota_group_race_test.go +++ b/internal/provider/quota_group_race_test.go @@ -1,54 +1,72 @@ package provider import ( + "context" "sync" "testing" "time" ) -// TestQuotaGroupIndexRaceCondition tests for race conditions when multiple -// goroutines attempt to create/access the quota group index simultaneously. -func TestQuotaGroupIndexRaceCondition(t *testing.T) { +// TestQuotaGroupIndexManagerProtection tests that quota group index operations +// are properly protected by the Manager's lock, simulating real usage patterns. +func TestQuotaGroupIndexManagerProtection(t *testing.T) { + // Simulate Manager holding lock during state updates + var managerMu sync.RWMutex + auth := &Auth{ ID: "test-auth", Provider: "antigravity", - Runtime: nil, // Start with nil runtime + Runtime: nil, } - // Register the Antigravity quota resolver RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) var wg sync.WaitGroup iterations := 100 - // Simulate concurrent quota group operations - for i := 0; i < iterations; i++ { + // Goroutines simulating MarkResult (write operations with full lock) + for i := 0; i < iterations/2; i++ { wg.Add(1) - go func(iteration int) { + go func() { defer wg.Done() - // Alternate between set and get operations - if iteration%2 == 0 { - idx := getOrCreateQuotaGroupIndex(auth) - if idx == nil { - t.Error("getOrCreateQuotaGroupIndex returned nil") - return - } + managerMu.Lock() // Manager holds lock during state updates + idx := getOrCreateQuotaGroupIndex(auth) + if idx != nil { idx.setGroupBlocked("claude", "claude-sonnet", time.Now().Add(time.Minute), time.Now().Add(time.Minute)) - } else { - idx := getQuotaGroupIndex(auth) - if idx != nil { - blocked, _ := idx.isGroupBlocked("claude", time.Now()) - _ = blocked // Just check if it works - } } - }(i) + managerMu.Unlock() + }() + } + + // Goroutines simulating selector Pick (read operations with read lock) + for i := 0; i < iterations/2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + + time.Sleep(time.Millisecond) // Let some writes happen first + + managerMu.RLock() // Selector gets auth pointer while holding read lock + // In real code, selector receives auth pointer from pickNext + idx := getQuotaGroupIndex(auth) + managerMu.RUnlock() + + // Selector can safely call isGroupBlocked on the index it received + // because quotaGroupIndex has its own internal mutex + if idx != nil { + _, _ = idx.isGroupBlocked("claude", time.Now()) + } + }() } wg.Wait() // Verify the index is in a valid state + managerMu.RLock() idx := getQuotaGroupIndex(auth) + managerMu.RUnlock() + if idx == nil { t.Fatal("Index should exist after concurrent operations") } @@ -57,10 +75,12 @@ func TestQuotaGroupIndexRaceCondition(t *testing.T) { _, _ = idx.isGroupBlocked("claude", time.Now()) } -// TestConcurrentQuotaPropagation tests concurrent quota propagation and clearing. +// TestConcurrentQuotaPropagation tests concurrent quota propagation and clearing +// with proper Manager lock protection. func TestConcurrentQuotaPropagation(t *testing.T) { RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + var managerMu sync.RWMutex auth := &Auth{ ID: "test-auth", Provider: "antigravity", @@ -76,11 +96,12 @@ func TestConcurrentQuotaPropagation(t *testing.T) { var wg sync.WaitGroup now := time.Now() - // Concurrent propagation + // Concurrent propagation (with lock, simulating MarkResult) for i := 0; i < 50; i++ { wg.Add(1) go func() { defer wg.Done() + managerMu.Lock() quota := QuotaState{ Exceeded: true, Reason: "quota", @@ -88,31 +109,39 @@ func TestConcurrentQuotaPropagation(t *testing.T) { BackoffLevel: 1, } _ = propagateQuotaToGroup(auth, "claude-sonnet-4", quota, now.Add(time.Minute), now) + managerMu.Unlock() }() } - // Concurrent clearing + // Concurrent clearing (with lock, simulating MarkResult) for i := 0; i < 50; i++ { wg.Add(1) go func() { defer wg.Done() - time.Sleep(10 * time.Millisecond) // Slight delay to let propagation happen first + time.Sleep(10 * time.Millisecond) + managerMu.Lock() _ = clearQuotaGroupOnSuccess(auth, "claude-opus-4", now) + managerMu.Unlock() }() } wg.Wait() // Verify auth is in a consistent state + managerMu.RLock() if auth.ModelStates == nil { + managerMu.RUnlock() t.Fatal("ModelStates should not be nil") } for model, state := range auth.ModelStates { if state == nil { + managerMu.RUnlock() t.Errorf("Model %s has nil state", model) + return } } + managerMu.RUnlock() } // TestQuotaGroupIndexConsistency tests that quota index and ModelStates stay in sync. @@ -158,3 +187,63 @@ func TestQuotaGroupIndexConsistency(t *testing.T) { t.Error("claude-opus-4 should not be blocked after quota clear") } } + +// TestSelectorWithQuotaGroups tests the full workflow: Manager passes auth clones +// to selector, which safely reads quota state. +func TestSelectorWithQuotaGroups(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + var managerMu sync.RWMutex + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + }, + } + + selector := &RoundRobinSelector{} + selector.Start() + defer selector.Stop() + + var wg sync.WaitGroup + + // Goroutine 1: Continuously set quota (simulating MarkResult) + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + managerMu.Lock() + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: time.Now().Add(time.Second), + BackoffLevel: 1, + } + propagateQuotaToGroup(auth, "claude-sonnet-4", quota, time.Now().Add(time.Second), time.Now()) + managerMu.Unlock() + time.Sleep(time.Millisecond) + } + }() + + // Goroutine 2: Continuously pick auth (simulating concurrent requests) + // This mimics what pickNext does: clone auth under RLock before passing to selector + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + managerMu.RLock() + // Manager clones auth before passing to selector (like pickNext does after our fix) + authClone := auth.Clone() + auths := []*Auth{authClone} + managerMu.RUnlock() + + // Selector picks and checks if auth is blocked + _, err := selector.Pick(context.Background(), "antigravity", "claude-opus-4", Options{}, auths) + _ = err // Error is expected when quota is exceeded + time.Sleep(time.Millisecond) + } + }() + + wg.Wait() +}