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 } // 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) } } // 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 } // 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 always // succeeds and whose get always errors. It isolates budgetViolation's own // storeDown branch (limits.go) from checkAndCount's earlier incrCounter // one: an erroringStore that fails both methods makes checkAndCount's // initial req:min/req:day increments 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 } // TestLimiter_FailClosed_BudgetReadReturnsStoreDownViolation covers the // budgetViolation fail-closed path specifically (as opposed to the // request-counter incrCounter 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 incrBy succeeds so checkAndCount's req:min/req:day // increments pass 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 } // 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/month, tokens/cost, user/group) must not pay a fresh // bounded wait against a store that just proved unreachable for each one. // failOpen=true is used so checkAndCount completes both its req:min and // req:day increments per call instead of short-circuiting after the // first — this is what makes the "zero additional store calls" assertion // meaningful, rather than trivially true because checkAndCount stops // after one op regardless of the latch. 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) } }) } }