mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-14 05:18:54 +00:00
feat: limit engine with fixed windows, token and cost budgets
This commit is contained in:
@@ -0,0 +1,374 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Window identifiers used throughout windowKey, counter metrics, and
|
||||
// retry-after calculation.
|
||||
const (
|
||||
windowMin = "min"
|
||||
windowDay = "day"
|
||||
windowMonth = "month"
|
||||
)
|
||||
|
||||
// Counter metric names embedded in windowKey.
|
||||
const (
|
||||
metricReq = "req"
|
||||
metricTok = "tok"
|
||||
metricCost = "cost"
|
||||
)
|
||||
|
||||
// TTLs applied to counter keys. They exceed their window's natural length
|
||||
// so a key stays readable for the whole window it belongs to; the bucket
|
||||
// embedded in the key (see windowKey) already makes a rolled-over window
|
||||
// use a different key, so these TTLs only bound how long a stale key
|
||||
// lingers in memoryStore before an opportunistic sweep reclaims it.
|
||||
const (
|
||||
minWindowTTL = 2 * time.Minute
|
||||
dayWindowTTL = 25 * time.Hour
|
||||
monthWindowTTL = 32 * 24 * time.Hour
|
||||
)
|
||||
|
||||
// usdToMicroFactor scales a USD amount to micro-USD (1e6 micro-USD per
|
||||
// USD), matching the integer accounting unit used by account and the
|
||||
// counterStore.
|
||||
const usdToMicroFactor = 1_000_000
|
||||
|
||||
// counterStore is the storage backend the limiter uses for atomic windowed
|
||||
// counters. memoryStore (below) is the in-process fallback; a distributed
|
||||
// (e.g. Redis-backed) implementation is wired in a later task.
|
||||
type counterStore interface {
|
||||
// incrBy adds n to key's counter, creating it with an expiry of ttl
|
||||
// from now if it does not exist or has expired, and returns the
|
||||
// counter's new value.
|
||||
incrBy(key string, n int64, ttl time.Duration) (int64, error)
|
||||
// get returns key's current counter value, or 0 if it does not exist
|
||||
// or has expired.
|
||||
get(key string) (int64, error)
|
||||
}
|
||||
|
||||
// memoryEntry is one counter's value and expiry in memoryStore.
|
||||
type memoryEntry struct {
|
||||
expiry time.Time
|
||||
value int64
|
||||
}
|
||||
|
||||
// memoryStoreSweepThreshold is the key count above which incrBy
|
||||
// opportunistically deletes expired entries before writing. Kept well
|
||||
// above typical live-scope counts so the sweep is rare on a healthy
|
||||
// gateway and only fires when stale keys are actually accumulating.
|
||||
const memoryStoreSweepThreshold = 1024
|
||||
|
||||
// memoryStore is an in-process counterStore: a mutex-guarded map with
|
||||
// per-key expiry. It has no background goroutine — a Yaegi middleware
|
||||
// instance is rebuilt on every config reload, and any ticker or goroutine
|
||||
// started here would leak on rebuild instead of being collected with the
|
||||
// rest of the old instance. Expired entries are instead swept
|
||||
// opportunistically from incrBy once the map grows past
|
||||
// memoryStoreSweepThreshold keys.
|
||||
type memoryStore struct {
|
||||
data map[string]*memoryEntry
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// newMemoryStore returns an empty memoryStore.
|
||||
func newMemoryStore() *memoryStore {
|
||||
return &memoryStore{data: make(map[string]*memoryEntry)}
|
||||
}
|
||||
|
||||
// incrBy implements counterStore.
|
||||
func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
|
||||
now := time.Now()
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
if len(m.data) > memoryStoreSweepThreshold {
|
||||
m.sweepLocked(now)
|
||||
}
|
||||
|
||||
e, ok := m.data[key]
|
||||
if !ok || !now.Before(e.expiry) {
|
||||
e = &memoryEntry{expiry: now.Add(ttl)}
|
||||
m.data[key] = e
|
||||
}
|
||||
e.value += n
|
||||
return e.value, nil
|
||||
}
|
||||
|
||||
// get implements counterStore.
|
||||
func (m *memoryStore) get(key string) (int64, error) {
|
||||
now := time.Now()
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
e, ok := m.data[key]
|
||||
if !ok || !now.Before(e.expiry) {
|
||||
return 0, nil
|
||||
}
|
||||
return e.value, nil
|
||||
}
|
||||
|
||||
// sweepLocked deletes every expired entry. Callers must hold m.mu.
|
||||
func (m *memoryStore) sweepLocked(now time.Time) {
|
||||
for k, e := range m.data {
|
||||
if !now.Before(e.expiry) {
|
||||
delete(m.data, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// usage holds one request's token counts. estimated marks counts derived
|
||||
// from a tokenizer estimate rather than the provider's reported usage
|
||||
// (e.g. for a streaming response cut off before an upstream usage frame
|
||||
// arrived).
|
||||
type usage struct {
|
||||
prompt, completion int64
|
||||
// estimated is set by a later task's streaming/estimation path; this
|
||||
// task only produces the field for that interface, never reads it.
|
||||
estimated bool //nolint:unused // interface field for a later task, see comment above
|
||||
}
|
||||
|
||||
// total returns the request's combined prompt and completion tokens.
|
||||
func (u usage) total() int64 {
|
||||
return u.prompt + u.completion
|
||||
}
|
||||
|
||||
// limitScope is one entity (a user or their group) whose limits apply to a
|
||||
// request. checkAndCount and account evaluate every scope in the slice
|
||||
// they are given, so a request is counted and checked against both its
|
||||
// user's and its group's limits in one call.
|
||||
type limitScope struct {
|
||||
limits *LimitsConfig
|
||||
kind string // "user" or "group"
|
||||
id string
|
||||
}
|
||||
|
||||
// limitViolation describes the first limit a request breached.
|
||||
type limitViolation struct {
|
||||
message string
|
||||
// retryAfter is the number of seconds until the violated window ends
|
||||
// — until the client can plausibly succeed again.
|
||||
retryAfter int
|
||||
// storeDown reports whether the violation was manufactured because the
|
||||
// configured counterStore was unreachable, rather than an actual limit
|
||||
// breach. Always false in this task; a later task wires failOpen /
|
||||
// fail-closed handling for a real remote store.
|
||||
storeDown bool
|
||||
}
|
||||
|
||||
// limiter enforces per-minute/day/month request, token, and cost limits
|
||||
// using fixed windows keyed by windowKey.
|
||||
type limiter struct {
|
||||
store counterStore // nil means always use fallback (see newLimiter)
|
||||
fallback *memoryStore // in-process counter store, always available
|
||||
nowFn func() time.Time // injected for tests; defaults to time.Now
|
||||
failOpen bool // wired for a later task's store-down handling
|
||||
}
|
||||
|
||||
// newLimiter returns a limiter. A nil store means every operation uses the
|
||||
// limiter's own in-process fallback memoryStore — the plugin still
|
||||
// enforces limits with no distributed backend configured, just without
|
||||
// sharing counters across Traefik instances.
|
||||
func newLimiter(store counterStore, failOpen bool) *limiter {
|
||||
return &limiter{
|
||||
store: store,
|
||||
fallback: newMemoryStore(),
|
||||
nowFn: time.Now,
|
||||
failOpen: failOpen,
|
||||
}
|
||||
}
|
||||
|
||||
// now returns the limiter's current time, via nowFn.
|
||||
func (l *limiter) now() time.Time {
|
||||
return l.nowFn()
|
||||
}
|
||||
|
||||
// activeStore returns the store operations should use: the configured
|
||||
// store when set, otherwise the in-process fallback.
|
||||
func (l *limiter) activeStore() counterStore {
|
||||
if l.store != nil {
|
||||
return l.store
|
||||
}
|
||||
return l.fallback
|
||||
}
|
||||
|
||||
// windowKey builds the counterStore key for one (kind, id, metric, window)
|
||||
// counter at time t: llmgw:{kind}:{id}:{metric}:{window}:{bucket}. bucket
|
||||
// is t.UTC() formatted to the window's granularity, so a counter's key
|
||||
// changes automatically when its window rolls over.
|
||||
func windowKey(kind, id, metric, window string, t time.Time) string {
|
||||
u := t.UTC()
|
||||
var bucket string
|
||||
switch window {
|
||||
case windowMin:
|
||||
bucket = u.Format("200601021504")
|
||||
case windowDay:
|
||||
bucket = u.Format("20060102")
|
||||
case windowMonth:
|
||||
bucket = u.Format("200601")
|
||||
default:
|
||||
// window is always one of the three constants above, set by this
|
||||
// file's own callers, never by request input — an unknown value
|
||||
// here is a programming error. ServeHTTP's recoverPanic turns
|
||||
// this into a logged 500 instead of crashing the process.
|
||||
panic(fmt.Sprintf("llmgateway: windowKey: unknown window %q", window))
|
||||
}
|
||||
return fmt.Sprintf("llmgw:%s:%s:%s:%s:%s", kind, id, metric, window, bucket)
|
||||
}
|
||||
|
||||
// windowEnd returns the UTC instant at which window's bucket containing t
|
||||
// ends: the next minute boundary, the next UTC midnight, or the first of
|
||||
// next month UTC.
|
||||
func windowEnd(t time.Time, window string) time.Time {
|
||||
u := t.UTC()
|
||||
switch window {
|
||||
case windowMin:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), u.Hour(), u.Minute(), 0, 0, time.UTC).Add(time.Minute)
|
||||
case windowDay:
|
||||
return time.Date(u.Year(), u.Month(), u.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, 1)
|
||||
case windowMonth:
|
||||
return time.Date(u.Year(), u.Month(), 1, 0, 0, 0, 0, time.UTC).AddDate(0, 1, 0)
|
||||
default:
|
||||
panic(fmt.Sprintf("llmgateway: windowEnd: unknown window %q", window))
|
||||
}
|
||||
}
|
||||
|
||||
// retryAfterSeconds returns the whole number of seconds from t until
|
||||
// window's bucket ends, rounded up so a client never retries a moment too
|
||||
// early.
|
||||
func retryAfterSeconds(t time.Time, window string) int {
|
||||
return int(math.Ceil(windowEnd(t, window).Sub(t.UTC()).Seconds()))
|
||||
}
|
||||
|
||||
// incrCounter increments the (kind, id, metric, window) counter at t by n
|
||||
// and returns its new value. A store error is treated as "no increment
|
||||
// observed" — the fallback memoryStore never errors, and a real store's
|
||||
// fail-open/fail-closed policy is wired in a later task.
|
||||
func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) int64 {
|
||||
v, err := l.activeStore().incrBy(windowKey(kind, id, metric, window, t), n, ttl)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// getCounter returns the (kind, id, metric, window) counter's value at t.
|
||||
func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, error) {
|
||||
return l.activeStore().get(windowKey(kind, id, metric, window, t))
|
||||
}
|
||||
|
||||
// checkAndCount increments every scope's req:min and req:day counters —
|
||||
// unconditionally, before any evaluation, so a request that ultimately
|
||||
// gets refused by one scope's limit still counts toward every other
|
||||
// scope's rate tracking — then evaluates each scope's set limits in order
|
||||
// and returns the first violation found, or nil if the request may
|
||||
// proceed.
|
||||
//
|
||||
// Requests limits compare against the value just incremented in this call.
|
||||
// Token and cost limits compare against the value already accumulated by
|
||||
// account (via get): a request that would start over budget is refused,
|
||||
// but because usage is only known after the upstream call completes, one
|
||||
// request may still push a scope over its budget — an accepted trade-off.
|
||||
func (l *limiter) checkAndCount(scopes []limitScope) *limitViolation {
|
||||
now := l.now()
|
||||
|
||||
minCounts := make([]int64, len(scopes))
|
||||
dayCounts := make([]int64, len(scopes))
|
||||
for i, sc := range scopes {
|
||||
minCounts[i] = l.incrCounter(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL)
|
||||
dayCounts[i] = l.incrCounter(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL)
|
||||
}
|
||||
|
||||
for i, sc := range scopes {
|
||||
if v := l.evaluateScope(sc, minCounts[i], dayCounts[i], now); v != nil {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// evaluateScope checks one scope's already-set limits against its
|
||||
// just-incremented request counts and its accumulated token/cost counts.
|
||||
// A nil limits field means nothing to evaluate — the scope's req counters
|
||||
// still incremented above, feeding other scopes' aggregates and future
|
||||
// observability, but this scope itself never blocks a request.
|
||||
func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now time.Time) *limitViolation {
|
||||
if sc.limits == nil {
|
||||
return nil
|
||||
}
|
||||
lim := sc.limits
|
||||
|
||||
if v := requestLimitViolation(sc, "requests-per-minute", lim.RequestsPerMinute, minCount, windowMin, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := requestLimitViolation(sc, "requests-per-day", lim.RequestsPerDay, dayCount, windowDay, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricTok, "tokens-per-day", lim.TokensPerDay, windowDay, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricTok, "tokens-per-month", lim.TokensPerMonth, windowMonth, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricCost, "cost-per-day", int64(lim.CostPerDayUSD*usdToMicroFactor), windowDay, now); v != nil {
|
||||
return v
|
||||
}
|
||||
if v := l.budgetViolation(sc, metricCost, "cost-per-month", int64(lim.CostPerMonthUSD*usdToMicroFactor), windowMonth, now); v != nil {
|
||||
return v
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// requestLimitViolation reports a violation when count (already
|
||||
// incremented for this request) exceeds limit. limit<=0 means unlimited.
|
||||
func requestLimitViolation(sc limitScope, name string, limit, count int64, window string, now time.Time) *limitViolation {
|
||||
if limit <= 0 || count <= limit {
|
||||
return nil
|
||||
}
|
||||
return &limitViolation{
|
||||
message: fmt.Sprintf("%s %q exceeded %s limit (%d)", sc.kind, sc.id, name, limit),
|
||||
retryAfter: retryAfterSeconds(now, window),
|
||||
}
|
||||
}
|
||||
|
||||
// budgetViolation reports a violation when the scope's accumulated
|
||||
// counter for metric/window is already at or above limit. limit<=0 means
|
||||
// unlimited. A store read error fails open — no violation reported —
|
||||
// consistent with this task's storeDown always being false; a later task
|
||||
// wires strict fail-closed handling for a real remote store.
|
||||
func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int64, window string, now time.Time) *limitViolation {
|
||||
if limit <= 0 {
|
||||
return nil
|
||||
}
|
||||
used, err := l.getCounter(sc.kind, sc.id, metric, window, now)
|
||||
if err != nil || used < limit {
|
||||
return nil
|
||||
}
|
||||
return &limitViolation{
|
||||
message: fmt.Sprintf("%s %q exceeded %s budget", sc.kind, sc.id, name),
|
||||
retryAfter: retryAfterSeconds(now, window),
|
||||
}
|
||||
}
|
||||
|
||||
// account records u's total tokens and costMicros against every scope's
|
||||
// day and month counters. A zero-valued metric is skipped entirely — no
|
||||
// store write for a metric this call has nothing to report.
|
||||
func (l *limiter) account(scopes []limitScope, u usage, costMicros int64) {
|
||||
now := l.now()
|
||||
total := u.total()
|
||||
|
||||
for _, sc := range scopes {
|
||||
if total != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricTok, windowDay, now, total, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricTok, windowMonth, now, total, monthWindowTTL)
|
||||
}
|
||||
if costMicros != 0 {
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowDay, now, costMicros, dayWindowTTL)
|
||||
l.incrCounter(sc.kind, sc.id, metricCost, windowMonth, now, costMicros, monthWindowTTL)
|
||||
}
|
||||
}
|
||||
}
|
||||
+286
@@ -0,0 +1,286 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 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.
|
||||
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
|
||||
}{
|
||||
{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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMemoryStoreSweepOnAccess grows the store past
|
||||
// memoryStoreSweepThreshold, backdates every entry's expiry directly
|
||||
// (rather than racing a short TTL against the insertion loop's own real
|
||||
// wall-clock time, which is inherently flaky), and asserts the next
|
||||
// incrBy's opportunistic sweep reclaims every expired entry.
|
||||
func TestMemoryStoreSweepOnAccess(t *testing.T) {
|
||||
m := newMemoryStore()
|
||||
for i := 0; i < memoryStoreSweepThreshold+10; i++ {
|
||||
key := fmt.Sprintf("k%d", i)
|
||||
if _, err := m.incrBy(key, 1, time.Hour); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
past := time.Now().Add(-time.Hour)
|
||||
for _, e := range m.data {
|
||||
e.expiry = past
|
||||
}
|
||||
setupCount := len(m.data)
|
||||
m.mu.Unlock()
|
||||
if setupCount != memoryStoreSweepThreshold+10 {
|
||||
t.Fatalf("setup: got %d keys, want %d", setupCount, memoryStoreSweepThreshold+10)
|
||||
}
|
||||
|
||||
if _, err := m.incrBy("trigger-sweep", 1, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
m.mu.Lock()
|
||||
n := len(m.data)
|
||||
m.mu.Unlock()
|
||||
if n != 1 {
|
||||
t.Errorf("expected sweep to reclaim every expired entry, %d keys remain (want 1: trigger-sweep)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccountSkipsZeroMetrics asserts account writes no tokens counter
|
||||
// when usage totals 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)
|
||||
|
||||
tok, err := l.getCounter("user", "z", metricTok, windowDay, now)
|
||||
if err != nil || tok != 0 {
|
||||
t.Errorf("tok counter = %d, %v, want 0, nil", tok, err)
|
||||
}
|
||||
cost, err := l.getCounter("user", "z", metricCost, windowDay, now)
|
||||
if err != nil || cost != 0 {
|
||||
t.Errorf("cost counter = %d, %v, want 0, nil", cost, err)
|
||||
}
|
||||
}
|
||||
+90
@@ -0,0 +1,90 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import "sync"
|
||||
|
||||
// tokensPerMillion is the denominator for the providers' per-1M-token
|
||||
// pricing unit.
|
||||
const tokensPerMillion = 1_000_000
|
||||
|
||||
// builtinPricing is the built-in per-model price table, in USD per 1M
|
||||
// tokens (InputPerM, OutputPerM). Source: public list prices, 2026-08;
|
||||
// verify before release.
|
||||
var builtinPricing = map[string]ModelPricing{
|
||||
"gpt-5": {InputPerM: 1.25, OutputPerM: 10},
|
||||
"gpt-5-mini": {InputPerM: 0.25, OutputPerM: 2},
|
||||
"gpt-4o": {InputPerM: 2.5, OutputPerM: 10},
|
||||
"o3": {InputPerM: 2, OutputPerM: 8},
|
||||
"claude-opus-4-5": {InputPerM: 5, OutputPerM: 25},
|
||||
"claude-sonnet-4-5": {InputPerM: 3, OutputPerM: 15},
|
||||
"claude-haiku-4-5": {InputPerM: 1, OutputPerM: 5},
|
||||
"gemini-2.5-pro": {InputPerM: 1.25, OutputPerM: 10},
|
||||
"gemini-2.5-flash": {InputPerM: 0.3, OutputPerM: 2.5},
|
||||
"grok-4": {InputPerM: 3, OutputPerM: 15},
|
||||
}
|
||||
|
||||
// pricingWarnFn is called once per unknown model id encountered by
|
||||
// costMicros, with the model id. The default is a no-op so pricing.go has
|
||||
// no logging dependency of its own; a later task wires this to the
|
||||
// gateway's logf, and tests can substitute their own func to capture
|
||||
// warnings.
|
||||
var pricingWarnFn = func(string) {}
|
||||
|
||||
// warnedModels tracks which unknown model ids have already triggered
|
||||
// pricingWarnFn, so a gateway processing many requests for the same
|
||||
// unpriced model logs the warning once per process lifetime rather than
|
||||
// once per request. Package-level and mutex-guarded by design: costMicros
|
||||
// is called from the request path with no natural place to thread a
|
||||
// longer-lived cache through, and the dedup set must survive across
|
||||
// requests and across limiter instances, not reset per call.
|
||||
var (
|
||||
warnedModelsMu sync.Mutex
|
||||
warnedModels = map[string]bool{}
|
||||
)
|
||||
|
||||
// costMicros computes the micro-USD cost of u's token usage for model,
|
||||
// looking up its price in overrides first (exact model id match), then
|
||||
// builtinPricing. An unknown model costs 0 and triggers one pricingWarnFn
|
||||
// call per model id.
|
||||
//
|
||||
// The math stays in integer micro-USD throughout to avoid float drift:
|
||||
// priceMicrosPerM is the per-1M-token price scaled to micro-USD, and each
|
||||
// of input/output cost is tokens*priceMicrosPerM/1e6, summed.
|
||||
func costMicros(model string, u usage, overrides map[string]*ModelPricing) int64 {
|
||||
price, ok := lookupPricing(model, overrides)
|
||||
if !ok {
|
||||
warnUnknownModel(model)
|
||||
return 0
|
||||
}
|
||||
|
||||
inputMicrosPerM := int64(price.InputPerM * tokensPerMillion)
|
||||
outputMicrosPerM := int64(price.OutputPerM * tokensPerMillion)
|
||||
|
||||
inputCost := u.prompt * inputMicrosPerM / tokensPerMillion
|
||||
outputCost := u.completion * outputMicrosPerM / tokensPerMillion
|
||||
return inputCost + outputCost
|
||||
}
|
||||
|
||||
// lookupPricing resolves model's price: overrides first (a nil entry does
|
||||
// not count as a match, so a caller can't accidentally price a model at
|
||||
// zero by leaving a map entry nil), then builtinPricing.
|
||||
func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPricing, bool) {
|
||||
if p, ok := overrides[model]; ok && p != nil {
|
||||
return *p, true
|
||||
}
|
||||
if p, ok := builtinPricing[model]; ok {
|
||||
return p, true
|
||||
}
|
||||
return ModelPricing{}, false
|
||||
}
|
||||
|
||||
// warnUnknownModel calls pricingWarnFn for model, once per process
|
||||
// lifetime.
|
||||
func warnUnknownModel(model string) {
|
||||
warnedModelsMu.Lock()
|
||||
defer warnedModelsMu.Unlock()
|
||||
if warnedModels[model] {
|
||||
return
|
||||
}
|
||||
warnedModels[model] = true
|
||||
pricingWarnFn(model)
|
||||
}
|
||||
+118
@@ -0,0 +1,118 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestCostMicrosBuiltin exercises the integer micro-USD math against
|
||||
// builtinPricing entries directly, so a regression to float-based pricing
|
||||
// (drift) or a wrong tokensPerMillion divisor shows up as an exact
|
||||
// mismatch, not a tolerance check.
|
||||
func TestCostMicrosBuiltin(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
u usage
|
||||
want int64
|
||||
}{
|
||||
// 1M prompt tokens @ $1.25/M + 1M completion tokens @ $10/M = $11.25.
|
||||
{"gpt-5", usage{prompt: 1_000_000, completion: 1_000_000}, 11_250_000},
|
||||
// 500 prompt @ $2.5/M + 200 completion @ $10/M.
|
||||
{"gpt-4o", usage{prompt: 500, completion: 200}, 1250 + 2000},
|
||||
// Zero usage costs zero regardless of price.
|
||||
{"claude-opus-4-5", usage{}, 0},
|
||||
// Prompt-only usage.
|
||||
{"o3", usage{prompt: 100_000}, 200_000}, // 100_000 * 2_000_000 / 1_000_000
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.model, func(t *testing.T) {
|
||||
if got := costMicros(c.model, c.u, nil); got != c.want {
|
||||
t.Errorf("costMicros(%q, %+v, nil) = %d, want %d", c.model, c.u, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestCostMicrosOverride asserts an exact-match override takes priority
|
||||
// over a builtin entry for the same model id, and that an override for a
|
||||
// model absent from builtinPricing still prices correctly.
|
||||
func TestCostMicrosOverride(t *testing.T) {
|
||||
overrides := map[string]*ModelPricing{
|
||||
"gpt-5": {InputPerM: 1, OutputPerM: 1}, // overrides the builtin gpt-5 price
|
||||
"in-house-model": {InputPerM: 1, OutputPerM: 2},
|
||||
}
|
||||
u := usage{prompt: 1_000_000, completion: 1_000_000}
|
||||
|
||||
if got, want := costMicros("gpt-5", u, overrides), int64(2_000_000); got != want {
|
||||
t.Errorf("override for known model: costMicros = %d, want %d", got, want)
|
||||
}
|
||||
if got, want := costMicros("in-house-model", u, overrides), int64(3_000_000); got != want {
|
||||
t.Errorf("override for unknown model: costMicros = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCostMicrosOverrideNilEntryFallsBack asserts a nil override map entry
|
||||
// does not shadow the builtin price with a silent zero — it falls through
|
||||
// to builtinPricing, same as if the key were absent.
|
||||
func TestCostMicrosOverrideNilEntryFallsBack(t *testing.T) {
|
||||
overrides := map[string]*ModelPricing{"gpt-5": nil}
|
||||
u := usage{prompt: 1_000_000, completion: 1_000_000}
|
||||
if got, want := costMicros("gpt-5", u, overrides), int64(11_250_000); got != want {
|
||||
t.Errorf("nil override entry: costMicros = %d, want %d (builtin price)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCostMicrosUnknownModelWarnsOnce asserts an unpriced model costs 0
|
||||
// and fires pricingWarnFn exactly once even across repeated calls, then
|
||||
// resumes warning for a second distinct unknown model.
|
||||
func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) {
|
||||
prevWarn, prevWarned := pricingWarnFn, warnedModels
|
||||
t.Cleanup(func() {
|
||||
pricingWarnFn = prevWarn
|
||||
warnedModels = prevWarned
|
||||
})
|
||||
warnedModels = map[string]bool{}
|
||||
|
||||
var warned []string
|
||||
pricingWarnFn = func(model string) { warned = append(warned, model) }
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
if got := costMicros("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 {
|
||||
t.Errorf("iteration %d: costMicros = %d, want 0", i, got)
|
||||
}
|
||||
}
|
||||
if got := costMicros("another-unpriced-model", usage{prompt: 1}, nil); got != 0 {
|
||||
t.Errorf("costMicros = %d, want 0", got)
|
||||
}
|
||||
|
||||
want := []string{"totally-unpriced-model", "another-unpriced-model"}
|
||||
if len(warned) != len(want) {
|
||||
t.Fatalf("warned = %v, want %v", warned, want)
|
||||
}
|
||||
for i, m := range want {
|
||||
if warned[i] != m {
|
||||
t.Errorf("warned[%d] = %q, want %q", i, warned[i], m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBuiltinPricingSeedComplete asserts every model the task brief seeds
|
||||
// is present in builtinPricing with a positive price on both sides — a
|
||||
// missing or zeroed entry would silently price real traffic at $0.
|
||||
func TestBuiltinPricingSeedComplete(t *testing.T) {
|
||||
want := []string{
|
||||
"gpt-5", "gpt-5-mini", "gpt-4o", "o3",
|
||||
"claude-opus-4-5", "claude-sonnet-4-5", "claude-haiku-4-5",
|
||||
"gemini-2.5-pro", "gemini-2.5-flash", "grok-4",
|
||||
}
|
||||
if len(builtinPricing) != len(want) {
|
||||
t.Errorf("builtinPricing has %d entries, want %d", len(builtinPricing), len(want))
|
||||
}
|
||||
for _, m := range want {
|
||||
p, ok := builtinPricing[m]
|
||||
if !ok {
|
||||
t.Errorf("builtinPricing missing %q", m)
|
||||
continue
|
||||
}
|
||||
if p.InputPerM <= 0 || p.OutputPerM <= 0 {
|
||||
t.Errorf("builtinPricing[%q] = %+v, want both prices > 0", m, p)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user