diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 6745b535f..ce4d222ae 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -88,18 +88,31 @@ func NewAPIKeyPool(keys []string) *APIKeyPool { } } -func (p *APIKeyPool) Get() string { +type APIKeyIterator struct { + pool *APIKeyPool + startIdx uint32 + attempt uint32 +} + +func (p *APIKeyPool) NewIterator() *APIKeyIterator { if len(p.keys) == 0 { - return "" - } - if len(p.keys) == 1 { - return p.keys[0] + return &APIKeyIterator{pool: p} } idx := atomic.AddUint32(&p.current, 1) - 1 - if idx >= uint32(len(p.keys))-1 { - atomic.CompareAndSwapUint32(&p.current, idx+1, 0) + return &APIKeyIterator{ + pool: p, + startIdx: idx, } - return p.keys[idx%uint32(len(p.keys))] +} + +func (it *APIKeyIterator) Next() (string, bool) { + length := uint32(len(it.pool.keys)) + if length == 0 || it.attempt >= length { + return "", false + } + key := it.pool.keys[(it.startIdx+it.attempt)%length] + it.attempt++ + return key, true } type SearchProvider interface { @@ -117,13 +130,13 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in url.QueryEscape(query), count) var lastErr error - maxAttempts := len(p.keyPool.keys) - if maxAttempts == 0 { - return "", errors.New("no api key available for Brave") - } + iter := p.keyPool.NewIterator() - for attempt := 0; attempt < maxAttempts; attempt++ { - apiKey := p.keyPool.Get() + for { + apiKey, ok := iter.Next() + if !ok { + break + } req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) if err != nil { @@ -210,13 +223,13 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i } var lastErr error - maxAttempts := len(p.keyPool.keys) - if maxAttempts == 0 { - return "", errors.New("no api key available for Tavily") - } + iter := p.keyPool.NewIterator() - for attempt := 0; attempt < maxAttempts; attempt++ { - apiKey := p.keyPool.Get() + for { + apiKey, ok := iter.Next() + if !ok { + break + } payload := map[string]any{ "api_key": apiKey, @@ -402,13 +415,13 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou searchURL := "https://api.perplexity.ai/chat/completions" var lastErr error - maxAttempts := len(p.keyPool.keys) - if maxAttempts == 0 { - return "", errors.New("no api key available for Perplexity") - } + iter := p.keyPool.NewIterator() - for attempt := 0; attempt < maxAttempts; attempt++ { - apiKey := p.keyPool.Get() + for { + apiKey, ok := iter.Next() + if !ok { + break + } payload := map[string]any{ "model": "sonar", diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 86f7e512e..188fb8adb 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -695,31 +695,48 @@ func TestAPIKeyPool(t *testing.T) { t.Fatalf("unexpected keys: %v", pool.keys) } - // Test Get() - if k := pool.Get(); k != "key1" { - t.Errorf("expected key1, got %s", k) + // Test Iterator: each iterator should cover all keys exactly once + iter := pool.NewIterator() + expected := []string{"key1", "key2", "key3"} + for i, want := range expected { + k, ok := iter.Next() + if !ok { + t.Fatalf("iter.Next() returned false at step %d", i) + } + if k != want { + t.Errorf("step %d: expected %s, got %s", i, want, k) + } } - if k := pool.Get(); k != "key2" { - t.Errorf("expected key2, got %s", k) - } - if k := pool.Get(); k != "key3" { - t.Errorf("expected key3, got %s", k) - } - if k := pool.Get(); k != "key1" { - t.Errorf("expected key1, got %s", k) + // Should be exhausted + if _, ok := iter.Next(); ok { + t.Errorf("expected iterator exhausted after all keys") } + // Second iterator starts at next position (load balancing) + iter2 := pool.NewIterator() + k, ok := iter2.Next() + if !ok { + t.Fatal("iter2.Next() returned false") + } + if k != "key2" { + t.Errorf("expected key2 (round-robin), got %s", k) + } + + // Empty pool emptyPool := NewAPIKeyPool([]string{}) - if k := emptyPool.Get(); k != "" { - t.Errorf("expected empty string, got %s", k) + emptyIter := emptyPool.NewIterator() + if _, ok := emptyIter.Next(); ok { + t.Errorf("expected false for empty pool") } + // Single key pool singlePool := NewAPIKeyPool([]string{"single"}) - if k := singlePool.Get(); k != "single" { - t.Errorf("expected single, got %s", k) + singleIter := singlePool.NewIterator() + if k, ok := singleIter.Next(); !ok || k != "single" { + t.Errorf("expected single, got %s (ok=%v)", k, ok) } - if k := singlePool.Get(); k != "single" { - t.Errorf("expected single, got %s", k) + if _, ok := singleIter.Next(); ok { + t.Errorf("expected exhausted after single key") } }