mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
Review fixes (2 Important + 4 folded): 1. GET /admin now serves the HTML shell with no authentication - a browser navigating directly to the URL cannot attach a custom header. The page reads an admin key from sessionStorage and sends it as x-api-key on /admin/api/* fetches, falling back to an inline key-entry form on a missing key or any 401/403. Both JSON routes keep their original 401/403/200 gate. Spec ss4 amended in place. 2. GET /admin/api/usage no longer serializes 6 GETs per user/group scope on the shared Redis connection. respClient gains getBatch (one pipelined round trip); counterStore gains getMulti; limiter gains storeGetMulti mirroring storeGet's fail-open/fail-closed policy. currentUsage now issues one getMulti per scope. 3. Provider baseUrl is stripped of userinfo and query before it is ever echoed in the overview JSON (sanitizeBaseURL), stronger than net/url's own Redacted (which leaves a bare username in place). 4. limiter.redisStatus now also returns the last failure's timestamp, so the dashboard can render "last error (Ns ago)" instead of a bare message that reads as currently-broken indefinitely. 5. Both JSON admin responses now set X-Content-Type-Options: nosniff, Cache-Control: no-store, and the same CSP the HTML page uses. 6. Added a test proving a file-sourced (replaceFileUsers) admin user gets the same access an inline admin user does. go test ./... -count=1 -race, gofmt, vet, make yaegi-check all pass.
759 lines
29 KiB
Go
759 lines
29 KiB
Go
package traefikllmgateway
|
||
|
||
import (
|
||
"fmt"
|
||
"math"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// storeErrorLogEvery bounds how often the limiter logs a configured
|
||
// store's operation error: once per this interval, regardless of how many
|
||
// operations fail in between. A backend outage under load must not flood
|
||
// the log with one line per request.
|
||
const storeErrorLogEvery = 30 * time.Second
|
||
|
||
// storeDownLatchFor bounds how long the limiter avoids a configured store
|
||
// after any operation against it fails: for this long after the failure,
|
||
// every store operation short-circuits straight to the fail-open/
|
||
// fail-closed policy without attempting the network call at all. Each
|
||
// respClient call is already bounded to ~respCallTimeout on its own
|
||
// (resp.go), but one request can still touch several counters (its own
|
||
// and its group's, request/token/cost, day/month) — without this latch,
|
||
// every one of those pays a fresh bounded wait against a store that just
|
||
// told the previous op it was unreachable. After the window, the next
|
||
// operation probes the real store again: a natural half-open retry, no
|
||
// separate recovery state needed.
|
||
const storeDownLatchFor = 5 * time.Second
|
||
|
||
// 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
|
||
|
||
// usdToMicros converts a USD amount to micro-USD, rounded to the nearest
|
||
// micro-USD. A bare int64(usd*usdToMicroFactor) truncates rather than
|
||
// rounds — float64 can represent a value like 8.2 as very slightly under
|
||
// its true value, so a truncating cast turns an intended $8.20 limit into
|
||
// 8_199_999 micro-USD instead of 8_200_000, one micro-USD too strict.
|
||
func usdToMicros(usd float64) int64 {
|
||
return int64(math.Round(usd * usdToMicroFactor))
|
||
}
|
||
|
||
// validate reports an error if lc holds a negative int64 count limit, or a
|
||
// cost limit that is negative, NaN, or ±Inf. A nil lc (no limits
|
||
// configured) is valid. Without this check, a negative int64 limit would
|
||
// silently behave as unlimited — every enforcement path below treats
|
||
// limit<=0 as "no limit" — and a NaN/Inf cost limit would propagate
|
||
// through usdToMicros into meaningless comparisons instead of failing
|
||
// loudly at construction time. Called by newAuthStore for every group's
|
||
// and every user's LimitsConfig, so a malformed limit is a constructor
|
||
// error rather than a silent misconfiguration.
|
||
func (lc *LimitsConfig) validate() error {
|
||
if lc == nil {
|
||
return nil
|
||
}
|
||
counts := []struct {
|
||
name string
|
||
v int64
|
||
}{
|
||
{"requestsPerMinute", lc.RequestsPerMinute},
|
||
{"requestsPerDay", lc.RequestsPerDay},
|
||
{"tokensPerDay", lc.TokensPerDay},
|
||
{"tokensPerMonth", lc.TokensPerMonth},
|
||
}
|
||
for _, c := range counts {
|
||
if c.v < 0 {
|
||
return fmt.Errorf("llmgateway: limits: %s must not be negative, got %d", c.name, c.v)
|
||
}
|
||
}
|
||
|
||
costs := []struct {
|
||
name string
|
||
v float64
|
||
}{
|
||
{"costPerDayUSD", lc.CostPerDayUSD},
|
||
{"costPerMonthUSD", lc.CostPerMonthUSD},
|
||
}
|
||
for _, c := range costs {
|
||
if math.IsNaN(c.v) || math.IsInf(c.v, 0) {
|
||
return fmt.Errorf("llmgateway: limits: %s must be a finite number, got %v", c.name, c.v)
|
||
}
|
||
if c.v < 0 {
|
||
return fmt.Errorf("llmgateway: limits: %s must not be negative, got %v", c.name, c.v)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// 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)
|
||
// getMulti reads every key in keys in one round trip where the
|
||
// backend supports it (redisStore pipelines via respClient.getBatch;
|
||
// memoryStore's in-process map needs no such optimization but
|
||
// implements the same contract for interface conformance), returning
|
||
// one value per key in the same order — 0 for a key that does not
|
||
// exist or has expired, matching get's contract. An error fails the
|
||
// whole batch (mirrored by the limiter's storeGetMulti as a single
|
||
// fail-open/fail-closed decision), never a partial result.
|
||
getMulti(keys []string) ([]int64, error)
|
||
}
|
||
|
||
// memoryEntry is one counter's value and expiry in memoryStore.
|
||
type memoryEntry struct {
|
||
expiry time.Time
|
||
value int64
|
||
}
|
||
|
||
// sweepEvery is the minimum real time between incrBy's opportunistic
|
||
// expired-entry sweeps, regardless of how many keys memoryStore holds. A
|
||
// size-gated sweep (only above N keys) was measured to scan the whole map
|
||
// on every incrBy once live keys exceeded that threshold, since a sweep
|
||
// that frees nothing (all keys still live) never brings the count back
|
||
// down — a 167x incrBy cliff. Gating on elapsed time instead bounds sweep
|
||
// frequency independent of key count: worst case is one full-map scan
|
||
// every 30s.
|
||
const sweepEvery = 30 * time.Second
|
||
|
||
// 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, at most once every sweepEvery.
|
||
type memoryStore struct {
|
||
data map[string]*memoryEntry
|
||
nowFn func() time.Time // injected for tests; defaults to time.Now
|
||
lastSweep time.Time // guarded by mu; zero value sweeps on the first incrBy
|
||
mu sync.Mutex
|
||
}
|
||
|
||
// newMemoryStore returns an empty memoryStore.
|
||
func newMemoryStore() *memoryStore {
|
||
return &memoryStore{data: make(map[string]*memoryEntry), nowFn: time.Now}
|
||
}
|
||
|
||
// now returns m's current time, via nowFn.
|
||
func (m *memoryStore) now() time.Time {
|
||
return m.nowFn()
|
||
}
|
||
|
||
// incrBy implements counterStore.
|
||
func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, error) {
|
||
now := m.now()
|
||
m.mu.Lock()
|
||
defer m.mu.Unlock()
|
||
|
||
if now.Sub(m.lastSweep) >= sweepEvery {
|
||
m.sweepLocked(now)
|
||
m.lastSweep = 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 := m.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
|
||
}
|
||
|
||
// getMulti implements counterStore by looping over get: memoryStore's
|
||
// in-process map access is already effectively free per key, so there is
|
||
// no round-trip cost to batch away — this exists purely so memoryStore
|
||
// satisfies counterStore's getMulti contract for the limiter's nil-store
|
||
// fallback path (storeGetMulti).
|
||
func (m *memoryStore) getMulti(keys []string) ([]int64, error) {
|
||
out := make([]int64, len(keys))
|
||
for i, k := range keys {
|
||
out[i], _ = m.get(k) // memoryStore.get never errors
|
||
}
|
||
return out, 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 true when prompt/completion were derived from a
|
||
// tokenizer estimate (request body size ÷ 4) rather than the
|
||
// provider's own reported usage. routes_unified.go's runUnified reads
|
||
// this to log the estimate explicitly, so a caller reading logs can
|
||
// tell an estimated cost from a provider-reported one, per the spec's
|
||
// "logged as estimated" requirement.
|
||
estimated bool
|
||
}
|
||
|
||
// 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. Left at its zero
|
||
// value for a storeDown violation, which carries no meaningful window.
|
||
retryAfter int
|
||
// storeDown reports whether the violation was manufactured because the
|
||
// configured counterStore was unreachable and failOpen is false,
|
||
// rather than an actual limit breach. Callers map this to an HTTP 503
|
||
// instead of 429 (Task 12).
|
||
storeDown bool
|
||
}
|
||
|
||
// limiter enforces per-minute/day/month request, token, and cost limits
|
||
// using fixed windows keyed by windowKey.
|
||
type limiter struct {
|
||
store counterStore // configured backend; 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
|
||
// logf is a bound method value (g.errorf), injectable for tests; defaults
|
||
// to a no-op. Its own call site (logStoreError) passes exactly one
|
||
// variadic argument — never extend that to two or more without first
|
||
// reading modelRegistry.log's doc comment in registry.go: a struct
|
||
// field of variadic func type crashes Yaegi v0.16.1's CFG builder past
|
||
// one variadic argument, even though the identical call through a
|
||
// method or interface method does not.
|
||
logf func(format string, args ...any)
|
||
lastLogAt time.Time // guarded by logMu; last time a store error was logged
|
||
lastStoreFailure time.Time // guarded by logMu; zero means the store-down latch is not open (see storeLatched)
|
||
// lastErrMsg is the message of the most recent store operation
|
||
// failure, guarded by logMu alongside lastStoreFailure. It is never
|
||
// cleared on a later success — "last store error" for the admin
|
||
// dashboard (spec §4, v0.2) means exactly that, a persisting fact,
|
||
// not "is the store currently failing" (storeLatched already answers
|
||
// that question for the enforcement path).
|
||
lastErrMsg string
|
||
logMu sync.Mutex
|
||
failOpen bool // store-error policy: true falls back to fallback, false refuses the request
|
||
}
|
||
|
||
// 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. failOpen governs what happens
|
||
// when a non-nil store errors: true silently falls back to the in-process
|
||
// memoryStore for that operation, false refuses the request (see
|
||
// storeIncrBy / storeGet).
|
||
func newLimiter(store counterStore, failOpen bool) *limiter {
|
||
return &limiter{
|
||
store: store,
|
||
fallback: newMemoryStore(),
|
||
nowFn: time.Now,
|
||
logf: func(string, ...any) {},
|
||
failOpen: failOpen,
|
||
}
|
||
}
|
||
|
||
// now returns the limiter's current time, via nowFn.
|
||
func (l *limiter) now() time.Time {
|
||
return l.nowFn()
|
||
}
|
||
|
||
// logStoreError logs err via l.logf, rate-limited to once per
|
||
// storeErrorLogEvery regardless of how many store operations fail in
|
||
// between.
|
||
func (l *limiter) logStoreError(err error) {
|
||
now := l.now()
|
||
l.logMu.Lock()
|
||
shouldLog := now.Sub(l.lastLogAt) >= storeErrorLogEvery
|
||
if shouldLog {
|
||
l.lastLogAt = now
|
||
}
|
||
l.logMu.Unlock()
|
||
if shouldLog {
|
||
l.logf("limit store error: %v", err)
|
||
}
|
||
}
|
||
|
||
// storeLatched reports whether l is within storeDownLatchFor of its last
|
||
// recorded store failure. When true, storeIncrBy/storeGet must skip the
|
||
// network call entirely and go straight to the fail-open/fail-closed
|
||
// policy — the same outcome a fresh call would reach anyway, just without
|
||
// paying respCallTimeout again to find out. Guarded by logMu, the same
|
||
// mutex lastLogAt already uses for this kind of small timestamp state.
|
||
func (l *limiter) storeLatched() bool {
|
||
l.logMu.Lock()
|
||
defer l.logMu.Unlock()
|
||
return !l.lastStoreFailure.IsZero() && l.now().Sub(l.lastStoreFailure) < storeDownLatchFor
|
||
}
|
||
|
||
// recordStoreFailure logs err (rate-limited, see logStoreError) and opens
|
||
// the store-down latch: every storeIncrBy/storeGet call for the next
|
||
// storeDownLatchFor skips the network call and applies the fail-open/
|
||
// fail-closed policy directly.
|
||
func (l *limiter) recordStoreFailure(err error) {
|
||
l.logStoreError(err)
|
||
l.logMu.Lock()
|
||
l.lastStoreFailure = l.now()
|
||
l.lastErrMsg = err.Error()
|
||
l.logMu.Unlock()
|
||
}
|
||
|
||
// redisStatus reports whether l has a configured (non-fallback-only)
|
||
// store, that store's most recent operation failure message, and when
|
||
// that failure happened — the admin dashboard's redis status line (spec
|
||
// §4, v0.2). lastErr is "" and lastErrAt is the zero time when no store
|
||
// operation has ever failed. Returning the timestamp alongside the
|
||
// message (controller-approved amendment, 2026-08-20 review) lets the
|
||
// dashboard render "last error (12s ago): ..." instead of a bare message
|
||
// that reads as "currently broken" even long after a single transient
|
||
// blip. Guarded by logMu, the same mutex lastStoreFailure/lastErrMsg
|
||
// already use.
|
||
func (l *limiter) redisStatus() (configured bool, lastErr string, lastErrAt time.Time) {
|
||
l.logMu.Lock()
|
||
defer l.logMu.Unlock()
|
||
return l.store != nil, l.lastErrMsg, l.lastStoreFailure
|
||
}
|
||
|
||
// storeIncrBy increments key by n with the given ttl, applying the
|
||
// limiter's fail-open/fail-closed policy when a configured store errors
|
||
// or when the store-down latch (storeLatched) is already open from a
|
||
// recent failure. ok is false only in the fail-closed case — a caller
|
||
// must refuse the request for that, rather than treating a zero value as
|
||
// a real counter reading.
|
||
func (l *limiter) storeIncrBy(key string, n int64, ttl time.Duration) (v int64, ok bool) {
|
||
if l.store == nil {
|
||
v, _ = l.fallback.incrBy(key, n, ttl) // fallback never errors
|
||
return v, true
|
||
}
|
||
if l.storeLatched() {
|
||
return l.failPolicyIncrBy(key, n, ttl)
|
||
}
|
||
v, err := l.store.incrBy(key, n, ttl)
|
||
if err == nil {
|
||
return v, true
|
||
}
|
||
l.recordStoreFailure(err)
|
||
return l.failPolicyIncrBy(key, n, ttl)
|
||
}
|
||
|
||
// storeGet mirrors storeIncrBy for a read: it applies the same fail-open/
|
||
// fail-closed/latched handling on a configured store's error.
|
||
func (l *limiter) storeGet(key string) (v int64, ok bool) {
|
||
if l.store == nil {
|
||
v, _ = l.fallback.get(key)
|
||
return v, true
|
||
}
|
||
if l.storeLatched() {
|
||
return l.failPolicyGet(key)
|
||
}
|
||
v, err := l.store.get(key)
|
||
if err == nil {
|
||
return v, true
|
||
}
|
||
l.recordStoreFailure(err)
|
||
return l.failPolicyGet(key)
|
||
}
|
||
|
||
// failPolicyIncrBy applies the limiter's fail-open/fail-closed policy for
|
||
// an incrBy the caller has decided not to attempt against the configured
|
||
// store (it just failed, or the store-down latch is open): failOpen=true
|
||
// counts key in the fallback instead; failOpen=false reports the
|
||
// operation as failed.
|
||
func (l *limiter) failPolicyIncrBy(key string, n int64, ttl time.Duration) (int64, bool) {
|
||
if !l.failOpen {
|
||
return 0, false
|
||
}
|
||
v, _ := l.fallback.incrBy(key, n, ttl)
|
||
return v, true
|
||
}
|
||
|
||
// failPolicyGet mirrors failPolicyIncrBy for a read.
|
||
func (l *limiter) failPolicyGet(key string) (int64, bool) {
|
||
if !l.failOpen {
|
||
return 0, false
|
||
}
|
||
v, _ := l.fallback.get(key)
|
||
return v, true
|
||
}
|
||
|
||
// storeGetMulti mirrors storeGet for a batch of keys: one round trip
|
||
// against the configured store (or the in-process fallback) for the
|
||
// whole slice, applying the identical fail-open/fail-closed/latched
|
||
// policy storeGet applies per key. ok is false only in the fail-closed
|
||
// case, matching storeGet's contract — a caller must not read the
|
||
// returned slice as real values when ok is false.
|
||
func (l *limiter) storeGetMulti(keys []string) (v []int64, ok bool) {
|
||
if l.store == nil {
|
||
v, _ = l.fallback.getMulti(keys) // fallback never errors
|
||
return v, true
|
||
}
|
||
if l.storeLatched() {
|
||
return l.failPolicyGetMulti(keys)
|
||
}
|
||
v, err := l.store.getMulti(keys)
|
||
if err == nil {
|
||
return v, true
|
||
}
|
||
l.recordStoreFailure(err)
|
||
return l.failPolicyGetMulti(keys)
|
||
}
|
||
|
||
// failPolicyGetMulti mirrors failPolicyGet for a batch read.
|
||
func (l *limiter) failPolicyGetMulti(keys []string) ([]int64, bool) {
|
||
if !l.failOpen {
|
||
return nil, false
|
||
}
|
||
v, _ := l.fallback.getMulti(keys)
|
||
return v, true
|
||
}
|
||
|
||
// storeDownViolation is the violation checkAndCount/budgetViolation return
|
||
// when the configured store is unreachable and failOpen is false: the
|
||
// request is refused instead of silently enforcing limits against a
|
||
// non-shared fallback, so a backend outage cannot let every configured
|
||
// limit go unenforced across a fleet of gateway instances.
|
||
func storeDownViolation() *limitViolation {
|
||
return &limitViolation{message: "limit store unavailable", storeDown: true}
|
||
}
|
||
|
||
// 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, together with whether the operation
|
||
// succeeded under the limiter's fail-open/fail-closed policy (see
|
||
// storeIncrBy). ok is false only in the fail-closed case: a configured
|
||
// store errored and failOpen is false.
|
||
func (l *limiter) incrCounter(kind, id, metric, window string, t time.Time, n int64, ttl time.Duration) (int64, bool) {
|
||
return l.storeIncrBy(windowKey(kind, id, metric, window, t), n, ttl)
|
||
}
|
||
|
||
// getCounter returns the (kind, id, metric, window) counter's value at t,
|
||
// together with whether the read succeeded under the limiter's fail-open/
|
||
// fail-closed policy (see storeGet).
|
||
func (l *limiter) getCounter(kind, id, metric, window string, t time.Time) (int64, bool) {
|
||
return l.storeGet(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. If any store operation fails closed (a configured store errored
|
||
// and failOpen is false), it returns a storeDownViolation immediately —
|
||
// the request is refused rather than evaluated against partial or
|
||
// fallback-only counters.
|
||
//
|
||
// 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 {
|
||
v, ok := l.incrCounter(sc.kind, sc.id, metricReq, windowMin, now, 1, minWindowTTL)
|
||
if !ok {
|
||
return storeDownViolation()
|
||
}
|
||
minCounts[i] = v
|
||
|
||
v, ok = l.incrCounter(sc.kind, sc.id, metricReq, windowDay, now, 1, dayWindowTTL)
|
||
if !ok {
|
||
return storeDownViolation()
|
||
}
|
||
dayCounts[i] = v
|
||
}
|
||
|
||
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 for this scope. In
|
||
// principle a nil-limits scope can still reach here with its req counters
|
||
// already incremented by checkAndCount, since that increment runs before
|
||
// any evaluation; in practice the unified route's buildLimitScopes
|
||
// (routes_unified.go, ruling e) never builds one — it omits a user or
|
||
// group scope from the slice entirely whenever that scope's own limits
|
||
// are nil — so this check is a defensive no-op today, not a path any
|
||
// current caller exercises.
|
||
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", usdToMicros(lim.CostPerDayUSD), windowDay, now); v != nil {
|
||
return v
|
||
}
|
||
if v := l.budgetViolation(sc, metricCost, "cost-per-month", usdToMicros(lim.CostPerMonthUSD), 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 failed-closed read (configured store errored, failOpen
|
||
// false) returns storeDownViolation rather than silently passing the
|
||
// check — a store outage must never look the same as staying under
|
||
// budget.
|
||
func (l *limiter) budgetViolation(sc limitScope, metric, name string, limit int64, window string, now time.Time) *limitViolation {
|
||
if limit <= 0 {
|
||
return nil
|
||
}
|
||
used, ok := l.getCounter(sc.kind, sc.id, metric, window, now)
|
||
if !ok {
|
||
return storeDownViolation()
|
||
}
|
||
if 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. A sample that
|
||
// hits a fail-closed store error (see storeIncrBy) is dropped silently:
|
||
// account has no error return to signal it, and dropping — rather than
|
||
// counting it in the fallback — avoids double-counting once the store
|
||
// recovers.
|
||
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)
|
||
}
|
||
}
|
||
}
|
||
|
||
// scopeUsage is one scope's (a user's or a group's) current-window
|
||
// counter values — the admin dashboard's usage table (spec §4, v0.2). The
|
||
// six value fields mirror LimitsConfig's six limit kinds exactly, so the
|
||
// admin API can echo a usage value beside its matching limit.
|
||
type scopeUsage struct {
|
||
kind string // "user" or "group", mirroring limitScope.kind
|
||
id string
|
||
requestsPerMinute int64
|
||
requestsPerDay int64
|
||
tokensPerDay int64
|
||
tokensPerMonth int64
|
||
costPerDayMicros int64
|
||
costPerMonthMicros int64
|
||
// storeDown reports whether reading any of the six counters above
|
||
// failed closed (a configured store errored and failOpen is false) —
|
||
// mirrors limitViolation.storeDown. Every value field is 0 in that
|
||
// case; a caller must not present them as "confirmed zero usage".
|
||
storeDown bool
|
||
}
|
||
|
||
// usageWindowKeys returns the six windowKey strings currentUsage reads
|
||
// for sc at time now, in a fixed order — req/min, req/day, tok/day,
|
||
// tok/month, cost/day, cost/month — matching scopeUsage's field order
|
||
// exactly, so currentUsage can map storeGetMulti's result slice back to
|
||
// named fields by plain index.
|
||
func usageWindowKeys(sc limitScope, now time.Time) []string {
|
||
return []string{
|
||
windowKey(sc.kind, sc.id, metricReq, windowMin, now),
|
||
windowKey(sc.kind, sc.id, metricReq, windowDay, now),
|
||
windowKey(sc.kind, sc.id, metricTok, windowDay, now),
|
||
windowKey(sc.kind, sc.id, metricTok, windowMonth, now),
|
||
windowKey(sc.kind, sc.id, metricCost, windowDay, now),
|
||
windowKey(sc.kind, sc.id, metricCost, windowMonth, now),
|
||
}
|
||
}
|
||
|
||
// currentUsage reads every scope's six current-window counters — the
|
||
// same (kind, id, metric, window) combinations checkAndCount/account
|
||
// already write — via the limiter's own storeGetMulti, so it applies the
|
||
// identical fail-open/fail-closed policy every enforcement read already
|
||
// does, and it is read-only: unlike checkAndCount, it never increments
|
||
// anything.
|
||
//
|
||
// One storeGetMulti call per scope (controller-approved amendment,
|
||
// 2026-08-20 review): the original per-metric l.getCounter loop issued
|
||
// six separate round trips per scope against the single shared Redis
|
||
// connection — GET /admin/api/usage serialized 6×(users+groups) of them.
|
||
// Batching all scopes' keys into one single storeGetMulti call would cut
|
||
// that further, to one round trip total, but was not taken here: it
|
||
// would need buildAdminUsage to flatten and later re-split two
|
||
// differently-sized scope lists (users, groups) around one shared call,
|
||
// for a further win only realized when both lists are non-trivially
|
||
// large. One call per scope already turns 6N round trips into N, keeps
|
||
// currentUsage's per-scope iteration shape identical to
|
||
// checkAndCount/evaluateScope's, and keeps a single scope's store
|
||
// failure from forcing every other scope in the same request to look
|
||
// storeDown too. Order is preserved: currentUsage(scopes)[i] corresponds
|
||
// to scopes[i].
|
||
func (l *limiter) currentUsage(scopes []limitScope) []scopeUsage {
|
||
now := l.now()
|
||
out := make([]scopeUsage, len(scopes))
|
||
for i, sc := range scopes {
|
||
vals, ok := l.storeGetMulti(usageWindowKeys(sc, now))
|
||
if !ok {
|
||
out[i] = scopeUsage{kind: sc.kind, id: sc.id, storeDown: true}
|
||
continue
|
||
}
|
||
out[i] = scopeUsage{
|
||
kind: sc.kind, id: sc.id,
|
||
requestsPerMinute: vals[0], requestsPerDay: vals[1],
|
||
tokensPerDay: vals[2], tokensPerMonth: vals[3],
|
||
costPerDayMicros: vals[4], costPerMonthMicros: vals[5],
|
||
}
|
||
}
|
||
return out
|
||
}
|