mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-13 04:58:50 +00:00
Re-review round on 5cc6814 found round 1's 48h fallback clamp
defective: it broke TTL > window-length for the MONTH window on
memoryStore, turning tokensPerMonth/costPerMonthUSD into a rolling
~48h budget (fail-open) whenever Redis is absent or down.
1. counterIncr gains enforceTTL time.Duration — the minimum retention
window's own enforcement correctness requires (enforceTTLFor:
min->2m, hour->2h, day->25h, month->32*24h, matching each window's
pre-history-bump TTL). newCounterIncr builds every checkAndCount/
account entry through one helper so a (window, ttl) pair can never
reach incrMulti without its matching floor. memoryStore's clampTTL
now takes (ttl, enforceTTL) and applies max(min(ttl,
memoryStoreMaxTTL), enforceTTL): the 48h ceiling still bounds
min/hour/day (unaffected/still-capped as before), but the floor
wins for month, landing it back at ~32 days (1-2 live keys per
scope/metric — negligible against the day-key savings that
motivated the ceiling). redisStore ignores enforceTTL; incrBy's own
single-key path passes enforceTTL=0 (ceiling only, unchanged
behavior — it has no window to derive a floor from and has no
production caller since the batching, see (3) below).
2. limits.go's TTL-invariant comment restated precisely: the consts
still satisfy TTL > window-length exactly as applied to Redis; that
did not automatically extend to memoryStore's own effective TTL
after its ceiling, which is what round 1 broke and round 2 fixes.
README: the Limits and accounting "Fallback" bullet now says
enforcement stays correct on the fallback at every window
(including tokensPerMonth/costPerMonthUSD), and that the 48h
ceiling applies to usage-history CHARTING retention only.
LimitsConfig's window table (tokensPerDay/tokensPerMonth rows)
cross-links it. The "resets on a rebuild" bullet's stale
limiter.storeIncrBy reference is corrected to storeIncrMulti and
now distinguishes rebuild-loss from the fallback's own per-window
TTL expiry.
3. counterStore.incrBy, limiter.storeIncrBy, and limiter.incrCounter's
doc comments now state plainly they have no production caller since
the batching (perf review, 2026-08-21) and are kept for direct
counter seeding in tests — 5 call sites, all in admin_test.go's
TestAdminUsage_MathAgainstSeededCounters, correcting the prior fix
report's "dozens of tests" overclaim to the actual count.
Tests: TestClampTTL_ExactMathTable pins the ceiling-vs-floor
precedence for every window's real (ttl, enforceTTL) pair.
TestEnforceTTLFor_MatchesTheOriginalPreHistoryBumpValues pins the four
floor values directly. TestMemoryStore_IncrMulti_
MonthCounterSurvivesPast48hAndStillEnforces drives a real
TokensPerMonth budget through account/checkAndCount on the fallback
5 days in and confirms both survival and enforcement.
TestMemoryStore_IncrMulti_DayCounterStillClampsAt48h confirms day
still resets at the 48h ceiling (its own floor sits below it).
Gate: go test ./... -count=1, -race, gofmt, vet, make yaegi-check —
all green.
1377 lines
54 KiB
Go
1377 lines
54 KiB
Go
package traefikllmgateway
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestLimiterRequestWindow is the brief's Step-1 case: a per-minute
|
|
// request limit blocks the third request in a window and resets on the
|
|
// next window.
|
|
func TestLimiterRequestWindow(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC)
|
|
l := newLimiter(nil, true) // nil store → uses fallback memoryStore
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "a", limits: &LimitsConfig{RequestsPerMinute: 2}}}
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("1st should pass")
|
|
}
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("2nd should pass")
|
|
}
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("3rd should violate")
|
|
}
|
|
now = now.Add(time.Minute) // next window
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("new window should pass")
|
|
}
|
|
}
|
|
|
|
// TestLimiterTokenBudget is the brief's Step-1 case: accounting tokens
|
|
// past a day budget makes the next check refuse.
|
|
func TestLimiterTokenBudget(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerDay: 100}}}
|
|
l.account(scopes, usage{prompt: 60, completion: 50}, 0) // 110 > 100
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("over token budget must refuse")
|
|
}
|
|
}
|
|
|
|
// TestLimiterCostBudget mirrors TestLimiterTokenBudget for a cost budget:
|
|
// accounting micro-USD cost past a day budget makes the next check refuse,
|
|
// and the violation carries the day window's Retry-After.
|
|
func TestLimiterCostBudget(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerDayUSD: 1.00}}}
|
|
|
|
l.account(scopes, usage{}, 1_500_000) // $1.50 > $1.00 budget
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("over cost budget must refuse")
|
|
}
|
|
wantRetry := int(time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC).Sub(now).Seconds())
|
|
if v.retryAfter != wantRetry {
|
|
t.Errorf("retryAfter = %d, want %d (seconds to next UTC midnight)", v.retryAfter, wantRetry)
|
|
}
|
|
if v.storeDown {
|
|
t.Error("storeDown must be false in this task")
|
|
}
|
|
}
|
|
|
|
// TestLimiterNilLimitsUnlimited asserts a scope with limits == nil never
|
|
// blocks a request, however many times it is checked — request counters
|
|
// still climb unboundedly but nothing evaluates them.
|
|
func TestLimiterNilLimitsUnlimited(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
scopes := []limitScope{{kind: "user", id: "unlimited", limits: nil}}
|
|
for i := 0; i < 1000; i++ {
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("iteration %d: nil limits must never violate, got %+v", i, v)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestLimiterUserOverridesAlongsideGroup asserts checkAndCount evaluates
|
|
// every scope it is given in one call: a tight user limit blocks the
|
|
// request even though the group scope in the same call is unlimited, and
|
|
// an unlimited user scope still lets a tight group limit block.
|
|
func TestLimiterUserOverridesAlongsideGroup(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
|
|
t.Run("tight user limit blocks despite unlimited group", func(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{
|
|
{kind: "user", id: "u1", limits: &LimitsConfig{RequestsPerMinute: 1}},
|
|
{kind: "group", id: "g1", limits: nil},
|
|
}
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("1st request should pass")
|
|
}
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("2nd request should violate the user's tight limit")
|
|
}
|
|
})
|
|
|
|
t.Run("tight group limit blocks despite unlimited user", func(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{
|
|
{kind: "user", id: "u2", limits: nil},
|
|
{kind: "group", id: "g2", limits: &LimitsConfig{RequestsPerMinute: 1}},
|
|
}
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("1st request should pass")
|
|
}
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("2nd request should violate the group's tight limit")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestLimiterRetryAfter checks the Retry-After computation directly for
|
|
// each window kind: seconds to the next minute boundary, next UTC
|
|
// midnight, and the first of next month UTC.
|
|
func TestLimiterRetryAfter(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
now time.Time
|
|
window string
|
|
want int
|
|
}{
|
|
{"minute boundary", time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC), windowMin, 30},
|
|
{"day boundary", time.Date(2026, 8, 20, 23, 59, 0, 0, time.UTC), windowDay, 60},
|
|
{"month boundary, 31-day month", time.Date(2026, 8, 31, 23, 59, 59, 0, time.UTC), windowMonth, 1},
|
|
{"month boundary, mid-month", time.Date(2026, 2, 15, 0, 0, 0, 0, time.UTC), windowMonth, int(14 * 24 * time.Hour / time.Second)}, // Feb 2026 (not a leap year) runs Feb1-Mar1 = 28 days; Feb15 is 14 days before Mar1
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
if got := retryAfterSeconds(c.now, c.window); got != c.want {
|
|
t.Errorf("retryAfterSeconds(%v, %q) = %d, want %d", c.now, c.window, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestLimiterRequestLimitRetryAfter asserts a request-limit violation
|
|
// (not a budget violation) also carries the correct Retry-After for its
|
|
// window.
|
|
func TestLimiterRequestLimitRetryAfter(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "a", limits: &LimitsConfig{RequestsPerMinute: 1}}}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("1st should pass")
|
|
}
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("2nd should violate")
|
|
}
|
|
if v.retryAfter != 30 {
|
|
t.Errorf("retryAfter = %d, want 30", v.retryAfter)
|
|
}
|
|
}
|
|
|
|
// TestWindowKeyFormat pins the exact key layout every store operation
|
|
// depends on: llmgw:{kind}:{id}:{metric}:{window}:{bucket}, with bucket
|
|
// formatted from t.UTC() at each window's granularity — including hour
|
|
// (v0.2 data-layer task).
|
|
func TestWindowKeyFormat(t *testing.T) {
|
|
tm := time.Date(2026, 8, 20, 10, 4, 59, 0, time.FixedZone("UTC+1", 3600)) // 09:04:59 UTC
|
|
cases := []struct {
|
|
window string
|
|
want string
|
|
}{
|
|
{windowHour, "llmgw:user:a:req:hour:2026082009"},
|
|
{windowMin, "llmgw:user:a:req:min:202608200904"},
|
|
{windowDay, "llmgw:user:a:req:day:20260820"},
|
|
{windowMonth, "llmgw:user:a:req:month:202608"},
|
|
}
|
|
for _, c := range cases {
|
|
if got := windowKey("user", "a", "req", c.window, tm); got != c.want {
|
|
t.Errorf("windowKey(..., %q, ...) = %q, want %q", c.window, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestWindowKeyFormat_HourUTCBoundary pins the hour bucket rolling over
|
|
// exactly at the UTC hour boundary, and proves a non-UTC input is
|
|
// converted to UTC first (an input at 00:04:59 in a UTC+1 zone is 23:04:59
|
|
// UTC the PREVIOUS day — a bug converting the hour field alone without
|
|
// re-deriving the date would produce "...23" glued onto the wrong day).
|
|
func TestWindowKeyFormat_HourUTCBoundary(t *testing.T) {
|
|
justBefore := time.Date(2026, 8, 20, 8, 59, 59, 0, time.UTC)
|
|
justAfter := time.Date(2026, 8, 20, 9, 0, 0, 0, time.UTC)
|
|
if got := windowKey("user", "a", "req", windowHour, justBefore); got != "llmgw:user:a:req:hour:2026082008" {
|
|
t.Errorf("just before the hour boundary: got %q", got)
|
|
}
|
|
if got := windowKey("user", "a", "req", windowHour, justAfter); got != "llmgw:user:a:req:hour:2026082009" {
|
|
t.Errorf("just after the hour boundary: got %q", got)
|
|
}
|
|
|
|
// 00:04:59 in UTC+1 is 23:04:59 UTC the previous day.
|
|
crossMidnight := time.Date(2026, 8, 21, 0, 4, 59, 0, time.FixedZone("UTC+1", 3600))
|
|
if got := windowKey("user", "a", "req", windowHour, crossMidnight); got != "llmgw:user:a:req:hour:2026082023" {
|
|
t.Errorf("cross-midnight non-UTC input: got %q, want the previous UTC day's hour 23 bucket", got)
|
|
}
|
|
}
|
|
|
|
// TestWindowEnd_HourBoundary asserts windowEnd's new hour case (v0.2
|
|
// data-layer task) returns the next UTC hour boundary, and that a value
|
|
// exactly on the boundary reports zero seconds remaining in that boundary
|
|
// (retryAfterSeconds rounds a zero duration up to 0, distinct from
|
|
// rounding a fractional second up to 1 — see TestLimiterRetryAfter for
|
|
// the day/month/min cases this mirrors).
|
|
func TestWindowEnd_HourBoundary(t *testing.T) {
|
|
t.Run("mid-hour", func(t *testing.T) {
|
|
got := windowEnd(time.Date(2026, 8, 20, 9, 30, 0, 0, time.UTC), windowHour)
|
|
want := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
if !got.Equal(want) {
|
|
t.Errorf("windowEnd = %v, want %v", got, want)
|
|
}
|
|
})
|
|
t.Run("exactly on the boundary", func(t *testing.T) {
|
|
got := windowEnd(time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC), windowHour)
|
|
want := time.Date(2026, 8, 20, 11, 0, 0, 0, time.UTC)
|
|
if !got.Equal(want) {
|
|
t.Errorf("windowEnd = %v, want %v", got, want)
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestHistoryStepBack_HourAndDay asserts the fixed-duration step-back
|
|
// cases (hour, day) land i units before now exactly, including across a
|
|
// UTC day boundary.
|
|
func TestHistoryStepBack_HourAndDay(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 1, 30, 0, 0, time.UTC)
|
|
|
|
if got := historyStepBack(now, windowHour, 0); !got.Equal(now) {
|
|
t.Errorf("i=0 must be now itself, got %v", got)
|
|
}
|
|
wantHour := time.Date(2026, 8, 19, 23, 30, 0, 0, time.UTC) // 2 hours back, crossing midnight
|
|
if got := historyStepBack(now, windowHour, 2); !got.Equal(wantHour) {
|
|
t.Errorf("i=2 hours back = %v, want %v", got, wantHour)
|
|
}
|
|
|
|
wantDay := time.Date(2026, 8, 17, 1, 30, 0, 0, time.UTC)
|
|
if got := historyStepBack(now, windowDay, 3); !got.Equal(wantDay) {
|
|
t.Errorf("i=3 days back = %v, want %v", got, wantDay)
|
|
}
|
|
}
|
|
|
|
// TestHistoryStepBack_MonthAnchorsOnDayOne is the month-arithmetic trap
|
|
// TestLimiterRetryAfter's own month case does not exercise: stepping back
|
|
// from a high day-of-month (the 31st) with a naive
|
|
// now.AddDate(0, -i, 0) overflows a shorter target month (e.g. "Mar 31"
|
|
// minus one month normalizes to "Mar 3", not "Feb 28") and would land in
|
|
// the WRONG calendar month's bucket. historyStepBack avoids this by
|
|
// anchoring on day 1 first — this test drives it from Mar 31 and Jan 31
|
|
// specifically because those are the inputs that would expose a
|
|
// regression back to the naive form.
|
|
func TestHistoryStepBack_MonthAnchorsOnDayOne(t *testing.T) {
|
|
cases := []struct {
|
|
now time.Time
|
|
name string
|
|
want string
|
|
i int
|
|
}{
|
|
{name: "Mar 31 minus 1 month = Feb", now: time.Date(2026, 3, 31, 12, 0, 0, 0, time.UTC), i: 1, want: "202602"},
|
|
{name: "Jan 31 minus 1 month = Dec (year rollover)", now: time.Date(2026, 1, 31, 0, 0, 0, 0, time.UTC), i: 1, want: "202512"},
|
|
{name: "May 31 minus 3 months = Feb", now: time.Date(2026, 5, 31, 0, 0, 0, 0, time.UTC), i: 3, want: "202602"},
|
|
{name: "i=0 is now's own month", now: time.Date(2026, 3, 31, 0, 0, 0, 0, time.UTC), i: 0, want: "202603"},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
got := bucketFor(historyStepBack(c.now, windowMonth, c.i), windowMonth)
|
|
if got != c.want {
|
|
t.Errorf("bucketFor(historyStepBack(%v, month, %d)) = %q, want %q", c.now, c.i, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestHistoryBucketKeys_OldestFirstInclusiveOfCurrent asserts
|
|
// historyBucketKeys returns span buckets oldest-first, with the LAST
|
|
// entry being now's own (current, possibly partial) bucket — the
|
|
// "inclusive of current bucket" contract GET /admin/api/usage/history
|
|
// promises.
|
|
func TestHistoryBucketKeys_OldestFirstInclusiveOfCurrent(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 14, 0, 0, 0, time.UTC)
|
|
keys, buckets := historyBucketKeys("total", "all", "req", windowHour, now, 4)
|
|
|
|
wantBuckets := []string{"2026082011", "2026082012", "2026082013", "2026082014"}
|
|
if len(buckets) != len(wantBuckets) {
|
|
t.Fatalf("len(buckets) = %d, want %d", len(buckets), len(wantBuckets))
|
|
}
|
|
for i, want := range wantBuckets {
|
|
if buckets[i] != want {
|
|
t.Errorf("buckets[%d] = %q, want %q", i, buckets[i], want)
|
|
}
|
|
wantKey := "llmgw:total:all:req:hour:" + want
|
|
if keys[i] != wantKey {
|
|
t.Errorf("keys[%d] = %q, want %q", i, keys[i], wantKey)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestLimiterHistory_OneBatchCallOldestFirst drives limiter.history
|
|
// against a counting stub store: it must issue exactly ONE getMulti call
|
|
// for the whole span, and the returned points must be oldest-first with
|
|
// the seeded values in the right buckets.
|
|
func TestLimiterHistory_OneBatchCallOldestFirst(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
countStore := &historyCountingStore{values: map[string]int64{
|
|
windowKey("total", "all", metricReq, windowHour, now.Add(-2*time.Hour)): 5,
|
|
windowKey("total", "all", metricReq, windowHour, now): 9,
|
|
}}
|
|
l := newLimiter(countStore, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
points, ok := l.history("total", "all", "req", windowHour, now, 3)
|
|
if !ok {
|
|
t.Fatal("want ok=true")
|
|
}
|
|
if countStore.getMultiCalls != 1 {
|
|
t.Errorf("getMultiCalls = %d, want 1", countStore.getMultiCalls)
|
|
}
|
|
if len(points) != 3 {
|
|
t.Fatalf("len(points) = %d, want 3", len(points))
|
|
}
|
|
if points[0].value != 5 || points[0].bucket != bucketFor(now.Add(-2*time.Hour), windowHour) {
|
|
t.Errorf("points[0] (oldest) = %+v, want value 5 at the -2h bucket", points[0])
|
|
}
|
|
if points[1].value != 0 {
|
|
t.Errorf("points[1] (unseeded -1h bucket) = %+v, want value 0", points[1])
|
|
}
|
|
if points[2].value != 9 || points[2].bucket != bucketFor(now, windowHour) {
|
|
t.Errorf("points[2] (newest, current bucket) = %+v, want value 9", points[2])
|
|
}
|
|
}
|
|
|
|
// historyCountingStore is a counterStore stub whose getMulti serves fixed
|
|
// values and counts its own calls — mirrors admin_test.go's
|
|
// countingMultiStore, defined separately here since limits_test.go must
|
|
// not depend on admin_test.go's test-only types.
|
|
type historyCountingStore struct {
|
|
values map[string]int64
|
|
getMultiCalls int
|
|
}
|
|
|
|
func (s *historyCountingStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
|
|
func (s *historyCountingStore) get(string) (int64, error) { return 0, nil }
|
|
func (s *historyCountingStore) getMulti(keys []string) ([]int64, error) {
|
|
s.getMultiCalls++
|
|
out := make([]int64, len(keys))
|
|
for i, k := range keys {
|
|
out[i] = s.values[k]
|
|
}
|
|
return out, nil
|
|
}
|
|
func (s *historyCountingStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
|
return make([]int64, len(entries)), nil
|
|
}
|
|
|
|
// TestLimiterHistory_StoreDown asserts limiter.history reports ok=false
|
|
// when the configured store errors and failOpen is false — the caller
|
|
// (GET /admin/api/usage/history) must answer 503, never a silently-zero
|
|
// series.
|
|
func TestLimiterHistory_StoreDown(t *testing.T) {
|
|
l := newLimiter(alwaysErrStore{}, false)
|
|
points, ok := l.history("total", "all", "req", windowDay, time.Now(), 5)
|
|
if ok {
|
|
t.Fatal("want ok=false when the configured store errors and failOpen is false")
|
|
}
|
|
if points != nil {
|
|
t.Errorf("points = %+v, want nil on a storeDown read", points)
|
|
}
|
|
}
|
|
|
|
// TestWithTotalScope_AppendsUnlimitedTotalScope asserts withTotalScope
|
|
// appends exactly one {kind: totalScopeKind, id: totalScopeID, limits:
|
|
// nil} scope after whatever buildLimitScopes already returned.
|
|
func TestWithTotalScope_AppendsUnlimitedTotalScope(t *testing.T) {
|
|
base := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 1}}}
|
|
got := withTotalScope(base)
|
|
|
|
if len(got) != 2 {
|
|
t.Fatalf("len(got) = %d, want 2", len(got))
|
|
}
|
|
if got[0].kind != "user" || got[0].id != "u" {
|
|
t.Errorf("got[0] = %+v, want the original user scope untouched", got[0])
|
|
}
|
|
total := got[1]
|
|
if total.kind != totalScopeKind || total.id != totalScopeID || total.limits != nil {
|
|
t.Errorf("got[1] = %+v, want kind=%q id=%q limits=nil", total, totalScopeKind, totalScopeID)
|
|
}
|
|
}
|
|
|
|
// TestWithTotalScope_CountedButNeverEvaluated drives checkAndCount many
|
|
// times against a scopes slice built via withTotalScope: the total
|
|
// scope's own limits are always nil, so no volume of traffic can ever
|
|
// make it violate — evaluateScope's nil-limits check skips it — while its
|
|
// req counter still climbs by exactly one per call, proving
|
|
// checkAndCount's counting loop does not skip a nil-limits scope the way
|
|
// evaluation does.
|
|
func TestWithTotalScope_CountedButNeverEvaluated(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
now := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
scopes := withTotalScope([]limitScope{{kind: "user", id: "u", limits: nil}})
|
|
const n = 250
|
|
for i := 0; i < n; i++ {
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("iteration %d: total scope must never violate (limits always nil), got %+v", i, v)
|
|
}
|
|
}
|
|
|
|
total, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowDay, now)
|
|
if !ok || total != n {
|
|
t.Errorf("total req:day = %d, ok=%v, want %d (counted every call despite never evaluating)", total, ok, n)
|
|
}
|
|
totalHour, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowHour, now)
|
|
if !ok || totalHour != n {
|
|
t.Errorf("total req:hour = %d, ok=%v, want %d", totalHour, ok, n)
|
|
}
|
|
}
|
|
|
|
// TestWithTotalScope_AccountWritesHourBuckets asserts account (v0.2
|
|
// data-layer task) writes the total scope's hour-window tokin/tokout/cost
|
|
// counters, not just day/month — checkAndCount's own hour write is
|
|
// covered by TestWithTotalScope_CountedButNeverEvaluated above.
|
|
func TestWithTotalScope_AccountWritesHourBuckets(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
now := time.Date(2026, 8, 20, 12, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
scopes := withTotalScope([]limitScope{{kind: "user", id: "u", limits: nil}})
|
|
l.account(scopes, usage{prompt: 4, completion: 6}, 700)
|
|
|
|
tokIn, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokIn, windowHour, now)
|
|
if !ok || tokIn != 4 {
|
|
t.Errorf("total tokin:hour = %d, ok=%v, want 4", tokIn, ok)
|
|
}
|
|
tokOut, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowHour, now)
|
|
if !ok || tokOut != 6 {
|
|
t.Errorf("total tokout:hour = %d, ok=%v, want 6", tokOut, ok)
|
|
}
|
|
cost, ok := l.getCounter(totalScopeKind, totalScopeID, metricCost, windowHour, now)
|
|
if !ok || cost != 700 {
|
|
t.Errorf("total cost:hour = %d, ok=%v, want 700", cost, ok)
|
|
}
|
|
}
|
|
|
|
// TestMemoryStoreIncrAndGet exercises the fallback store directly: fresh
|
|
// keys start at n, repeated incrBy accumulates, and get on an unknown key
|
|
// returns 0 with no error.
|
|
func TestMemoryStoreIncrAndGet(t *testing.T) {
|
|
m := newMemoryStore()
|
|
|
|
v, err := m.incrBy("k", 5, time.Minute)
|
|
if err != nil || v != 5 {
|
|
t.Fatalf("incrBy = %d, %v, want 5, nil", v, err)
|
|
}
|
|
v, err = m.incrBy("k", 3, time.Minute)
|
|
if err != nil || v != 8 {
|
|
t.Fatalf("incrBy = %d, %v, want 8, nil", v, err)
|
|
}
|
|
|
|
got, err := m.get("k")
|
|
if err != nil || got != 8 {
|
|
t.Fatalf("get = %d, %v, want 8, nil", got, err)
|
|
}
|
|
|
|
got, err = m.get("missing")
|
|
if err != nil || got != 0 {
|
|
t.Fatalf("get missing = %d, %v, want 0, nil", got, err)
|
|
}
|
|
}
|
|
|
|
// TestMemoryStoreExpiry asserts a counter reset to 0 (as a fresh key) once
|
|
// its TTL has elapsed, rather than continuing to accumulate forever.
|
|
func TestMemoryStoreExpiry(t *testing.T) {
|
|
m := newMemoryStore()
|
|
if _, err := m.incrBy("k", 5, time.Nanosecond); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
|
|
got, err := m.get("k")
|
|
if err != nil || got != 0 {
|
|
t.Fatalf("get after expiry = %d, %v, want 0, nil", got, err)
|
|
}
|
|
|
|
v, err := m.incrBy("k", 2, time.Minute)
|
|
if err != nil || v != 2 {
|
|
t.Fatalf("incrBy after expiry = %d, %v, want 2, nil (fresh counter, not 7)", v, err)
|
|
}
|
|
}
|
|
|
|
// TestMemoryStore_TTLClampedToMax is the perf-review (2026-08-21) case:
|
|
// a TTL beyond memoryStoreMaxTTL (48h) — e.g. dayWindowTTL at 35 days or
|
|
// monthWindowTTL at 400 days — must be clamped down to 48h before it is
|
|
// ever applied, so the fallback's own live-key count stays bounded
|
|
// independent of which window a caller writes.
|
|
func TestMemoryStore_TTLClampedToMax(t *testing.T) {
|
|
m := newMemoryStore()
|
|
fake := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC)
|
|
m.nowFn = func() time.Time { return fake }
|
|
|
|
cases := []struct {
|
|
name string
|
|
ttl time.Duration
|
|
}{
|
|
{"dayWindowTTL (35d)", dayWindowTTL},
|
|
{"monthWindowTTL (400d)", monthWindowTTL},
|
|
{"an arbitrary huge ttl", 10 * 365 * 24 * time.Hour},
|
|
}
|
|
for i, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
key := fmt.Sprintf("k%d", i)
|
|
if _, err := m.incrBy(key, 1, c.ttl); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m.mu.Lock()
|
|
got := m.data[key].expiry
|
|
m.mu.Unlock()
|
|
want := fake.Add(memoryStoreMaxTTL)
|
|
if !got.Equal(want) {
|
|
t.Errorf("expiry = %v, want %v (clamped to memoryStoreMaxTTL, not the requested %v)", got, want, c.ttl)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestMemoryStore_TTLBelowMaxIsUnaffected asserts the clamp is a ceiling,
|
|
// not a rewrite: a TTL already at or under memoryStoreMaxTTL passes
|
|
// through exactly as requested — hourWindowTTL (48h) sits precisely at
|
|
// the clamp boundary and must round-trip unchanged, not get pushed under
|
|
// it by an off-by-one comparison.
|
|
func TestMemoryStore_TTLBelowMaxIsUnaffected(t *testing.T) {
|
|
m := newMemoryStore()
|
|
fake := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC)
|
|
m.nowFn = func() time.Time { return fake }
|
|
|
|
cases := []struct {
|
|
name string
|
|
ttl time.Duration
|
|
}{
|
|
{"minWindowTTL (2m)", minWindowTTL},
|
|
{"hourWindowTTL (48h, exactly at the clamp)", hourWindowTTL},
|
|
}
|
|
for i, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
key := fmt.Sprintf("k%d", i)
|
|
if _, err := m.incrBy(key, 1, c.ttl); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
m.mu.Lock()
|
|
got := m.data[key].expiry
|
|
m.mu.Unlock()
|
|
want := fake.Add(c.ttl)
|
|
if !got.Equal(want) {
|
|
t.Errorf("expiry = %v, want %v (ttl under the clamp must pass through unchanged)", got, want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestMemoryStoreSweepTimeGated asserts the sweep gate is purely
|
|
// time-based (sweepEvery), independent of how many keys memoryStore
|
|
// holds: an incrBy call before the gate elapses leaves an already-expired
|
|
// key physically in place (no scan happened), and the first incrBy at or
|
|
// past the gate reclaims it. Time is injected via m.nowFn so the test
|
|
// needs no real sleep and no size threshold to trigger the behavior.
|
|
func TestMemoryStoreSweepTimeGated(t *testing.T) {
|
|
m := newMemoryStore()
|
|
fake := time.Date(2026, 8, 20, 0, 0, 0, 0, time.UTC)
|
|
m.nowFn = func() time.Time { return fake }
|
|
|
|
// lastSweep starts at the zero time, so this first incrBy always
|
|
// sweeps (a no-op on an empty store) and sets lastSweep to fake.
|
|
if _, err := m.incrBy("k1", 1, time.Nanosecond); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
// k1 is already expired (ttl was 1ns), but stay well inside the
|
|
// sweepEvery gate — the next incrBy must not scan, so k1 survives as
|
|
// a physical (if stale) map entry.
|
|
fake = fake.Add(time.Millisecond)
|
|
if _, err := m.incrBy("k2", 1, time.Minute); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !memoryStoreHasKey(m, "k1") {
|
|
t.Fatal("k1 must still be present: the sweep gate has not elapsed yet")
|
|
}
|
|
|
|
// Cross the sweepEvery boundary from the last sweep: this incrBy must
|
|
// scan and reclaim k1.
|
|
fake = fake.Add(sweepEvery)
|
|
if _, err := m.incrBy("k3", 1, time.Minute); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if memoryStoreHasKey(m, "k1") {
|
|
t.Error("k1 should have been swept once sweepEvery elapsed")
|
|
}
|
|
}
|
|
|
|
// memoryStoreHasKey reports whether key is physically present in m.data,
|
|
// bypassing get's expiry check — used to observe sweep timing directly.
|
|
func memoryStoreHasKey(m *memoryStore, key string) bool {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
_, ok := m.data[key]
|
|
return ok
|
|
}
|
|
|
|
// TestLimiterCostBudgetRoundsUSDConversion pins usdToMicros' rounding
|
|
// against a value (8.2) where float64 representation sits just under the
|
|
// true value: a truncating int64(8.2*1e6) would undershoot to 8_199_999,
|
|
// making a spend of exactly $8.199999 look already over an intended
|
|
// $8.20 budget. With rounding, 8_199_999 must stay under budget and only
|
|
// the exact $8.20 spend violates.
|
|
func TestLimiterCostBudgetRoundsUSDConversion(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerDayUSD: 8.2}}}
|
|
|
|
l.account(scopes, usage{}, 8_199_999)
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("spending $8.199999 against an $8.20 budget must not violate, got %+v", v)
|
|
}
|
|
|
|
l.account(scopes, usage{}, 1) // now exactly at $8.20
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("spending exactly $8.20 against an $8.20 budget must violate")
|
|
}
|
|
}
|
|
|
|
// TestAccountSkipsZeroMetrics asserts account writes no tokin/tokout
|
|
// counter when the matching usage direction is 0, and no cost counter
|
|
// when costMicros is 0 — each metric's absence is independently
|
|
// observable via get.
|
|
func TestAccountSkipsZeroMetrics(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "z", limits: nil}}
|
|
|
|
l.account(scopes, usage{}, 0)
|
|
|
|
tokIn, ok := l.getCounter("user", "z", metricTokIn, windowDay, now)
|
|
if !ok || tokIn != 0 {
|
|
t.Errorf("tokin counter = %d, ok=%v, want 0, true", tokIn, ok)
|
|
}
|
|
tokOut, ok := l.getCounter("user", "z", metricTokOut, windowDay, now)
|
|
if !ok || tokOut != 0 {
|
|
t.Errorf("tokout counter = %d, ok=%v, want 0, true", tokOut, ok)
|
|
}
|
|
cost, ok := l.getCounter("user", "z", metricCost, windowDay, now)
|
|
if !ok || cost != 0 {
|
|
t.Errorf("cost counter = %d, ok=%v, want 0, true", cost, ok)
|
|
}
|
|
}
|
|
|
|
// TestAccountSplitsTokensIndependently asserts account writes prompt
|
|
// tokens to metricTokIn and completion tokens to metricTokOut
|
|
// independently — including the asymmetric case (prompt only, no
|
|
// completion yet) the estimation fallback and a dropped stream both rely
|
|
// on: a nonzero prompt with a zero completion must still write tokin
|
|
// alone, never skip it because completion is 0.
|
|
func TestAccountSplitsTokensIndependently(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "split", limits: nil}}
|
|
|
|
l.account(scopes, usage{prompt: 25}, 0) // completion left at zero
|
|
|
|
tokIn, ok := l.getCounter("user", "split", metricTokIn, windowDay, now)
|
|
if !ok || tokIn != 25 {
|
|
t.Errorf("tokin counter = %d, ok=%v, want 25, true", tokIn, ok)
|
|
}
|
|
tokOut, ok := l.getCounter("user", "split", metricTokOut, windowDay, now)
|
|
if !ok || tokOut != 0 {
|
|
t.Errorf("tokout counter = %d, ok=%v, want 0, true (no completion tokens reported)", tokOut, ok)
|
|
}
|
|
|
|
l.account(scopes, usage{completion: 9}, 0) // prompt left at zero this time
|
|
tokIn, ok = l.getCounter("user", "split", metricTokIn, windowDay, now)
|
|
if !ok || tokIn != 25 {
|
|
t.Errorf("tokin counter = %d, ok=%v, want still 25 (this call reported no prompt tokens)", tokIn, ok)
|
|
}
|
|
tokOut, ok = l.getCounter("user", "split", metricTokOut, windowDay, now)
|
|
if !ok || tokOut != 9 {
|
|
t.Errorf("tokout counter = %d, ok=%v, want 9", tokOut, ok)
|
|
}
|
|
}
|
|
|
|
// TestTokenBudgetSumsInAndOut asserts a TokensPerDay/TokensPerMonth limit
|
|
// enforces a TOTAL budget across metricTokIn and metricTokOut combined —
|
|
// the operator directive's "limits stay total-token (limiter sums both at
|
|
// read)": neither direction alone crosses the budget, but their sum does.
|
|
func TestTokenBudgetSumsInAndOut(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerDay: 100}}}
|
|
|
|
l.account(scopes, usage{prompt: 60, completion: 45}, 0) // 105 > 100, split 60/45
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("60 tokin + 45 tokout = 105 > 100 must refuse, even though neither direction alone exceeds 100")
|
|
}
|
|
}
|
|
|
|
// TestLimiterRequestPerDayLimit is carried-item (b): a per-day request
|
|
// limit blocks once the day's count exceeds it, independent of the
|
|
// (unset, unlimited) per-minute limit.
|
|
func TestLimiterRequestPerDayLimit(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "a", limits: &LimitsConfig{RequestsPerDay: 2}}}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("1st should pass")
|
|
}
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatal("2nd should pass")
|
|
}
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("3rd should violate the day limit")
|
|
}
|
|
if v.storeDown {
|
|
t.Error("storeDown must be false for a real limit breach")
|
|
}
|
|
}
|
|
|
|
// TestLimiterTokensPerMonthBudget mirrors TestLimiterTokenBudget for the
|
|
// month window: accounting tokens past a month budget makes the next
|
|
// check refuse.
|
|
func TestLimiterTokensPerMonthBudget(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerMonth: 100}}}
|
|
l.account(scopes, usage{prompt: 60, completion: 50}, 0) // 110 > 100
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("over month token budget must refuse")
|
|
}
|
|
}
|
|
|
|
// TestLimiterCostPerMonthBudget mirrors TestLimiterCostBudget for the
|
|
// month window, and checks the violation carries the month window's
|
|
// Retry-After.
|
|
func TestLimiterCostPerMonthBudget(t *testing.T) {
|
|
now := time.Date(2026, 8, 20, 10, 0, 30, 0, time.UTC)
|
|
l := newLimiter(nil, true)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerMonthUSD: 1.00}}}
|
|
|
|
l.account(scopes, usage{}, 1_500_000) // $1.50 > $1.00 budget
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("over month cost budget must refuse")
|
|
}
|
|
wantRetry := int(time.Date(2026, 9, 1, 0, 0, 0, 0, time.UTC).Sub(now).Seconds())
|
|
if v.retryAfter != wantRetry {
|
|
t.Errorf("retryAfter = %d, want %d (seconds to next UTC month start)", v.retryAfter, wantRetry)
|
|
}
|
|
if v.storeDown {
|
|
t.Error("storeDown must be false for a real limit breach")
|
|
}
|
|
}
|
|
|
|
// erroringStore is a counterStore stub that always fails, used to test the
|
|
// limiter's fail-open/fail-closed handling on a store error without a
|
|
// real network dependency (see redis_store_test.go for the equivalent
|
|
// end-to-end test against an unreachable redisStore).
|
|
type erroringStore struct {
|
|
err error
|
|
}
|
|
|
|
func (s *erroringStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, s.err }
|
|
func (s *erroringStore) get(string) (int64, error) { return 0, s.err }
|
|
func (s *erroringStore) getMulti([]string) ([]int64, error) { return nil, s.err }
|
|
func (s *erroringStore) incrMulti([]counterIncr) ([]int64, error) { return nil, s.err }
|
|
|
|
// TestLimiter_FailOpen_StoreErrorUsesFallback is carried-item (c): a store
|
|
// error with failOpen=true must not block the request — the operation
|
|
// transparently falls back to the limiter's in-process memoryStore, and
|
|
// the fallback counter actually advances (proving the write really
|
|
// happened somewhere, not just that no violation was reported).
|
|
func TestLimiter_FailOpen_StoreErrorUsesFallback(t *testing.T) {
|
|
store := &erroringStore{err: errors.New("boom")}
|
|
l := newLimiter(store, true)
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("want no violation with failOpen=true on a store error, got %+v", v)
|
|
}
|
|
got, err := l.fallback.get(windowKey("user", "u", metricReq, windowDay, l.now()))
|
|
if err != nil {
|
|
t.Fatalf("fallback.get: %v", err)
|
|
}
|
|
if got != 1 {
|
|
t.Errorf("fallback req:day counter = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
// TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation is
|
|
// carried-item (c): failOpen=false refuses the request outright on a
|
|
// store error, rather than silently using the fallback.
|
|
func TestLimiter_FailClosed_StoreErrorReturnsStoreDownViolation(t *testing.T) {
|
|
store := &erroringStore{err: errors.New("boom")}
|
|
l := newLimiter(store, false)
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 100}}}
|
|
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("want a storeDown violation on a store error with failOpen=false")
|
|
}
|
|
if !v.storeDown {
|
|
t.Error("v.storeDown = false, want true")
|
|
}
|
|
if v.message != "limit store unavailable" {
|
|
t.Errorf("v.message = %q, want %q", v.message, "limit store unavailable")
|
|
}
|
|
}
|
|
|
|
// succeedIncrFailGetStore is a counterStore stub whose incrBy/incrMulti
|
|
// always succeed and whose get/getMulti always error. It isolates
|
|
// budgetViolation's own storeDown branch (limits.go) from checkAndCount's
|
|
// earlier incrMulti one: an erroringStore that fails every method makes
|
|
// checkAndCount's batched req:min/req:day/req:hour increment fail closed
|
|
// and return before evaluateScope — and therefore budgetViolation — ever
|
|
// runs, so a test built on it cannot actually prove budgetViolation's own
|
|
// branch works.
|
|
type succeedIncrFailGetStore struct {
|
|
getErr error
|
|
}
|
|
|
|
func (s *succeedIncrFailGetStore) incrBy(string, int64, time.Duration) (int64, error) {
|
|
return 1, nil
|
|
}
|
|
func (s *succeedIncrFailGetStore) get(string) (int64, error) { return 0, s.getErr }
|
|
func (s *succeedIncrFailGetStore) getMulti([]string) ([]int64, error) {
|
|
return nil, s.getErr
|
|
}
|
|
func (s *succeedIncrFailGetStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
|
out := make([]int64, len(entries))
|
|
for i := range out {
|
|
out[i] = 1
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation covers the
|
|
// budgetViolation fail-closed path specifically (as opposed to the
|
|
// request-counter incrMulti path TestLimiter_FailClosed_
|
|
// StoreErrorReturnsStoreDownViolation covers above): a store error on a
|
|
// token/cost budget read also refuses the request when failOpen is false.
|
|
// The store's incrMulti succeeds so checkAndCount's batched
|
|
// req:min/req:day/req:hour increment passes and evaluateScope actually
|
|
// reaches budgetViolation; a TokensPerDay limit makes evaluateScope call
|
|
// it.
|
|
func TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation(t *testing.T) {
|
|
store := &succeedIncrFailGetStore{getErr: errors.New("boom")}
|
|
l := newLimiter(store, false)
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{TokensPerDay: 100}}}
|
|
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("want a storeDown violation from the budget-read path with failOpen=false")
|
|
}
|
|
if !v.storeDown {
|
|
t.Error("v.storeDown = false, want true")
|
|
}
|
|
}
|
|
|
|
// TestLimiter_FailClosed_AccountDropsSampleOnStoreError asserts account
|
|
// drops a sample entirely — writes it nowhere — when failOpen=false and
|
|
// the store errors, rather than falling back to counting it locally (which
|
|
// would double-count once the store recovers). account has no error
|
|
// return to signal the drop, so this is observed via the fallback staying
|
|
// at 0.
|
|
func TestLimiter_FailClosed_AccountDropsSampleOnStoreError(t *testing.T) {
|
|
store := &erroringStore{err: errors.New("boom")}
|
|
l := newLimiter(store, false)
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: nil}}
|
|
|
|
l.account(scopes, usage{prompt: 10, completion: 10}, 500)
|
|
|
|
gotIn, err := l.fallback.get(windowKey("user", "u", metricTokIn, windowDay, l.now()))
|
|
if err != nil {
|
|
t.Fatalf("fallback.get: %v", err)
|
|
}
|
|
if gotIn != 0 {
|
|
t.Errorf("fallback tokin:day counter = %d, want 0 (sample must be dropped, not counted locally)", gotIn)
|
|
}
|
|
gotOut, err := l.fallback.get(windowKey("user", "u", metricTokOut, windowDay, l.now()))
|
|
if err != nil {
|
|
t.Fatalf("fallback.get: %v", err)
|
|
}
|
|
if gotOut != 0 {
|
|
t.Errorf("fallback tokout:day counter = %d, want 0 (sample must be dropped, not counted locally)", gotOut)
|
|
}
|
|
}
|
|
|
|
// TestLimiter_LogsStoreErrorOncePerRateLimit asserts a store error is
|
|
// logged via the limiter's injectable log func, rate-limited to once per
|
|
// storeErrorLogEvery — a store outage under load must not flood the log
|
|
// with one line per request, however many store operations fail within
|
|
// the window.
|
|
func TestLimiter_LogsStoreErrorOncePerRateLimit(t *testing.T) {
|
|
store := &erroringStore{err: errors.New("boom")}
|
|
l := newLimiter(store, true)
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
var logged []string
|
|
l.logf = func(format string, args ...any) { logged = append(logged, fmt.Sprintf(format, args...)) }
|
|
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: nil}}
|
|
|
|
l.checkAndCount(scopes) // two store ops (min, day) within one call
|
|
if len(logged) != 1 {
|
|
t.Fatalf("logged = %d lines after the first checkAndCount, want 1 (rate-limited)", len(logged))
|
|
}
|
|
|
|
l.checkAndCount(scopes) // still inside the rate-limit window
|
|
if len(logged) != 1 {
|
|
t.Fatalf("logged = %d lines after a second call inside the window, want 1", len(logged))
|
|
}
|
|
|
|
now = now.Add(storeErrorLogEvery)
|
|
l.checkAndCount(scopes)
|
|
if len(logged) != 2 {
|
|
t.Fatalf("logged = %d lines after the rate-limit window elapsed, want 2", len(logged))
|
|
}
|
|
}
|
|
|
|
// countingErrorStore is a counterStore stub that always errors and counts
|
|
// how many times each method is actually invoked. It proves the
|
|
// store-down latch (limits.go) skips the network call entirely during
|
|
// its window, rather than merely tolerating the error each time.
|
|
type countingErrorStore struct {
|
|
err error
|
|
incrCalls int
|
|
getCalls int
|
|
}
|
|
|
|
func (s *countingErrorStore) incrBy(string, int64, time.Duration) (int64, error) {
|
|
s.incrCalls++
|
|
return 0, s.err
|
|
}
|
|
|
|
func (s *countingErrorStore) get(string) (int64, error) {
|
|
s.getCalls++
|
|
return 0, s.err
|
|
}
|
|
|
|
func (s *countingErrorStore) getMulti([]string) ([]int64, error) {
|
|
s.getCalls++
|
|
return nil, s.err
|
|
}
|
|
|
|
func (s *countingErrorStore) incrMulti([]counterIncr) ([]int64, error) {
|
|
s.incrCalls++
|
|
return nil, s.err
|
|
}
|
|
|
|
// TestLimiter_StoreDownLatch_SkipsStoreCallsWithinWindow is review-round-3
|
|
// item 2: after a store operation fails, every operation for the next
|
|
// storeDownLatchFor must short-circuit straight to the fail-open policy
|
|
// without calling the store at all — a request that touches several
|
|
// counters (min/day/hour, tokens/cost, user/group) must not pay a fresh
|
|
// bounded wait against a store that just proved unreachable for each one.
|
|
// checkAndCount batches its whole req:min/req:day/req:hour increment into
|
|
// ONE incrMulti call (perf review, 2026-08-21), so "the store must not be
|
|
// called" now means that single call is skipped entirely during the latch
|
|
// window, not that some subset of several per-key calls is. failOpen=true
|
|
// is required so every checkAndCount call below returns no violation
|
|
// (ok=true via the fallback) instead of a storeDownViolation — the
|
|
// incrCalls assertions only make sense if checkAndCount actually
|
|
// completes normally each time.
|
|
func TestLimiter_StoreDownLatch_SkipsStoreCallsWithinWindow(t *testing.T) {
|
|
store := &countingErrorStore{err: errors.New("boom")}
|
|
l := newLimiter(store, true)
|
|
now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "user", id: "u", limits: nil}}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("want no violation with failOpen=true, got %+v", v)
|
|
}
|
|
firstCalls := store.incrCalls
|
|
if firstCalls == 0 {
|
|
t.Fatal("want the first call to actually reach the store at least once")
|
|
}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("want no violation with failOpen=true, got %+v", v)
|
|
}
|
|
if store.incrCalls != firstCalls {
|
|
t.Errorf("incrBy calls = %d after a second checkAndCount within the latch window, want unchanged at %d (the store must not be called)", store.incrCalls, firstCalls)
|
|
}
|
|
|
|
now = now.Add(storeDownLatchFor) // latch expires
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("want no violation with failOpen=true, got %+v", v)
|
|
}
|
|
if store.incrCalls <= firstCalls {
|
|
t.Errorf("incrBy calls = %d after the latch window elapsed, want more than %d (the store must be probed again)", store.incrCalls, firstCalls)
|
|
}
|
|
}
|
|
|
|
// TestLimitsConfigValidate is carried-item (a): validate rejects
|
|
// NaN/±Inf/negative cost limits and negative int64 count limits, accepts
|
|
// zero/positive values, and treats a nil receiver (no limits configured
|
|
// at all) as valid.
|
|
func TestLimitsConfigValidate(t *testing.T) {
|
|
cases := []struct {
|
|
lc *LimitsConfig
|
|
name string
|
|
wantErr bool
|
|
}{
|
|
{name: "nil is valid", lc: nil, wantErr: false},
|
|
{name: "zero value is valid", lc: &LimitsConfig{}, wantErr: false},
|
|
{name: "positive values are valid", lc: &LimitsConfig{
|
|
RequestsPerMinute: 1, RequestsPerDay: 1, TokensPerDay: 1, TokensPerMonth: 1,
|
|
CostPerDayUSD: 1.5, CostPerMonthUSD: 1.5,
|
|
}, wantErr: false},
|
|
{name: "negative RequestsPerMinute", lc: &LimitsConfig{RequestsPerMinute: -1}, wantErr: true},
|
|
{name: "negative RequestsPerDay", lc: &LimitsConfig{RequestsPerDay: -1}, wantErr: true},
|
|
{name: "negative TokensPerDay", lc: &LimitsConfig{TokensPerDay: -1}, wantErr: true},
|
|
{name: "negative TokensPerMonth", lc: &LimitsConfig{TokensPerMonth: -1}, wantErr: true},
|
|
{name: "negative CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: -0.01}, wantErr: true},
|
|
{name: "negative CostPerMonthUSD", lc: &LimitsConfig{CostPerMonthUSD: -0.01}, wantErr: true},
|
|
{name: "NaN CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: math.NaN()}, wantErr: true},
|
|
{name: "+Inf CostPerMonthUSD", lc: &LimitsConfig{CostPerMonthUSD: math.Inf(1)}, wantErr: true},
|
|
{name: "-Inf CostPerDayUSD", lc: &LimitsConfig{CostPerDayUSD: math.Inf(-1)}, wantErr: true},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
err := c.lc.validate()
|
|
if (err != nil) != c.wantErr {
|
|
t.Errorf("validate() error = %v, wantErr %v", err, c.wantErr)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestLimiter_FailPolicyGet covers failPolicyGet's own fail-open/fail-closed
|
|
// branching in isolation (storeGet's error/latched paths already drive it
|
|
// indirectly elsewhere in this file): failOpen=true reads through to the
|
|
// in-process fallback, failOpen=false refuses outright.
|
|
func TestLimiter_FailPolicyGet(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
failOpen bool
|
|
wantOK bool
|
|
}{
|
|
{name: "failOpen true reads the fallback", failOpen: true, wantOK: true},
|
|
{name: "failOpen false refuses", failOpen: false, wantOK: false},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
l := newLimiter(nil, c.failOpen)
|
|
v, ok := l.failPolicyGet("k")
|
|
require.Equal(t, c.wantOK, ok)
|
|
if c.wantOK {
|
|
assert.Equal(t, int64(0), v)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// --- perf review, 2026-08-21: single-batch counter writes ---
|
|
|
|
// countingIncrStore is a counterStore stub whose incrMulti/get/getMulti
|
|
// all succeed against a real backing map (so repeated calls actually
|
|
// accumulate, proving checkAndCount/account read back correct values
|
|
// after batching, not just that batching happened at all), counting how
|
|
// many times incrMulti and getMulti were each called. incrBy is never
|
|
// exercised by checkAndCount/account after the perf-review batching and
|
|
// just returns a zero value.
|
|
type countingIncrStore struct {
|
|
values map[string]int64
|
|
incrMultiCalls int
|
|
getMultiCalls int
|
|
}
|
|
|
|
func (s *countingIncrStore) incrBy(string, int64, time.Duration) (int64, error) { return 0, nil }
|
|
|
|
func (s *countingIncrStore) get(key string) (int64, error) {
|
|
if s.values == nil {
|
|
return 0, nil
|
|
}
|
|
return s.values[key], nil
|
|
}
|
|
|
|
func (s *countingIncrStore) getMulti(keys []string) ([]int64, error) {
|
|
s.getMultiCalls++
|
|
out := make([]int64, len(keys))
|
|
for i, k := range keys {
|
|
if s.values != nil {
|
|
out[i] = s.values[k]
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (s *countingIncrStore) incrMulti(entries []counterIncr) ([]int64, error) {
|
|
s.incrMultiCalls++
|
|
if s.values == nil {
|
|
s.values = make(map[string]int64)
|
|
}
|
|
out := make([]int64, len(entries))
|
|
for i, e := range entries {
|
|
s.values[e.key] += e.delta
|
|
out[i] = s.values[e.key]
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount is the
|
|
// perf-review (2026-08-21) case for checkAndCount: a 3-scope request
|
|
// (user, group, total — 9 counters, min/day/hour x 3 scopes) must cost
|
|
// exactly ONE incrMulti round trip, not one per counter, and the batched
|
|
// reply must still drive correct evaluation — a real requestsPerMinute
|
|
// violation on the 3rd call, with the exact right counter value.
|
|
func TestCheckAndCount_OneIncrMultiCallRegardlessOfScopeCount(t *testing.T) {
|
|
store := &countingIncrStore{}
|
|
l := newLimiter(store, true)
|
|
now := time.Date(2026, 8, 21, 9, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
scopes := []limitScope{
|
|
{kind: "user", id: "u", limits: &LimitsConfig{RequestsPerMinute: 2}},
|
|
{kind: "group", id: "g", limits: nil},
|
|
{kind: totalScopeKind, id: totalScopeID, limits: nil},
|
|
}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("1st call should pass, got %+v", v)
|
|
}
|
|
if store.incrMultiCalls != 1 {
|
|
t.Errorf("incrMultiCalls after 1st checkAndCount = %d, want 1", store.incrMultiCalls)
|
|
}
|
|
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("2nd call should pass, got %+v", v)
|
|
}
|
|
if store.incrMultiCalls != 2 {
|
|
t.Errorf("incrMultiCalls after 2nd checkAndCount = %d, want 2 (one call per checkAndCount, not per scope or per counter)", store.incrMultiCalls)
|
|
}
|
|
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("3rd call should violate u's requestsPerMinute:2")
|
|
}
|
|
if v.storeDown {
|
|
t.Error("v.storeDown must be false — this is a real limit breach, not a store failure")
|
|
}
|
|
if store.incrMultiCalls != 3 {
|
|
t.Errorf("incrMultiCalls after 3rd checkAndCount = %d, want 3", store.incrMultiCalls)
|
|
}
|
|
|
|
// Values correct: batching three scopes into one call must not
|
|
// corrupt any one scope's own counter.
|
|
uMin, ok := l.getCounter("user", "u", metricReq, windowMin, now)
|
|
if !ok || uMin != 3 {
|
|
t.Errorf("user req:min = %d, ok=%v, want 3", uMin, ok)
|
|
}
|
|
totalMin, ok := l.getCounter(totalScopeKind, totalScopeID, metricReq, windowMin, now)
|
|
if !ok || totalMin != 3 {
|
|
t.Errorf("total req:min = %d, ok=%v, want 3", totalMin, ok)
|
|
}
|
|
gDay, ok := l.getCounter("group", "g", metricReq, windowDay, now)
|
|
if !ok || gDay != 3 {
|
|
t.Errorf("group req:day = %d, ok=%v, want 3", gDay, ok)
|
|
}
|
|
}
|
|
|
|
// TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount is the
|
|
// perf-review (2026-08-21) case for account: a 3-scope request with all
|
|
// three metrics nonzero (27 counters: 3 scopes x 3 metrics x 3 windows)
|
|
// must cost exactly ONE incrMulti round trip, and the batched reply must
|
|
// land every value in the right scope/metric/window slot.
|
|
func TestAccount_OneIncrMultiCallRegardlessOfScopeOrMetricCount(t *testing.T) {
|
|
store := &countingIncrStore{}
|
|
l := newLimiter(store, true)
|
|
now := time.Date(2026, 8, 21, 9, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
scopes := []limitScope{
|
|
{kind: "user", id: "u"},
|
|
{kind: "group", id: "g"},
|
|
{kind: totalScopeKind, id: totalScopeID},
|
|
}
|
|
l.account(scopes, usage{prompt: 10, completion: 5}, 700)
|
|
|
|
if store.incrMultiCalls != 1 {
|
|
t.Errorf("incrMultiCalls = %d, want 1 (one batch for 27 counters)", store.incrMultiCalls)
|
|
}
|
|
|
|
uTokInDay, ok := l.getCounter("user", "u", metricTokIn, windowDay, now)
|
|
if !ok || uTokInDay != 10 {
|
|
t.Errorf("user tokin:day = %d, ok=%v, want 10", uTokInDay, ok)
|
|
}
|
|
totalTokOutHour, ok := l.getCounter(totalScopeKind, totalScopeID, metricTokOut, windowHour, now)
|
|
if !ok || totalTokOutHour != 5 {
|
|
t.Errorf("total tokout:hour = %d, ok=%v, want 5", totalTokOutHour, ok)
|
|
}
|
|
gCostMonth, ok := l.getCounter("group", "g", metricCost, windowMonth, now)
|
|
if !ok || gCostMonth != 700 {
|
|
t.Errorf("group cost:month = %d, ok=%v, want 700", gCostMonth, ok)
|
|
}
|
|
|
|
// A second call accumulates rather than overwrites, and still costs
|
|
// exactly one more round trip.
|
|
l.account(scopes, usage{prompt: 2}, 0)
|
|
if store.incrMultiCalls != 2 {
|
|
t.Errorf("incrMultiCalls after 2nd account = %d, want 2", store.incrMultiCalls)
|
|
}
|
|
uTokInDay, ok = l.getCounter("user", "u", metricTokIn, windowDay, now)
|
|
if !ok || uTokInDay != 12 {
|
|
t.Errorf("user tokin:day after 2nd account = %d, ok=%v, want 12", uTokInDay, ok)
|
|
}
|
|
}
|
|
|
|
// TestTokenBudget_ExactBoundaryViolates asserts budgetViolation/
|
|
// tokenBudgetViolation's ">= limit" semantics at the exact boundary: a
|
|
// scope whose accumulated tokin+tokout lands EXACTLY on the limit (not
|
|
// merely over it) must still violate. budgetViolation's own "used <
|
|
// limit passes" check makes "used == limit" the smallest value that
|
|
// violates; this pins that boundary directly rather than only exercising
|
|
// values already over it (TestLimiterTokenBudget, TestTokenBudgetSumsInAndOut).
|
|
func TestTokenBudget_ExactBoundaryViolates(t *testing.T) {
|
|
l := newLimiter(nil, true)
|
|
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerDay: 100}}}
|
|
|
|
l.account(scopes, usage{prompt: 63, completion: 36}, 0) // 99 < 100: must not violate yet
|
|
if v := l.checkAndCount(scopes); v != nil {
|
|
t.Fatalf("99 < 100 must not violate, got %+v", v)
|
|
}
|
|
|
|
l.account(scopes, usage{completion: 1}, 0) // pushes tokout to 37: total now exactly 100
|
|
v := l.checkAndCount(scopes)
|
|
if v == nil {
|
|
t.Fatal("tokin+tokout landing exactly on the 100 budget must violate (>= limit, not only > limit)")
|
|
}
|
|
if v.storeDown {
|
|
t.Error("storeDown must be false for a real limit breach")
|
|
}
|
|
}
|
|
|
|
// --- review round 2, 2026-08-21: enforcement-aware fallback TTL clamp ---
|
|
|
|
// TestClampTTL_ExactMathTable pins clampTTL's ceiling-vs-floor precedence
|
|
// for every window's real (ttl, enforceTTL) pair. Round 1's clamp applied
|
|
// memoryStoreMaxTTL unconditionally; round 2 corrects it so a window's
|
|
// own enforceTTL floor always wins a conflict with that ceiling — this is
|
|
// the exact math that keeps min/hour/day capped at 48h while month lands
|
|
// back at its own ~32-day floor.
|
|
func TestClampTTL_ExactMathTable(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
ttl time.Duration
|
|
enforceTTL time.Duration
|
|
want time.Duration
|
|
}{
|
|
{"min: 2m sits under both ceiling and floor, unchanged", minWindowTTL, enforceTTLFor(windowMin), minWindowTTL},
|
|
{"hour: 48h ttl equals the 48h ceiling; its own 2h floor never applies", hourWindowTTL, enforceTTLFor(windowHour), memoryStoreMaxTTL},
|
|
{"day: 35d ttl clamped down to the 48h ceiling, still above the 25h floor", dayWindowTTL, enforceTTLFor(windowDay), memoryStoreMaxTTL},
|
|
{"month: 400d ttl would clamp to 48h, but the 32d floor wins instead", monthWindowTTL, enforceTTLFor(windowMonth), enforceTTLFor(windowMonth)},
|
|
{"enforceTTL=0 (incrBy's own single-key path): pure ceiling, no floor", 400 * 24 * time.Hour, 0, memoryStoreMaxTTL},
|
|
{"ttl already under both ceiling and floor passes through unchanged", 90 * time.Minute, time.Hour, 90 * time.Minute},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
if got := clampTTL(c.ttl, c.enforceTTL); got != c.want {
|
|
t.Errorf("clampTTL(%v, %v) = %v, want %v", c.ttl, c.enforceTTL, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestEnforceTTLFor_MatchesTheOriginalPreHistoryBumpValues pins
|
|
// enforceTTLFor's four values directly — they are not arbitrary, they
|
|
// match dayWindowTTL/monthWindowTTL's values from before the v0.2
|
|
// data-layer task's history-retention bump (25h/32d), which were
|
|
// themselves already windowLength+margin for correct single-window
|
|
// enforcement.
|
|
func TestEnforceTTLFor_MatchesTheOriginalPreHistoryBumpValues(t *testing.T) {
|
|
cases := []struct {
|
|
window string
|
|
want time.Duration
|
|
}{
|
|
{windowMin, 2 * time.Minute},
|
|
{windowHour, 2 * time.Hour},
|
|
{windowDay, 25 * time.Hour},
|
|
{windowMonth, 32 * 24 * time.Hour},
|
|
}
|
|
for _, c := range cases {
|
|
if got := enforceTTLFor(c.window); got != c.want {
|
|
t.Errorf("enforceTTLFor(%q) = %v, want %v", c.window, got, c.want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestMemoryStore_IncrMulti_MonthCounterSurvivesPast48hAndStillEnforces
|
|
// is the review-round-2 regression case: a month counter written via
|
|
// incrMulti — the real production path, checkAndCount/account — must
|
|
// stay alive, not silently reset to 0, well past 48h, and a
|
|
// TokensPerMonth limit must still enforce correctly at that point. This
|
|
// is exactly what round 1's unconditional 48h ceiling broke: a month
|
|
// counter's key would have expired and reset roughly every two days,
|
|
// turning the budget into a rolling ~48h one instead of a true
|
|
// calendar-month one whenever Redis is absent or down.
|
|
func TestMemoryStore_IncrMulti_MonthCounterSurvivesPast48hAndStillEnforces(t *testing.T) {
|
|
l := newLimiter(nil, true) // nil store -> every operation uses the in-process fallback
|
|
now := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
|
l.nowFn = func() time.Time { return now }
|
|
scopes := []limitScope{{kind: "group", id: "g", limits: &LimitsConfig{TokensPerMonth: 100}}}
|
|
|
|
l.account(scopes, usage{prompt: 60, completion: 20}, 0) // 80 tokens on day 1
|
|
|
|
// Advance 5 days: well past memoryStoreMaxTTL's 48h ceiling, far short
|
|
// of monthWindowTTL's own natural rollover.
|
|
now = now.Add(5 * 24 * time.Hour)
|
|
l.nowFn = func() time.Time { return now }
|
|
|
|
tokIn, ok := l.getCounter("group", "g", metricTokIn, windowMonth, now)
|
|
if !ok || tokIn != 60 {
|
|
t.Fatalf("tokin:month after 5 days = %d, ok=%v, want 60 (the key must not have expired at 48h)", tokIn, ok)
|
|
}
|
|
|
|
// Push the total over budget and confirm enforcement still fires —
|
|
// proves the surviving counter is actually read by
|
|
// tokenBudgetViolation, not merely present.
|
|
l.account(scopes, usage{completion: 21}, 0) // 60 + 41 = 101 > 100
|
|
if v := l.checkAndCount(scopes); v == nil {
|
|
t.Fatal("want a TokensPerMonth violation 5 days in; the round-1 clamp bug would have reset the counter well before this")
|
|
}
|
|
}
|
|
|
|
// TestMemoryStore_IncrMulti_DayCounterStillClampsAt48h asserts the
|
|
// round-1 clamp's original goal — bounding the fallback's own live-key
|
|
// growth — still holds for day, whose own enforcement floor (25h) sits
|
|
// below the 48h ceiling: a day counter written via incrMulti expires and
|
|
// resets once its key crosses 48h, well before dayWindowTTL's own 35-day
|
|
// retention would otherwise have kept it alive.
|
|
func TestMemoryStore_IncrMulti_DayCounterStillClampsAt48h(t *testing.T) {
|
|
m := newMemoryStore()
|
|
now := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
|
|
m.nowFn = func() time.Time { return now }
|
|
|
|
entry := counterIncr{key: "k", delta: 10, ttl: dayWindowTTL, enforceTTL: enforceTTLFor(windowDay)}
|
|
if _, err := m.incrMulti([]counterIncr{entry}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
now = now.Add(49 * time.Hour) // just past the 48h ceiling
|
|
m.nowFn = func() time.Time { return now }
|
|
|
|
got, err := m.get("k")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got != 0 {
|
|
t.Errorf("day counter at 49h past write = %d, want 0 (expired at the 48h ceiling; day's own 25h floor never raises it)", got)
|
|
}
|
|
}
|