Files
traefik-llmgateway/limits_test.go
T
lukaszraczylo 5cc681496c perf(limits): single-batch counter writes, bounded fallback retention
Review fix round (2 Important + 5 folded) on 08cfdb4/b3fb705.

1. counterStore gains incrMulti(entries []counterIncr) ([]int64, error).
   redisStore.incrMulti sends every entry's INCRBY+EXPIRE pair as ONE
   pipeline (ttlToSeconds extracted from incrBy, now shared);
   memoryStore.incrMulti loops incrBy for interface conformance.
   checkAndCount batches every scope's req:min/req:day/req:hour into one
   storeIncrMulti call (down from up to 9 round trips for a 3-scope
   request) and reads req:min/req:day back from the reply, no extra
   GETs. account batches every scope's nonzero tok-in/tok-out/cost
   entries into one call (down from up to 27). Violation math and
   fail-open/fail-closed/latched semantics are unchanged — only the
   round-trip count moves.

2. memoryStoreMaxTTL (48h) caps every TTL memoryStore actually applies,
   clamped inside incrBy so incrMulti's loop inherits it for free. The
   in-process fallback is continuity, not history: applying the real
   hour/day/month TTLs there was growing live-key count ~30x, multiplying
   sweepEvery's own documented "167x incrBy cliff" by the same factor.
   Sweep cadence is unchanged (still 30s-gated); the clamp shrinks what
   there is to sweep instead.

Folded: historyMaxSpan's margin comment corrected (~4 days, not a
month); checkAndCount's godoc names req:hour; the integration
current-hour assertion sums every point instead of reading only the
last one (an hour-rollover race the suite cannot control); a new
exact-boundary test proves tokin+tokout landing exactly on a
TokensPerDay limit still violates; README documents the history route
as unthrottled (admin-gate is its only trust boundary) and the
fallback's 48h retention ceiling.

Tests: countingIncrStore proves exactly one incrMulti call per
checkAndCount and per account, with correct per-scope/metric/window
values and unchanged violation math; a fake-RESP-server test pins
incrMulti's wire format (one pipeline, entry-ordered replies); TTL
clamp tests cover both the ceiling and the pass-through case at the
boundary.

Gates: go test ./... -count=1, -race, gofmt, vet, golangci-lint,
make yaegi-check all green. make integration: every test passes
except the pre-existing TestUsersFileHotReload flake, which is NOT
caused by this change — verified via an isolated worktree at the
prior commit (b3fb705) run under identical conditions (full suite,
same host), which fails identically. This host's Docker Desktop is
currently slow to propagate the bind-mounted users.json change into
the containers; auth.go/users_file.go are untouched by this diff.
2026-08-21 01:02:28 +01:00

1259 lines
48 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")
}
}