fix(ratelimit): enforce negative token validation before quota lookup
This commit is contained in:
parent
2ad4870bd0
commit
781e0ee3d6
1 changed files with 23 additions and 16 deletions
39
ratelimit.go
39
ratelimit.go
|
|
@ -249,7 +249,12 @@ func (rl *RateLimiter) Load() error {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
return yaml.Unmarshal([]byte(content), rl)
|
if err := yaml.Unmarshal([]byte(content), rl); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
ensureMaps(rl)
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// loadSQLite reads quotas and state from the SQLite backend.
|
// loadSQLite reads quotas and state from the SQLite backend.
|
||||||
|
|
@ -596,6 +601,14 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
decision := Decision{}
|
decision := Decision{}
|
||||||
|
|
||||||
|
if estimatedTokens < 0 {
|
||||||
|
decision.Allowed = false
|
||||||
|
decision.Code = DecisionInvalidTokens
|
||||||
|
decision.Reason = "estimated tokens must be non-negative"
|
||||||
|
decision.Stats = rl.snapshotLocked(model)
|
||||||
|
return decision
|
||||||
|
}
|
||||||
|
|
||||||
quota, ok := rl.Quotas[model]
|
quota, ok := rl.Quotas[model]
|
||||||
if !ok {
|
if !ok {
|
||||||
decision.Allowed = true
|
decision.Allowed = true
|
||||||
|
|
@ -606,13 +619,6 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
|
||||||
}
|
}
|
||||||
|
|
||||||
if quota.MaxRPM == 0 && quota.MaxTPM == 0 && quota.MaxRPD == 0 {
|
if quota.MaxRPM == 0 && quota.MaxTPM == 0 && quota.MaxRPD == 0 {
|
||||||
if estimatedTokens < 0 {
|
|
||||||
decision.Allowed = false
|
|
||||||
decision.Code = DecisionInvalidTokens
|
|
||||||
decision.Reason = "estimated tokens must be non-negative"
|
|
||||||
decision.Stats = rl.snapshotLocked(model)
|
|
||||||
return decision
|
|
||||||
}
|
|
||||||
decision.Allowed = true
|
decision.Allowed = true
|
||||||
decision.Code = DecisionUnlimited
|
decision.Code = DecisionUnlimited
|
||||||
decision.Reason = "all limits are unlimited"
|
decision.Reason = "all limits are unlimited"
|
||||||
|
|
@ -627,14 +633,6 @@ func (rl *RateLimiter) Decide(model string, estimatedTokens int) Decision {
|
||||||
rl.State[model] = stats
|
rl.State[model] = stats
|
||||||
}
|
}
|
||||||
|
|
||||||
if estimatedTokens < 0 {
|
|
||||||
decision.Allowed = false
|
|
||||||
decision.Code = DecisionInvalidTokens
|
|
||||||
decision.Reason = "estimated tokens must be non-negative"
|
|
||||||
decision.Stats = rl.snapshotLocked(model)
|
|
||||||
return decision
|
|
||||||
}
|
|
||||||
|
|
||||||
decision.Stats = rl.snapshotLocked(model)
|
decision.Stats = rl.snapshotLocked(model)
|
||||||
|
|
||||||
if quota.MaxRPD > 0 && stats.DayCount >= quota.MaxRPD {
|
if quota.MaxRPD > 0 && stats.DayCount >= quota.MaxRPD {
|
||||||
|
|
@ -836,6 +834,15 @@ func newConfiguredRateLimiter(cfg Config) *RateLimiter {
|
||||||
return rl
|
return rl
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func ensureMaps(rl *RateLimiter) {
|
||||||
|
if rl.Quotas == nil {
|
||||||
|
rl.Quotas = make(map[string]ModelQuota)
|
||||||
|
}
|
||||||
|
if rl.State == nil {
|
||||||
|
rl.State = make(map[string]*UsageStats)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func applyConfig(rl *RateLimiter, cfg Config) {
|
func applyConfig(rl *RateLimiter, cfg Config) {
|
||||||
profiles := DefaultProfiles()
|
profiles := DefaultProfiles()
|
||||||
providers := cfg.Providers
|
providers := cfg.Providers
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue