fix(agent): preserve per-model identity in rate limiting and fallback
This commit is contained in:
afjcjsbx
parent
3f7cf1a0b6
commit
51d61ac970
10 changed files with 385 additions and 94 deletions
|
|
@ -74,7 +74,9 @@ model_list:
|
|||
- 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.
|
||||
If the current candidate's bucket is empty and there are more candidates available, PicoClaw skips the locally saturated candidate and tries the next fallback immediately. Only the last remaining candidate waits for a token to refill. If the context deadline is hit while waiting on that last candidate, the wait error propagates.
|
||||
|
||||
For `model_list` aliases that resolve to the same underlying provider/model, rate limiting is keyed by the stable config identity (for example `model_name`) rather than the resolved runtime model string. This preserves distinct RPM settings for multi-key and alias-based configurations.
|
||||
|
||||
### Burst behaviour
|
||||
|
||||
|
|
@ -89,5 +91,5 @@ To reduce burstiness for strict APIs, set a lower `rpm` and rely on the steady-s
|
|||
| `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/model_resolution.go` | Resolves candidates from `model_list`, preserving stable config identity and propagating `RPM` into `FallbackCandidate` |
|
||||
| `pkg/agent/loop.go` | Build `RateLimiterRegistry`, register all agents' candidates, pass to `NewFallbackChain` |
|
||||
|
|
|
|||
|
|
@ -165,6 +165,58 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_PreservesDistinctLimiterIdentityForSharedResolvedModel(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
ModelName: "glm-4.7",
|
||||
ModelFallbacks: []string{"glm-4.7__key_1"},
|
||||
},
|
||||
},
|
||||
ModelList: []*config.ModelConfig{
|
||||
{
|
||||
ModelName: "glm-4.7",
|
||||
Model: "zhipu/glm-4.7",
|
||||
RPM: 1,
|
||||
},
|
||||
{
|
||||
ModelName: "glm-4.7__key_1",
|
||||
Model: "zhipu/glm-4.7",
|
||||
RPM: 3,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
|
||||
if len(agent.Candidates) != 2 {
|
||||
t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates))
|
||||
}
|
||||
|
||||
first := agent.Candidates[0]
|
||||
second := agent.Candidates[1]
|
||||
if first.Provider != "zhipu" || first.Model != "glm-4.7" {
|
||||
t.Fatalf("first candidate = %s/%s, want zhipu/glm-4.7", first.Provider, first.Model)
|
||||
}
|
||||
if second.Provider != "zhipu" || second.Model != "glm-4.7" {
|
||||
t.Fatalf("second candidate = %s/%s, want zhipu/glm-4.7", second.Provider, second.Model)
|
||||
}
|
||||
if first.IdentityKey != "model_name:glm-4.7" {
|
||||
t.Fatalf("first identity key = %q, want %q", first.IdentityKey, "model_name:glm-4.7")
|
||||
}
|
||||
if second.IdentityKey != "model_name:glm-4.7__key_1" {
|
||||
t.Fatalf("second identity key = %q, want %q", second.IdentityKey, "model_name:glm-4.7__key_1")
|
||||
}
|
||||
if first.RPM != 1 {
|
||||
t.Fatalf("first RPM = %d, want 1", first.RPM)
|
||||
}
|
||||
if second.RPM != 3 {
|
||||
t.Fatalf("second RPM = %d, want 3", second.RPM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
|
||||
workspace := t.TempDir()
|
||||
mediaDir := media.TempDir()
|
||||
|
|
|
|||
|
|
@ -3453,7 +3453,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
|
|||
return "", fmt.Errorf("failed to initialize model %q: %w", value, err)
|
||||
}
|
||||
|
||||
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, modelCfg.Model, agent.Fallbacks)
|
||||
nextCandidates := resolveModelCandidates(cfg, cfg.Agents.Defaults.Provider, value, agent.Fallbacks)
|
||||
if len(nextCandidates) == 0 {
|
||||
return "", fmt.Errorf("model %q did not resolve to any provider candidates", value)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,8 +8,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
)
|
||||
|
||||
func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool) {
|
||||
ensureProtocol := func(model string) string {
|
||||
func ensureProtocolModel(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
return ""
|
||||
|
|
@ -18,34 +17,93 @@ func buildModelListResolver(cfg *config.Config) func(raw string) (string, bool)
|
|||
return model
|
||||
}
|
||||
return "openai/" + model
|
||||
}
|
||||
|
||||
func modelConfigIdentityKey(mc *config.ModelConfig) string {
|
||||
if mc == nil {
|
||||
return ""
|
||||
}
|
||||
if name := strings.TrimSpace(mc.ModelName); name != "" {
|
||||
return "model_name:" + name
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func candidateFromModelConfig(
|
||||
defaultProvider string,
|
||||
mc *config.ModelConfig,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
if mc == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return func(raw string) (string, bool) {
|
||||
ref := providers.ParseModelRef(ensureProtocolModel(mc.Model), defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
RPM: mc.RPM,
|
||||
IdentityKey: modelConfigIdentityKey(mc),
|
||||
}, true
|
||||
}
|
||||
|
||||
func lookupModelConfigByRef(cfg *config.Config, raw string) *config.ModelConfig {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" || cfg == nil {
|
||||
return "", false
|
||||
return nil
|
||||
}
|
||||
|
||||
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
|
||||
return ensureProtocol(mc.Model), true
|
||||
return mc
|
||||
}
|
||||
|
||||
for i := range cfg.ModelList {
|
||||
fullModel := strings.TrimSpace(cfg.ModelList[i].Model)
|
||||
mc := cfg.ModelList[i]
|
||||
if mc == nil {
|
||||
continue
|
||||
}
|
||||
fullModel := strings.TrimSpace(mc.Model)
|
||||
if fullModel == "" {
|
||||
continue
|
||||
}
|
||||
if fullModel == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
return mc
|
||||
}
|
||||
_, modelID := providers.ExtractProtocol(fullModel)
|
||||
if modelID == raw {
|
||||
return ensureProtocol(fullModel), true
|
||||
return mc
|
||||
}
|
||||
}
|
||||
|
||||
return "", false
|
||||
return nil
|
||||
}
|
||||
|
||||
func resolveModelCandidate(
|
||||
cfg *config.Config,
|
||||
defaultProvider string,
|
||||
raw string,
|
||||
) (providers.FallbackCandidate, bool) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
if mc := lookupModelConfigByRef(cfg, raw); mc != nil {
|
||||
return candidateFromModelConfig(defaultProvider, mc)
|
||||
}
|
||||
|
||||
ref := providers.ParseModelRef(raw, defaultProvider)
|
||||
if ref == nil {
|
||||
return providers.FallbackCandidate{}, false
|
||||
}
|
||||
|
||||
return providers.FallbackCandidate{
|
||||
Provider: ref.Provider,
|
||||
Model: ref.Model,
|
||||
}, true
|
||||
}
|
||||
|
||||
func resolveModelCandidates(
|
||||
|
|
@ -54,45 +112,31 @@ func resolveModelCandidates(
|
|||
primary string,
|
||||
fallbacks []string,
|
||||
) []providers.FallbackCandidate {
|
||||
candidates := providers.ResolveCandidatesWithLookup(
|
||||
providers.ModelConfig{
|
||||
Primary: primary,
|
||||
Fallbacks: fallbacks,
|
||||
},
|
||||
defaultProvider,
|
||||
buildModelListResolver(cfg),
|
||||
)
|
||||
seen := make(map[string]bool)
|
||||
candidates := make([]providers.FallbackCandidate, 0, 1+len(fallbacks))
|
||||
|
||||
// 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
|
||||
addCandidate := func(raw string) {
|
||||
candidate, ok := resolveModelCandidate(cfg, defaultProvider, raw)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
key := candidate.StableKey()
|
||||
if seen[key] {
|
||||
return
|
||||
}
|
||||
seen[key] = true
|
||||
candidates = append(candidates, candidate)
|
||||
}
|
||||
|
||||
addCandidate(primary)
|
||||
for _, fallback := range fallbacks {
|
||||
addCandidate(fallback)
|
||||
}
|
||||
|
||||
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 {
|
||||
if len(candidates) > 0 && strings.TrimSpace(candidates[0].Model) != "" {
|
||||
return candidates[0].Model
|
||||
|
|
|
|||
|
|
@ -208,7 +208,10 @@ func (p *Provider) Chat(
|
|||
if err != nil {
|
||||
// Check for SSO token expiration errors and provide actionable guidance
|
||||
if isSSOTokenError(err) {
|
||||
return nil, fmt.Errorf("bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w", err)
|
||||
return nil, fmt.Errorf(
|
||||
"bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
return nil, fmt.Errorf("bedrock converse: %w", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -584,12 +584,16 @@ func TestIsSSOTokenError(t *testing.T) {
|
|||
},
|
||||
{
|
||||
name: "full SSO error message",
|
||||
err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token"),
|
||||
err: fmt.Errorf(
|
||||
"get identity: get credentials: failed to refresh cached credentials, refresh cached SSO token failed, unable to refresh SSO token",
|
||||
),
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "SSO token file missing",
|
||||
err: fmt.Errorf("get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory"),
|
||||
err: fmt.Errorf(
|
||||
"get identity: get credentials: failed to refresh cached credentials, failed to read cached SSO token file, open ~/.aws/sso/cache/abc123.json: no such file or directory",
|
||||
),
|
||||
expected: true,
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,16 @@ type FallbackCandidate struct {
|
|||
Provider string
|
||||
Model string
|
||||
RPM int // requests per minute; 0 means unrestricted
|
||||
IdentityKey string // optional stable config identity for cooldown/rate limiting
|
||||
}
|
||||
|
||||
// StableKey returns the candidate's config-level identity when available,
|
||||
// otherwise it falls back to the runtime provider/model key.
|
||||
func (c FallbackCandidate) StableKey() string {
|
||||
if key := strings.TrimSpace(c.IdentityKey); key != "" {
|
||||
return key
|
||||
}
|
||||
return ModelKey(c.Provider, c.Model)
|
||||
}
|
||||
|
||||
// FallbackResult contains the successful response and metadata about all attempts.
|
||||
|
|
@ -120,9 +130,9 @@ func (fc *FallbackChain) Execute(
|
|||
return nil, context.Canceled
|
||||
}
|
||||
|
||||
// Check cooldown (per provider/model, not just provider).
|
||||
// This allows multi-key failover where different keys use different model names.
|
||||
cooldownKey := ModelKey(candidate.Provider, candidate.Model)
|
||||
// Check cooldown per stable candidate identity, not just provider/model.
|
||||
// This allows aliases and multi-key configs to fail over independently.
|
||||
cooldownKey := candidate.StableKey()
|
||||
if !fc.cooldown.IsAvailable(cooldownKey) {
|
||||
remaining := fc.cooldown.CooldownRemaining(cooldownKey)
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
|
|
@ -140,11 +150,31 @@ func (fc *FallbackChain) Execute(
|
|||
}
|
||||
|
||||
// Enforce per-candidate rate limit before calling the provider.
|
||||
// If this candidate is locally saturated, try other candidates first.
|
||||
if fc.rl != nil {
|
||||
if !fc.rl.TryAcquire(cooldownKey) {
|
||||
if i < len(candidates)-1 {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: fmt.Errorf("%s waiting for local rate limit token", cooldownKey),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if waitErr := fc.rl.Wait(ctx, cooldownKey); waitErr != nil {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: waitErr,
|
||||
})
|
||||
return nil, waitErr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Execute the run function.
|
||||
start := time.Now()
|
||||
|
|
@ -240,12 +270,32 @@ func (fc *FallbackChain) ExecuteImage(
|
|||
}
|
||||
|
||||
// Enforce per-candidate rate limit before calling the provider.
|
||||
imageKey := ModelKey(candidate.Provider, candidate.Model)
|
||||
// If this candidate is locally saturated, try other candidates first.
|
||||
imageKey := candidate.StableKey()
|
||||
if fc.rl != nil {
|
||||
if !fc.rl.TryAcquire(imageKey) {
|
||||
if i < len(candidates)-1 {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: fmt.Errorf("%s waiting for local rate limit token", imageKey),
|
||||
})
|
||||
continue
|
||||
}
|
||||
if waitErr := fc.rl.Wait(ctx, imageKey); waitErr != nil {
|
||||
result.Attempts = append(result.Attempts, FallbackAttempt{
|
||||
Provider: candidate.Provider,
|
||||
Model: candidate.Model,
|
||||
Skipped: true,
|
||||
Reason: FailoverRateLimit,
|
||||
Error: waitErr,
|
||||
})
|
||||
return nil, waitErr
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := run(ctx, candidate.Provider, candidate.Model)
|
||||
|
|
|
|||
|
|
@ -293,6 +293,42 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
rl := NewRateLimiterRegistry()
|
||||
rl.Register("model_name:primary", 1)
|
||||
if err := rl.Wait(context.Background(), "model_name:primary"); err != nil {
|
||||
t.Fatalf("failed to pre-drain primary limiter: %v", err)
|
||||
}
|
||||
|
||||
fc := NewFallbackChain(ct, rl)
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", IdentityKey: "model_name:primary"},
|
||||
{Provider: "anthropic", Model: "claude", IdentityKey: "model_name:fallback"},
|
||||
}
|
||||
|
||||
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||
if provider != "anthropic" || model != "claude" {
|
||||
t.Fatalf("expected fallback candidate to run, got %s/%s", provider, model)
|
||||
}
|
||||
return &LLMResponse{Content: "fallback ok", FinishReason: "stop"}, nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
result, err := fc.Execute(ctx, candidates, run)
|
||||
if err != nil {
|
||||
t.Fatalf("expected fallback success, got error: %v", err)
|
||||
}
|
||||
if result.Provider != "anthropic" || result.Model != "claude" {
|
||||
t.Fatalf("result = %s/%s, want anthropic/claude", result.Provider, result.Model)
|
||||
}
|
||||
if len(result.Attempts) != 1 || !result.Attempts[0].Skipped {
|
||||
t.Fatalf("expected one skipped primary attempt, got %+v", result.Attempts)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Image Fallback Tests ---
|
||||
|
||||
func TestImageFallback_Success(t *testing.T) {
|
||||
|
|
@ -384,6 +420,42 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestImageFallback_LocalRateLimitSkipsToHealthyFallback(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
rl := NewRateLimiterRegistry()
|
||||
rl.Register("model_name:primary-image", 1)
|
||||
if err := rl.Wait(context.Background(), "model_name:primary-image"); err != nil {
|
||||
t.Fatalf("failed to pre-drain primary image limiter: %v", err)
|
||||
}
|
||||
|
||||
fc := NewFallbackChain(ct, rl)
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", IdentityKey: "model_name:primary-image"},
|
||||
{Provider: "anthropic", Model: "claude-sonnet", IdentityKey: "model_name:fallback-image"},
|
||||
}
|
||||
|
||||
run := func(ctx context.Context, provider, model string) (*LLMResponse, error) {
|
||||
if provider != "anthropic" || model != "claude-sonnet" {
|
||||
t.Fatalf("expected image fallback candidate to run, got %s/%s", provider, model)
|
||||
}
|
||||
return &LLMResponse{Content: "image fallback ok", FinishReason: "stop"}, nil
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
result, err := fc.ExecuteImage(ctx, candidates, run)
|
||||
if err != nil {
|
||||
t.Fatalf("expected image fallback success, got error: %v", err)
|
||||
}
|
||||
if result.Provider != "anthropic" || result.Model != "claude-sonnet" {
|
||||
t.Fatalf("result = %s/%s, want anthropic/claude-sonnet", result.Provider, result.Model)
|
||||
}
|
||||
if len(result.Attempts) != 1 || !result.Attempts[0].Skipped {
|
||||
t.Fatalf("expected one skipped primary attempt, got %+v", result.Attempts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageFallback_NoCandidates(t *testing.T) {
|
||||
ct := NewCooldownTracker()
|
||||
fc := NewFallbackChain(ct, nil)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,15 @@ type RateLimiter struct {
|
|||
nowFunc func() time.Time // for testing
|
||||
}
|
||||
|
||||
func (rl *RateLimiter) refillLocked(now time.Time) {
|
||||
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)
|
||||
}
|
||||
|
||||
// newRateLimiter creates a RateLimiter that allows rpm requests/minute.
|
||||
func newRateLimiter(rpm int) *RateLimiter {
|
||||
return &RateLimiter{
|
||||
|
|
@ -35,12 +44,7 @@ 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)
|
||||
rl.refillLocked(now)
|
||||
|
||||
if rl.tokens >= 1.0 {
|
||||
rl.tokens--
|
||||
|
|
@ -53,15 +57,32 @@ func (rl *RateLimiter) Wait(ctx context.Context) error {
|
|||
waitSec := deficit / (float64(rl.rpm) / 60.0)
|
||||
rl.mu.Unlock()
|
||||
|
||||
timer := time.NewTimer(time.Duration(waitSec * float64(time.Second)))
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if !timer.Stop() {
|
||||
<-timer.C
|
||||
}
|
||||
return ctx.Err()
|
||||
case <-time.After(time.Duration(waitSec * float64(time.Second))):
|
||||
case <-timer.C:
|
||||
// Loop to re-check (another goroutine may have consumed the token).
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TryAcquire attempts to consume a token without blocking.
|
||||
func (rl *RateLimiter) TryAcquire() bool {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
rl.refillLocked(rl.nowFunc())
|
||||
if rl.tokens < 1.0 {
|
||||
return false
|
||||
}
|
||||
rl.tokens--
|
||||
return true
|
||||
}
|
||||
|
||||
// RateLimiterRegistry holds per-candidate rate limiters.
|
||||
// Candidates with RPM=0 are unrestricted.
|
||||
// Thread-safe for concurrent reads/writes.
|
||||
|
|
@ -100,12 +121,24 @@ func (r *RateLimiterRegistry) Wait(ctx context.Context, key string) error {
|
|||
return rl.Wait(ctx)
|
||||
}
|
||||
|
||||
// TryAcquire attempts to consume a token for the given key without blocking.
|
||||
// If no limiter is registered for key, it returns true.
|
||||
func (r *RateLimiterRegistry) TryAcquire(key string) bool {
|
||||
r.mu.RLock()
|
||||
rl := r.limiters[key]
|
||||
r.mu.RUnlock()
|
||||
if rl == nil {
|
||||
return true
|
||||
}
|
||||
return rl.TryAcquire()
|
||||
}
|
||||
|
||||
// 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)
|
||||
r.Register(c.StableKey(), c.RPM)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -149,6 +149,37 @@ func TestRateLimiterRegistry_RegisterCandidates(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestRateLimiterRegistry_RegisterCandidatesUsesStableIdentity(t *testing.T) {
|
||||
r := NewRateLimiterRegistry()
|
||||
candidates := []FallbackCandidate{
|
||||
{Provider: "openai", Model: "gpt-4o", RPM: 1, IdentityKey: "model_name:primary"},
|
||||
{Provider: "openai", Model: "gpt-4o", RPM: 2, IdentityKey: "model_name:fallback"},
|
||||
}
|
||||
r.RegisterCandidates(candidates)
|
||||
|
||||
if err := r.Wait(context.Background(), "model_name:primary"); err != nil {
|
||||
t.Fatalf("primary first call should pass: %v", err)
|
||||
}
|
||||
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||
t.Fatalf("fallback first call should pass: %v", err)
|
||||
}
|
||||
if err := r.Wait(context.Background(), "model_name:fallback"); err != nil {
|
||||
t.Fatalf("fallback second call should pass: %v", err)
|
||||
}
|
||||
|
||||
ctxPrimary, cancelPrimary := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancelPrimary()
|
||||
if err := r.Wait(ctxPrimary, "model_name:primary"); err == nil {
|
||||
t.Fatal("primary second call should have been limited")
|
||||
}
|
||||
|
||||
ctxFallback, cancelFallback := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancelFallback()
|
||||
if err := r.Wait(ctxFallback, "model_name:fallback"); err == nil {
|
||||
t.Fatal("fallback third call should have been limited")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimiter_Concurrency verifies thread safety under concurrent access.
|
||||
func TestRateLimiter_Concurrency(t *testing.T) {
|
||||
rpm := 20
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue