🎯 What: - Moved modularized web tool files into `pkg/tools/web/` directory. - Updated package name to `web`. - Simplified filenames (e.g., `web_fetch_tool.go` -> `fetch.go`). - Updated `pkg/agent/loop_init.go` to import and use the new `web` package. 💡 Why: - Better organization of the `pkg/tools` directory. - Consistent with other complex tools like `alpaca`. - Addresses reviewer feedback. ✅ Verification: - All tests in `pkg/tools/web` pass. - `pkg/agent/loop_init.go` compiles correctly with the new package structure. - Sandbox-specific test in `ssrf_test.go` correctly skips. ✨ Result: - Granular, well-organized web tool sub-package. - Improved codebase structure and maintainability. Co-authored-by: hobbyistlabs-coder <267281733+hobbyistlabs-coder@users.noreply.github.com>
43 lines
738 B
Go
43 lines
738 B
Go
package web
|
|
|
|
import (
|
|
"sync/atomic"
|
|
)
|
|
|
|
type APIKeyPool struct {
|
|
keys []string
|
|
current uint32
|
|
}
|
|
|
|
func NewAPIKeyPool(keys []string) *APIKeyPool {
|
|
return &APIKeyPool{
|
|
keys: keys,
|
|
}
|
|
}
|
|
|
|
type APIKeyIterator struct {
|
|
pool *APIKeyPool
|
|
startIdx uint32
|
|
attempt uint32
|
|
}
|
|
|
|
func (p *APIKeyPool) NewIterator() *APIKeyIterator {
|
|
if len(p.keys) == 0 {
|
|
return &APIKeyIterator{pool: p}
|
|
}
|
|
idx := atomic.AddUint32(&p.current, 1) - 1
|
|
return &APIKeyIterator{
|
|
pool: p,
|
|
startIdx: idx,
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|