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) } }