new iter to get api key

pkg/tools/web.go

@@ -88,18 +88,31 @@ func NewAPIKeyPool(keys []string) *APIKeyPool {
 	}
 }
 
-func (p *APIKeyPool) Get() string {
-	if len(p.keys) == 0 {
-		return ""
+type APIKeyIterator struct {
+	pool     *APIKeyPool
+	startIdx uint32
+	attempt  uint32
 }
-	if len(p.keys) == 1 {
-		return p.keys[0]
+
+func (p *APIKeyPool) NewIterator() *APIKeyIterator {
+	if len(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",

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 := pool.Get(); k != "key2" {
-		t.Errorf("expected key2, got %s", k)
+		if k != want {
+			t.Errorf("step %d: expected %s, got %s", i, want, 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")
 	}
 }