fix(limits): cap unknown-model warnings, time-gated sweep, rounded micro-USD conversion

limits.go

@@ -38,6 +38,15 @@ const (
 // 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))
+}
+
 // 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.
@@ -57,37 +66,48 @@ type memoryEntry struct {
 	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
+// 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 once the map grows past
-// memoryStoreSweepThreshold keys.
+// opportunistically from incrBy, at most once every sweepEvery.
 type memoryStore struct {
-	data map[string]*memoryEntry
-	mu   sync.Mutex
+	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)}
+	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 := time.Now()
+	now := m.now()
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	if len(m.data) > memoryStoreSweepThreshold {
+	if now.Sub(m.lastSweep) >= sweepEvery {
 		m.sweepLocked(now)
+		m.lastSweep = now
 	}
 
 	e, ok := m.data[key]
@@ -101,7 +121,7 @@ func (m *memoryStore) incrBy(key string, n int64, ttl time.Duration) (int64, err
 
 // get implements counterStore.
 func (m *memoryStore) get(key string) (int64, error) {
-	now := time.Now()
+	now := m.now()
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
@@ -314,10 +334,10 @@ func (l *limiter) evaluateScope(sc limitScope, minCount, dayCount int64, now tim
 	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 {
+	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", int64(lim.CostPerMonthUSD*usdToMicroFactor), windowMonth, now); v != nil {
+	if v := l.budgetViolation(sc, metricCost, "cost-per-month", usdToMicros(lim.CostPerMonthUSD), windowMonth, now); v != nil {
 		return v
 	}
 	return nil

limits_test.go

@@ -1,7 +1,6 @@
 package traefikllmgateway
 
 import (
-	"fmt"
 	"testing"
 	"time"
 )
@@ -227,40 +226,74 @@ func TestMemoryStoreExpiry(t *testing.T) {
 	}
 }
 
-// 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) {
+// TestMemoryStoreSweepTimeGated asserts the sweep gate is purely
+// time-based (sweepEvery), independent of how many keys memoryStore
+// holds: an incrBy call before the gate elapses leaves an already-expired
+// key physically in place (no scan happened), and the first incrBy at or
+// past the gate reclaims it. Time is injected via m.nowFn so the test
+// needs no real sleep and no size threshold to trigger the behavior.
+func TestMemoryStoreSweepTimeGated(t *testing.T) {
 	m := newMemoryStore()
-	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)
-		}
-	}
+	fake := time.Date(2026, 8, 20, 0, 0, 0, 0, time.UTC)
+	m.nowFn = func() time.Time { return fake }
 
-	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 {
+	// lastSweep starts at the zero time, so this first incrBy always
+	// sweeps (a no-op on an empty store) and sets lastSweep to fake.
+	if _, err := m.incrBy("k1", 1, time.Nanosecond); err != nil {
 		t.Fatal(err)
 	}
 
+	// k1 is already expired (ttl was 1ns), but stay well inside the
+	// sweepEvery gate — the next incrBy must not scan, so k1 survives as
+	// a physical (if stale) map entry.
+	fake = fake.Add(time.Millisecond)
+	if _, err := m.incrBy("k2", 1, time.Minute); err != nil {
+		t.Fatal(err)
+	}
+	if !memoryStoreHasKey(m, "k1") {
+		t.Fatal("k1 must still be present: the sweep gate has not elapsed yet")
+	}
+
+	// Cross the sweepEvery boundary from the last sweep: this incrBy must
+	// scan and reclaim k1.
+	fake = fake.Add(sweepEvery)
+	if _, err := m.incrBy("k3", 1, time.Minute); err != nil {
+		t.Fatal(err)
+	}
+	if memoryStoreHasKey(m, "k1") {
+		t.Error("k1 should have been swept once sweepEvery elapsed")
+	}
+}
+
+// memoryStoreHasKey reports whether key is physically present in m.data,
+// bypassing get's expiry check — used to observe sweep timing directly.
+func memoryStoreHasKey(m *memoryStore, key string) bool {
 	m.mu.Lock()
-	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)
+	defer m.mu.Unlock()
+	_, ok := m.data[key]
+	return ok
+}
+
+// TestLimiterCostBudgetRoundsUSDConversion pins usdToMicros' rounding
+// against a value (8.2) where float64 representation sits just under the
+// true value: a truncating int64(8.2*1e6) would undershoot to 8_199_999,
+// making a spend of exactly $8.199999 look already over an intended
+// $8.20 budget. With rounding, 8_199_999 must stay under budget and only
+// the exact $8.20 spend violates.
+func TestLimiterCostBudgetRoundsUSDConversion(t *testing.T) {
+	now := time.Date(2026, 8, 20, 10, 0, 0, 0, time.UTC)
+	l := newLimiter(nil, true)
+	l.nowFn = func() time.Time { return now }
+	scopes := []limitScope{{kind: "user", id: "u", limits: &LimitsConfig{CostPerDayUSD: 8.2}}}
+
+	l.account(scopes, usage{}, 8_199_999)
+	if v := l.checkAndCount(scopes); v != nil {
+		t.Fatalf("spending $8.199999 against an $8.20 budget must not violate, got %+v", v)
+	}
+
+	l.account(scopes, usage{}, 1) // now exactly at $8.20
+	if v := l.checkAndCount(scopes); v == nil {
+		t.Fatal("spending exactly $8.20 against an $8.20 budget must violate")
 	}
 }
 

pricing.go

@@ -1,6 +1,9 @@
 package traefikllmgateway
 
-import "sync"
+import (
+	"math"
+	"sync"
+)
 
 // tokensPerMillion is the denominator for the providers' per-1M-token
 // pricing unit.
@@ -22,33 +25,55 @@ var builtinPricing = map[string]ModelPricing{
 	"grok-4":            {InputPerM: 3, OutputPerM: 15},
 }
 
+// unknownModelWarnCap bounds warnedModels. The model id in costMicros'
+// calls comes from client-supplied request bodies (a group with an empty
+// Models allowlist accepts any model string), so without a cap an attacker
+// or a buggy client cycling through model names could grow warnedModels
+// without bound. Once the cap is reached, costMicros stops recording new
+// names and pricingWarnFn fires once more with a summary message instead
+// of a per-model one.
+const unknownModelWarnCap = 128
+
+// warnCapMessage is the single summary warning fired once the dedup set
+// in warnedModels reaches unknownModelWarnCap.
+const warnCapMessage = "pricing: unknown-model warning cap reached; further unknown models unwarned"
+
 // 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.
+// costMicros, with the model id — or, once unknownModelWarnCap is
+// reached, once more with warnCapMessage. 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 via setPricingWarnFn, and tests can
+// substitute their own func to capture warnings the same way. Read inside
+// warnUnknownModel under warnedModelsMu — always assign through
+// setPricingWarnFn, never by a bare `pricingWarnFn = …`, or a concurrent
+// costMicros call can race the assignment.
 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.
+// once per request. warnCapNotified tracks whether the one-time summary
+// warning has already fired once the cap is reached. 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{}
+	warnedModelsMu  sync.Mutex
+	warnedModels    = map[string]bool{}
+	warnCapNotified 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.
+// call per model id, up to unknownModelWarnCap distinct names.
 //
 // 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.
+// priceMicrosPerM is the per-1M-token price scaled to micro-USD and
+// rounded to the nearest micro-USD (a bare truncating cast can undershoot
+// a price like 8.2 by one micro-USD due to float64 representation), 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 {
@@ -56,8 +81,8 @@ func costMicros(model string, u usage, overrides map[string]*ModelPricing) int64
 		return 0
 	}
 
-	inputMicrosPerM := int64(price.InputPerM * tokensPerMillion)
-	outputMicrosPerM := int64(price.OutputPerM * tokensPerMillion)
+	inputMicrosPerM := int64(math.Round(price.InputPerM * tokensPerMillion))
+	outputMicrosPerM := int64(math.Round(price.OutputPerM * tokensPerMillion))
 
 	inputCost := u.prompt * inputMicrosPerM / tokensPerMillion
 	outputCost := u.completion * outputMicrosPerM / tokensPerMillion
@@ -77,14 +102,34 @@ func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPrici
 	return ModelPricing{}, false
 }
 
+// setPricingWarnFn replaces pricingWarnFn, guarded by the same mutex
+// warnUnknownModel reads it under. A later task calls this at New() to
+// wire warnings to the gateway's logf; tests call it to capture warnings.
+// Always use this instead of assigning pricingWarnFn directly.
+func setPricingWarnFn(fn func(string)) {
+	warnedModelsMu.Lock()
+	defer warnedModelsMu.Unlock()
+	pricingWarnFn = fn
+}
+
 // warnUnknownModel calls pricingWarnFn for model, once per process
-// lifetime.
+// lifetime. Once warnedModels reaches unknownModelWarnCap, it stops
+// recording new model ids and instead fires a single one-time summary
+// warning via warnCapMessage — after that, further unknown models are
+// silently unwarned rather than growing the set forever.
 func warnUnknownModel(model string) {
 	warnedModelsMu.Lock()
 	defer warnedModelsMu.Unlock()
 	if warnedModels[model] {
 		return
 	}
+	if len(warnedModels) >= unknownModelWarnCap {
+		if !warnCapNotified {
+			warnCapNotified = true
+			pricingWarnFn(warnCapMessage)
+		}
+		return
+	}
 	warnedModels[model] = true
 	pricingWarnFn(model)
 }

pricing_test.go

@@ -1,6 +1,10 @@
 package traefikllmgateway
 
-import "testing"
+import (
+	"strconv"
+	"sync"
+	"testing"
+)
 
 // TestCostMicrosBuiltin exercises the integer micro-USD math against
 // builtinPricing entries directly, so a regression to float-based pricing
@@ -59,19 +63,42 @@ func TestCostMicrosOverrideNilEntryFallsBack(t *testing.T) {
 	}
 }
 
+// resetPricingWarnState snapshots and clears the package-level unknown-
+// model dedup state (warnedModels, warnCapNotified) and pricingWarnFn, and
+// registers a t.Cleanup to restore them — so tests exercising the
+// warn-once/warn-cap behavior are independent of each other and of test
+// order. All access goes through setPricingWarnFn / warnedModelsMu so
+// these tests stay race-clean.
+func resetPricingWarnState(t *testing.T) *[]string {
+	t.Helper()
+	warnedModelsMu.Lock()
+	prevWarn, prevWarned, prevCapNotified := pricingWarnFn, warnedModels, warnCapNotified
+	warnedModels = map[string]bool{}
+	warnCapNotified = false
+	warnedModelsMu.Unlock()
+
+	t.Cleanup(func() {
+		warnedModelsMu.Lock()
+		warnedModels, warnCapNotified = prevWarned, prevCapNotified
+		warnedModelsMu.Unlock()
+		setPricingWarnFn(prevWarn)
+	})
+
+	var warned []string
+	var mu sync.Mutex
+	setPricingWarnFn(func(model string) {
+		mu.Lock()
+		defer mu.Unlock()
+		warned = append(warned, model)
+	})
+	return &warned
+}
+
 // 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) }
+	warned := resetPricingWarnState(t)
 
 	for i := 0; i < 3; i++ {
 		if got := costMicros("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 {
@@ -83,16 +110,63 @@ func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) {
 	}
 
 	want := []string{"totally-unpriced-model", "another-unpriced-model"}
-	if len(warned) != len(want) {
-		t.Fatalf("warned = %v, want %v", warned, want)
+	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)
+		if (*warned)[i] != m {
+			t.Errorf("warned[%d] = %q, want %q", i, (*warned)[i], m)
 		}
 	}
 }
 
+// TestCostMicrosUnknownModelWarnCap asserts warnedModels never grows past
+// unknownModelWarnCap, however many distinct unknown model ids costMicros
+// sees — the model id comes from client-controlled request bodies (a
+// group with an empty Models allowlist accepts any string), so an
+// unbounded dedup set would be a memory-growth vector. Once the cap is
+// reached, one summary warning fires and no more do.
+func TestCostMicrosUnknownModelWarnCap(t *testing.T) {
+	warned := resetPricingWarnState(t)
+
+	const distinctModels = 200
+	for i := 0; i < distinctModels; i++ {
+		model := "unpriced-model-" + strconv.Itoa(i)
+		if got := costMicros(model, usage{prompt: 1}, nil); got != 0 {
+			t.Fatalf("costMicros(%q, ...) = %d, want 0", model, got)
+		}
+	}
+
+	warnedModelsMu.Lock()
+	setSize := len(warnedModels)
+	warnedModelsMu.Unlock()
+	if setSize > unknownModelWarnCap {
+		t.Errorf("warnedModels size = %d, want <= %d", setSize, unknownModelWarnCap)
+	}
+
+	if got, max := len(*warned), unknownModelWarnCap+1; got > max {
+		t.Errorf("warn count = %d, want <= %d (cap distinct warnings + 1 summary)", got, max)
+	}
+
+	last := (*warned)[len(*warned)-1]
+	if last != warnCapMessage {
+		t.Errorf("last warning = %q, want the cap-reached summary %q", last, warnCapMessage)
+	}
+}
+
+// TestCostMicrosRoundsPriceConversion pins costMicros' per-1M-token price
+// conversion to round-to-nearest rather than truncate: at $8.20/M, a
+// truncating int64(8.2*1e6) undershoots to 8_199_999 (float64 can
+// represent 8.2 as very slightly under its true value); rounding must
+// yield the intended 8_200_000.
+func TestCostMicrosRoundsPriceConversion(t *testing.T) {
+	overrides := map[string]*ModelPricing{"round-test": {InputPerM: 8.2, OutputPerM: 8.2}}
+	u := usage{prompt: 1_000_000, completion: 1_000_000}
+	if got, want := costMicros("round-test", u, overrides), int64(16_400_000); got != want {
+		t.Errorf("costMicros with an 8.2 price = %d, want %d", got, want)
+	}
+}
+
 // 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.