From ac11023a4507f576183918df1c2326746cac26a0 Mon Sep 17 00:00:00 2001 From: Lukasz Raczylo Date: Thu, 20 Aug 2026 01:28:13 +0100 Subject: [PATCH] feat: limit engine with fixed windows, token and cost budgets --- limits.go | 374 ++++++++++++++++++++++++++++++++++++++++++++++++ limits_test.go | 286 ++++++++++++++++++++++++++++++++++++ pricing.go | 90 ++++++++++++ pricing_test.go | 118 +++++++++++++++ 4 files changed, 868 insertions(+) create mode 100644 limits.go create mode 100644 limits_test.go create mode 100644 pricing.go create mode 100644 pricing_test.go diff --git a/limits.go b/limits.go new file mode 100644 index 0000000..2535b00 --- /dev/null +++ b/limits.go @@ -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) + } + } +} diff --git a/limits_test.go b/limits_test.go new file mode 100644 index 0000000..2335aac --- /dev/null +++ b/limits_test.go @@ -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) + } +} diff --git a/pricing.go b/pricing.go new file mode 100644 index 0000000..1d001b8 --- /dev/null +++ b/pricing.go @@ -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) +} diff --git a/pricing_test.go b/pricing_test.go new file mode 100644 index 0000000..a4c7c0a --- /dev/null +++ b/pricing_test.go @@ -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) + } + } +}