feat(model): rate limiting

This commit is contained in:
afjcjsbx 2026-03-30 18:43:03 +02:00
parent 7b3f47128f
commit 3f7cf1a0b6
8 changed files with 475 additions and 30 deletions

93
docs/rate-limiting.md Normal file
View file

@ -0,0 +1,93 @@
# Dynamic Rate Limiting
PicoClaw prevents 429 errors from LLM provider APIs by enforcing configurable per-model request-rate limits **before** sending each request. Unlike the reactive cooldown/fallback system (which activates *after* a 429 is received), rate limiting is **proactive**: it keeps outbound QPS within the provider's free-tier or plan limits.
## How it works
### Token-bucket algorithm
Each rate-limited model gets a token bucket:
- **Capacity** = `rpm` (burst size equals the per-minute limit)
- **Refill rate** = `rpm / 60` tokens per second
- Tokens are consumed one per LLM call; if the bucket is empty, the call blocks until a token refills or the request context is cancelled
### Call chain integration
```
AgentLoop.callLLM()
└─ FallbackChain.Execute() ← iterate candidates
├─ CooldownTracker.IsAvailable() ← skip if post-429 cooldown active
├─ RateLimiterRegistry.Wait() ← NEW: block until token available
└─ provider.Chat() ← actual LLM HTTP call
```
The rate limiter runs **after** the cooldown check and **before** the provider call, so:
- Candidates already in cooldown are skipped entirely (no token consumed)
- Candidates that are available get throttled to the configured RPM
The same check applies in `ExecuteImage`.
### Thread safety
`RateLimiterRegistry` is safe for concurrent use. The per-limiter token bucket uses a fine-grained mutex so concurrent goroutines each acquire their own token independently.
## Configuration
Set `rpm` on any model in `model_list`:
```yaml
model_list:
- model_name: gpt-4o-free
model: openai/gpt-4o
api_base: https://api.openai.com/v1
rpm: 3 # max 3 requests per minute
api_keys:
- sk-...
- model_name: claude-haiku
model: anthropic/claude-haiku-4-5
rpm: 60 # 60 rpm (Anthropic free tier)
api_keys:
- sk-ant-...
- model_name: local-llm
model: openai/llama3
api_base: http://localhost:11434/v1
# no rpm → unrestricted
```
| Field | Type | Default | Description |
|---|---|---|---|
| `rpm` | `int` | `0` | Requests per minute. `0` means no limit. |
### Interaction with fallbacks
When a model has fallbacks configured, each candidate is rate-limited **independently**:
```yaml
model_list:
- model_name: gpt4-with-fallback
model: openai/gpt-4o
rpm: 5
fallbacks:
- gpt-4o-mini # must also be in model_list; its own rpm applies
```
If the primary candidate's bucket is empty the call blocks until a token is available (or context is cancelled). If context deadline is hit while waiting, the error propagates and the agent retries normally.
### Burst behaviour
The bucket starts **full** (burst = RPM). For `rpm: 3`, the first 3 requests fire instantly; subsequent requests are spaced ~20 s apart.
To reduce burstiness for strict APIs, set a lower `rpm` and rely on the steady-state refill.
## Files changed
| File | What |
|---|---|
| `pkg/providers/ratelimiter.go` | `RateLimiter` (token bucket) + `RateLimiterRegistry` |
| `pkg/providers/ratelimiter_test.go` | Unit tests for limiter and registry |
| `pkg/providers/fallback.go` | `FallbackCandidate.RPM` field; `FallbackChain.rl`; `Wait()` call in `Execute`/`ExecuteImage` |
| `pkg/agent/model_resolution.go` | `lookupModelConfigByProtocolModel`; propagates `RPM` into `FallbackCandidate` |
| `pkg/agent/loop.go` | Build `RateLimiterRegistry`, register all agents' candidates, pass to `NewFallbackChain` |

View file

@ -118,9 +118,18 @@ func NewAgentLoop(
) *AgentLoop {
registry := NewAgentRegistry(cfg, provider)
// Set up shared fallback chain
// Set up shared fallback chain with rate limiting.
cooldown := providers.NewCooldownTracker()
fallbackChain := providers.NewFallbackChain(cooldown)
rl := providers.NewRateLimiterRegistry()
// Register rate limiters for all agents' candidates so that RPM limits
// configured in ModelConfig are enforced before each LLM call.
for _, agentID := range registry.ListAgentIDs() {
if agent, ok := registry.GetAgent(agentID); ok {
rl.RegisterCandidates(agent.Candidates)
rl.RegisterCandidates(agent.LightCandidates)
}
}
fallbackChain := providers.NewFallbackChain(cooldown, rl)
// Create state manager using default agent's workspace for channel recording
defaultAgent := registry.GetDefaultAgent()
@ -1000,8 +1009,15 @@ func (al *AgentLoop) ReloadProviderAndConfig(
al.cfg = cfg
al.registry = registry
// Also update fallback chain with new config
al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker())
// Also update fallback chain with new config; rebuild rate limiter registry.
newRL := providers.NewRateLimiterRegistry()
for _, agentID := range registry.ListAgentIDs() {
if agent, ok := registry.GetAgent(agentID); ok {
newRL.RegisterCandidates(agent.Candidates)
newRL.RegisterCandidates(agent.LightCandidates)
}
}
al.fallback = providers.NewFallbackChain(providers.NewCooldownTracker(), newRL)
al.mu.Unlock()

View file

@ -54,7 +54,7 @@ func resolveModelCandidates(
primary string,
fallbacks []string,
) []providers.FallbackCandidate {
return providers.ResolveCandidatesWithLookup(
candidates := providers.ResolveCandidatesWithLookup(
providers.ModelConfig{
Primary: primary,
Fallbacks: fallbacks,
@ -62,6 +62,35 @@ func resolveModelCandidates(
defaultProvider,
buildModelListResolver(cfg),
)
// Propagate RPM from ModelConfig to each candidate.
if cfg != nil {
for i := range candidates {
mc := lookupModelConfigByProtocolModel(cfg, candidates[i].Provider, candidates[i].Model)
if mc != nil && mc.RPM > 0 {
candidates[i].RPM = mc.RPM
}
}
}
return candidates
}
// lookupModelConfigByProtocolModel finds a ModelConfig whose resolved provider/model matches.
func lookupModelConfigByProtocolModel(cfg *config.Config, provider, model string) *config.ModelConfig {
for _, mc := range cfg.ModelList {
if mc == nil || strings.TrimSpace(mc.Model) == "" {
continue
}
p, m := providers.ExtractProtocol(mc.Model)
if p == "" {
p = "openai"
}
if p == provider && m == model {
return mc
}
}
return nil
}
func resolvedCandidateModel(candidates []providers.FallbackCandidate, fallback string) string {

View file

@ -10,12 +10,14 @@ import (
// FallbackChain orchestrates model fallback across multiple candidates.
type FallbackChain struct {
cooldown *CooldownTracker
rl *RateLimiterRegistry
}
// FallbackCandidate represents one model/provider to try.
type FallbackCandidate struct {
Provider string
Model string
RPM int // requests per minute; 0 means unrestricted
}
// FallbackResult contains the successful response and metadata about all attempts.
@ -36,9 +38,10 @@ type FallbackAttempt struct {
Skipped bool // true if skipped due to cooldown
}
// NewFallbackChain creates a new fallback chain with the given cooldown tracker.
func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain {
return &FallbackChain{cooldown: cooldown}
// NewFallbackChain creates a new fallback chain with the given cooldown tracker
// and rate limiter registry.
func NewFallbackChain(cooldown *CooldownTracker, rl *RateLimiterRegistry) *FallbackChain {
return &FallbackChain{cooldown: cooldown, rl: rl}
}
// ResolveCandidates parses model config into a deduplicated candidate list.
@ -136,6 +139,13 @@ func (fc *FallbackChain) Execute(
continue
}
// Enforce per-candidate rate limit before calling the provider.
if fc.rl != nil {
if waitErr := fc.rl.Wait(ctx, cooldownKey); waitErr != nil {
return nil, waitErr
}
}
// Execute the run function.
start := time.Now()
resp, err := run(ctx, candidate.Provider, candidate.Model)
@ -229,6 +239,14 @@ func (fc *FallbackChain) ExecuteImage(
return nil, context.Canceled
}
// Enforce per-candidate rate limit before calling the provider.
imageKey := ModelKey(candidate.Provider, candidate.Model)
if fc.rl != nil {
if waitErr := fc.rl.Wait(ctx, imageKey); waitErr != nil {
return nil, waitErr
}
}
start := time.Now()
resp, err := run(ctx, candidate.Provider, candidate.Model)
elapsed := time.Since(start)

View file

@ -25,7 +25,7 @@ func TestMultiKeyFailover(t *testing.T) {
// Create fallback chain
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Mock run function: first call fails with 429, second succeeds
callCount := 0
@ -82,7 +82,7 @@ func TestMultiKeyFailoverAllFail(t *testing.T) {
candidates := ResolveCandidates(cfg, "zhipu")
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Mock run function: all calls fail with rate limit
callCount := 0
@ -127,7 +127,7 @@ func TestMultiKeyFailoverCooldown(t *testing.T) {
candidates := ResolveCandidates(cfg, "zhipu")
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Put the first model in cooldown (using ModelKey now, not just provider)
cooldownKey := ModelKey(candidates[0].Provider, candidates[0].Model)
@ -183,7 +183,7 @@ func TestMultiKeyFailoverWithFormatError(t *testing.T) {
candidates := ResolveCandidates(cfg, "zhipu")
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Mock run function: first call fails with format error (bad request)
callCount := 0
@ -263,7 +263,7 @@ func TestMultiKeyWithModelFallback(t *testing.T) {
}
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Mock run function: first two fail, third succeeds (model fallback)
callCount := 0
@ -337,7 +337,7 @@ func TestMultiKeyFailoverMixedErrors(t *testing.T) {
candidates := ResolveCandidates(cfg, "zhipu")
cooldown := NewCooldownTracker()
chain := NewFallbackChain(cooldown)
chain := NewFallbackChain(cooldown, nil)
// Mock run function: different errors for each key
callCount := 0

View file

@ -19,7 +19,7 @@ func successRun(content string) func(ctx context.Context, provider, model string
func TestFallback_SingleCandidate_Success(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
result, err := fc.Execute(context.Background(), candidates, successRun("hello"))
@ -36,7 +36,7 @@ func TestFallback_SingleCandidate_Success(t *testing.T) {
func TestFallback_SecondCandidateSuccess(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"),
@ -69,7 +69,7 @@ func TestFallback_SecondCandidateSuccess(t *testing.T) {
func TestFallback_AllFail(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"),
@ -96,7 +96,7 @@ func TestFallback_AllFail(t *testing.T) {
func TestFallback_ContextCanceled(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
ctx, cancel := context.WithCancel(context.Background())
candidates := []FallbackCandidate{
@ -123,7 +123,7 @@ func TestFallback_ContextCanceled(t *testing.T) {
func TestFallback_NonRetriableError(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"),
@ -155,7 +155,7 @@ func TestFallback_NonRetriableError(t *testing.T) {
func TestFallback_CooldownSkip(t *testing.T) {
now := time.Now()
ct, _ := newTestTracker(now)
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
// Put openai/gpt-4 in cooldown (using ModelKey now)
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
@ -193,7 +193,7 @@ func TestFallback_CooldownSkip(t *testing.T) {
func TestFallback_AllInCooldown(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
// Put all models in cooldown (using ModelKey now)
ct.MarkFailure(ModelKey("openai", "gpt-4"), FailoverRateLimit)
@ -221,7 +221,7 @@ func TestFallback_AllInCooldown(t *testing.T) {
func TestFallback_NoCandidates(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
_, err := fc.Execute(context.Background(), nil, successRun("ok"))
if err == nil {
@ -232,7 +232,7 @@ func TestFallback_NoCandidates(t *testing.T) {
func TestFallback_EmptyFallbacks(t *testing.T) {
// Single primary, no fallbacks: should work like direct call
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
result, err := fc.Execute(context.Background(), candidates, successRun("ok"))
@ -246,7 +246,7 @@ func TestFallback_EmptyFallbacks(t *testing.T) {
func TestFallback_UnclassifiedError(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4"),
@ -270,7 +270,7 @@ func TestFallback_UnclassifiedError(t *testing.T) {
func TestFallback_SuccessResetsCooldown(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")}
modelKey := ModelKey("openai", "gpt-4")
@ -297,7 +297,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
func TestImageFallback_Success(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4o")}
result, err := fc.ExecuteImage(context.Background(), candidates, successRun("image result"))
@ -311,7 +311,7 @@ func TestImageFallback_Success(t *testing.T) {
func TestImageFallback_DimensionError(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4o"),
@ -335,7 +335,7 @@ func TestImageFallback_DimensionError(t *testing.T) {
func TestImageFallback_SizeError(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4o"),
@ -359,7 +359,7 @@ func TestImageFallback_SizeError(t *testing.T) {
func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
candidates := []FallbackCandidate{
makeCandidate("openai", "gpt-4o"),
@ -386,7 +386,7 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
func TestImageFallback_NoCandidates(t *testing.T) {
ct := NewCooldownTracker()
fc := NewFallbackChain(ct)
fc := NewFallbackChain(ct, nil)
_, err := fc.ExecuteImage(context.Background(), nil, successRun("ok"))
if err == nil {

View file

@ -0,0 +1,111 @@
package providers
import (
"context"
"sync"
"time"
)
// RateLimiter implements a token-bucket rate limiter for a single key.
// Allows up to RPM requests per minute with a burst equal to RPM.
// Thread-safe.
type RateLimiter struct {
mu sync.Mutex
rpm int
tokens float64
maxBurst float64
lastTick time.Time
nowFunc func() time.Time // for testing
}
// newRateLimiter creates a RateLimiter that allows rpm requests/minute.
func newRateLimiter(rpm int) *RateLimiter {
return &RateLimiter{
rpm: rpm,
tokens: float64(rpm), // start full
maxBurst: float64(rpm),
lastTick: time.Now(),
nowFunc: time.Now,
}
}
// Wait blocks until a token is available or ctx is cancelled.
// Returns ctx.Err() if cancelled while waiting.
func (rl *RateLimiter) Wait(ctx context.Context) error {
for {
rl.mu.Lock()
now := rl.nowFunc()
elapsed := now.Sub(rl.lastTick).Seconds()
rl.lastTick = now
// Refill tokens proportional to elapsed time.
refill := elapsed * float64(rl.rpm) / 60.0
rl.tokens = min(rl.maxBurst, rl.tokens+refill)
if rl.tokens >= 1.0 {
rl.tokens--
rl.mu.Unlock()
return nil
}
// Calculate how long until a token is available.
deficit := 1.0 - rl.tokens
waitSec := deficit / (float64(rl.rpm) / 60.0)
rl.mu.Unlock()
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(time.Duration(waitSec * float64(time.Second))):
// Loop to re-check (another goroutine may have consumed the token).
}
}
}
// RateLimiterRegistry holds per-candidate rate limiters.
// Candidates with RPM=0 are unrestricted.
// Thread-safe for concurrent reads/writes.
type RateLimiterRegistry struct {
mu sync.RWMutex
limiters map[string]*RateLimiter
}
// NewRateLimiterRegistry creates an empty registry.
func NewRateLimiterRegistry() *RateLimiterRegistry {
return &RateLimiterRegistry{
limiters: make(map[string]*RateLimiter),
}
}
// Register adds a rate limiter for the given key at the given RPM.
// If rpm <= 0, no limiter is registered (unrestricted).
func (r *RateLimiterRegistry) Register(key string, rpm int) {
if rpm <= 0 {
return
}
r.mu.Lock()
defer r.mu.Unlock()
r.limiters[key] = newRateLimiter(rpm)
}
// Wait acquires a token for the given key, blocking if needed.
// If no limiter is registered for key, returns immediately.
func (r *RateLimiterRegistry) Wait(ctx context.Context, key string) error {
r.mu.RLock()
rl := r.limiters[key]
r.mu.RUnlock()
if rl == nil {
return nil
}
return rl.Wait(ctx)
}
// RegisterCandidates registers rate limiters for all candidates that have RPM > 0.
// Candidates with RPM == 0 are ignored (no restriction).
func (r *RateLimiterRegistry) RegisterCandidates(candidates []FallbackCandidate) {
for _, c := range candidates {
if c.RPM > 0 {
r.Register(ModelKey(c.Provider, c.Model), c.RPM)
}
}
}

View file

@ -0,0 +1,178 @@
package providers
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
)
// TestRateLimiter_AllowsUpToRPM verifies that up to RPM requests pass immediately
// (burst capacity) and the (RPM+1)-th request is delayed.
func TestRateLimiter_AllowsUpToRPM(t *testing.T) {
rpm := 5
rl := newRateLimiter(rpm)
// All rpm tokens should be available immediately (bucket starts full).
for i := 0; i < rpm; i++ {
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
if err := rl.Wait(ctx); err != nil {
t.Fatalf("request %d should pass immediately, got: %v", i+1, err)
}
cancel()
}
// The next request must wait; cancel it to confirm it blocks.
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
err := rl.Wait(ctx)
if err == nil {
t.Fatal("expected request beyond RPM to block, but it passed immediately")
}
}
// TestRateLimiter_ContextCancellation verifies that a blocked Wait respects cancellation.
func TestRateLimiter_ContextCancellation(t *testing.T) {
rl := newRateLimiter(1)
// Drain the one token.
ctx := context.Background()
if err := rl.Wait(ctx); err != nil {
t.Fatalf("first request failed: %v", err)
}
// Second request should block; cancel it.
cancelCtx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
defer cancel()
err := rl.Wait(cancelCtx)
if err == nil {
t.Fatal("expected cancellation error, got nil")
}
}
// TestRateLimiter_TokenRefill verifies that tokens refill over time.
func TestRateLimiter_TokenRefill(t *testing.T) {
rpm := 60 // 1 token per second
rl := newRateLimiter(rpm)
// Drain all tokens.
for i := 0; i < rpm; i++ {
rl.Wait(context.Background()) //nolint:errcheck
}
// Advance time via nowFunc: simulate 2 seconds passing (should give 2 tokens).
start := time.Now()
rl.nowFunc = func() time.Time { return start.Add(2 * time.Second) }
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
if err := rl.Wait(ctx); err != nil {
t.Fatalf("expected refilled token to be available: %v", err)
}
}
// TestRateLimiterRegistry_NoLimiter verifies that keys without a registered limiter pass freely.
func TestRateLimiterRegistry_NoLimiter(t *testing.T) {
r := NewRateLimiterRegistry()
ctx := context.Background()
for i := 0; i < 100; i++ {
if err := r.Wait(ctx, "unregistered/key"); err != nil {
t.Fatalf("unregistered key should not block: %v", err)
}
}
}
// TestRateLimiterRegistry_ZeroRPM verifies that RPM=0 means no limiter is registered.
func TestRateLimiterRegistry_ZeroRPM(t *testing.T) {
r := NewRateLimiterRegistry()
r.Register("some/key", 0)
ctx := context.Background()
for i := 0; i < 50; i++ {
if err := r.Wait(ctx, "some/key"); err != nil {
t.Fatalf("zero-RPM key should not block: %v", err)
}
}
}
// TestRateLimiterRegistry_Enforcement verifies the registry enforces RPM per key.
func TestRateLimiterRegistry_Enforcement(t *testing.T) {
r := NewRateLimiterRegistry()
r.Register("openai/gpt-4o", 3)
// First 3 calls should pass (burst = RPM).
for i := 0; i < 3; i++ {
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
if err := r.Wait(ctx, "openai/gpt-4o"); err != nil {
t.Fatalf("call %d should pass: %v", i+1, err)
}
cancel()
}
// 4th call should block.
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
if err := r.Wait(ctx, "openai/gpt-4o"); err == nil {
t.Fatal("4th call should have been rate-limited")
}
}
// TestRateLimiterRegistry_RegisterCandidates verifies that RegisterCandidates
// correctly picks up RPM from FallbackCandidate.
func TestRateLimiterRegistry_RegisterCandidates(t *testing.T) {
r := NewRateLimiterRegistry()
candidates := []FallbackCandidate{
{Provider: "openai", Model: "gpt-4o", RPM: 2},
{Provider: "anthropic", Model: "claude-3", RPM: 0}, // no limit
}
r.RegisterCandidates(candidates)
// openai/gpt-4o: 2 tokens burst, 3rd should block.
for i := 0; i < 2; i++ {
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
if err := r.Wait(ctx, "openai/gpt-4o"); err != nil {
t.Fatalf("openai call %d should pass: %v", i+1, err)
}
cancel()
}
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
if err := r.Wait(ctx, "openai/gpt-4o"); err == nil {
t.Fatal("openai 3rd call should have been limited")
}
// anthropic/claude-3: no limit, should always pass.
for i := 0; i < 10; i++ {
if err := r.Wait(context.Background(), "anthropic/claude-3"); err != nil {
t.Fatalf("anthropic call should not be limited: %v", err)
}
}
}
// TestRateLimiter_Concurrency verifies thread safety under concurrent access.
func TestRateLimiter_Concurrency(t *testing.T) {
rpm := 20
rl := newRateLimiter(rpm)
var passed atomic.Int64
var wg sync.WaitGroup
// Launch 30 goroutines; only ~20 should pass immediately.
for i := 0; i < 30; i++ {
wg.Add(1)
go func() {
defer wg.Done()
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
defer cancel()
if rl.Wait(ctx) == nil {
passed.Add(1)
}
}()
}
wg.Wait()
got := passed.Load()
// Allow small timing slack: between rpm-2 and rpm+2.
if got < int64(rpm-2) || got > int64(rpm+2) {
t.Fatalf("expected ~%d immediate passes, got %d", rpm, got)
}
}