diff --git a/docs/rate-limiting.md b/docs/rate-limiting.md new file mode 100644 index 000000000..f2403ca73 --- /dev/null +++ b/docs/rate-limiting.md @@ -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` | diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index d7461e76f..7fcc117f3 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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() diff --git a/pkg/agent/model_resolution.go b/pkg/agent/model_resolution.go index 140cff718..fcae29bcc 100644 --- a/pkg/agent/model_resolution.go +++ b/pkg/agent/model_resolution.go @@ -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 { diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index 549ec7837..b9e237752 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -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) diff --git a/pkg/providers/fallback_multikey_test.go b/pkg/providers/fallback_multikey_test.go index 9ed8fa73c..10481ec61 100644 --- a/pkg/providers/fallback_multikey_test.go +++ b/pkg/providers/fallback_multikey_test.go @@ -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 diff --git a/pkg/providers/fallback_test.go b/pkg/providers/fallback_test.go index 1a1118e33..bd0d56e12 100644 --- a/pkg/providers/fallback_test.go +++ b/pkg/providers/fallback_test.go @@ -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 { diff --git a/pkg/providers/ratelimiter.go b/pkg/providers/ratelimiter.go new file mode 100644 index 000000000..d95f0feac --- /dev/null +++ b/pkg/providers/ratelimiter.go @@ -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) + } + } +} diff --git a/pkg/providers/ratelimiter_test.go b/pkg/providers/ratelimiter_test.go new file mode 100644 index 000000000..12734431d --- /dev/null +++ b/pkg/providers/ratelimiter_test.go @@ -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) + } +}