mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +00:00
fix(gateway): close billing, limit, routing and translation gaps
Billing: - Passthrough SSE streams now bill the usage the upstream reports. The gateway falls back to a request-size estimate only on a 2xx stream or a truncated JSON body. - A unified stream that fails after it starts is estimated, not billed at zero. - Billing uses the built-in LiteLLM price table and honours modelMeta free: true. - Pricing entries are validated. Cost arithmetic rounds the total and saturates on overflow. - Gemini thinking tokens count as completion tokens. Limits and storage: - A rejected request rolls back its counts on the other scopes, so one user cannot lock out a group. - Redis EXPIRE replies are checked. Keys recreated without a TTL are repaired at once. - A mutating pipeline is not retried after reply bytes arrive. - History, model totals and target usage report a store outage. Routing and translation: - Discovered ids with a provider-like prefix are listed only in owner/id form, so listing and routing agree. - Anthropic and Gemini discovery follows pagination. - The developer role maps to system. - stream_options.include_usage is honoured on translated streams. - tool_choice none keeps tools when history holds tool blocks. - refusal maps to content_filter. Data URI parameters are stripped. Streaming, retry and failover: - SSE accepts CR, LF and CRLF line ends. This closes a CR injection path. - Retry-After is never jittered. A cancelled wait returns the context error. - The watchdog timer stops at end of body. - Oversized bodies return 413. - A failed stream ends with an error event. - Failover restores pre-existing headers. Client cancels do not count as provider successes. Also: "admin" is a reserved provider name. metrics.path must start with "/". Pricing warnings go to the right gateway instance. The admin targets view sends access as null or [] so the two states differ. Upstream organisation headers are stripped. The users-file reload error is logged once per error. Dead code is removed, and test-only helpers move to test files. Claude-Session: https://claude.ai/code/session_01V5jYJNxpDFdJkyyExpADrv
This commit is contained in:
+220
-75
@@ -1,7 +1,9 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"sync"
|
||||
)
|
||||
|
||||
@@ -27,36 +29,25 @@ 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.
|
||||
// unknownModelWarnCap bounds warnedModels. The model id in
|
||||
// costMicrosKnownFor's 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,
|
||||
// warnUnknownModelFor stops recording new names and fires its warn func
|
||||
// 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 — 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
|
||||
// a warn func, so a gateway processing many requests for the same
|
||||
// unpriced model logs the warning once per process lifetime rather than
|
||||
// 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
|
||||
// mutex-guarded by design: costMicrosKnownFor 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.
|
||||
@@ -66,59 +57,162 @@ var (
|
||||
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, up to unknownModelWarnCap distinct names.
|
||||
// costMicrosKnownFor computes the micro-USD cost of u's token usage for
|
||||
// model, and whether a price was actually FOUND. An unknown model returns
|
||||
// (0, false) — a zero that means "cannot be priced", which is NOT the
|
||||
// same value a genuinely free model produces, even though both are 0.
|
||||
// Looks up model's price via lookupPricing (overrides, then
|
||||
// builtinPricing, then builtinModelMetaTable); an unknown model triggers
|
||||
// one warn 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
|
||||
// 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 {
|
||||
cost, _ := costMicrosKnown(model, u, overrides)
|
||||
return cost
|
||||
}
|
||||
|
||||
// costMicrosKnown is costMicros plus the one bit costMicros throws away:
|
||||
// whether a price was actually FOUND. An unknown model returns (0, false)
|
||||
// — a zero that means "cannot be priced", which is NOT the same value a
|
||||
// genuinely free model produces, even though both are 0.
|
||||
// Security audit run-1, finding F-1 (high): an earlier single-int64-
|
||||
// returning costMicros conflated "unpriced" with "genuinely zero cost",
|
||||
// and limiter.account's own `if costMicros != 0` guard then skipped every
|
||||
// cost-counter write for an unpriced model. checkAndCount reads exactly
|
||||
// those never-written counters, so `used < limit` stayed permanently true
|
||||
// and a configured costPerDayUSD/costPerMonthUSD could never fire. Any
|
||||
// caller gating a SPEND control must handle ok == false.
|
||||
//
|
||||
// Security audit run-1, finding F-1 (high): costMicros' single-int64
|
||||
// return conflated those two cases, and limiter.account's own
|
||||
// `if costMicros != 0` guard then skipped every cost-counter write for an
|
||||
// unpriced model. checkAndCount reads exactly those never-written
|
||||
// counters, so `used < limit` stayed permanently true and a configured
|
||||
// costPerDayUSD/costPerMonthUSD could never fire. Any caller gating a
|
||||
// SPEND control must use this function and handle ok == false; costMicros
|
||||
// above remains for callers that only need the number (the accounting
|
||||
// write itself, where a zero is harmless).
|
||||
//
|
||||
// This is the distinction resolvedModelMeta.CostKnown (modelmeta.go)
|
||||
// already draws for the metadata path, and the one candidateCostAllowed
|
||||
// (failover.go) already relies on to protect real money — its own comment
|
||||
// calls recording zero for an unpriced model "an accounting shortcut, not
|
||||
// proof the provider invoices zero".
|
||||
func costMicrosKnown(model string, u usage, overrides map[string]*ModelPricing) (int64, bool) {
|
||||
// warn is the caller's own per-instance sink (a *Gateway's g.pricingWarn,
|
||||
// logger.go), never a package global: Traefik constructs one Gateway per
|
||||
// middleware instance, and a shared global made the most recently
|
||||
// constructed instance log every other instance's warnings (finding F11 /
|
||||
// review-auth F8, 2026-09 review). The unknown-model dedup (warnedModels)
|
||||
// stays process-global, so a model already warned about under one
|
||||
// instance is not warned about again under another.
|
||||
func costMicrosKnownFor(model string, u usage, overrides map[string]*ModelPricing, warn func(string)) (int64, bool) {
|
||||
price, ok := lookupPricing(model, overrides)
|
||||
if !ok {
|
||||
warnUnknownModel(model)
|
||||
warnUnknownModelFor(model, warn)
|
||||
return 0, false
|
||||
}
|
||||
return computeCostMicros(price, u), true
|
||||
}
|
||||
|
||||
// computeCostMicros is costMicrosKnownFor's arithmetic: price's two per-1M-token rates, each 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), multiplied by their own token count, SUMMED, and only
|
||||
// THEN divided back down and rounded to the nearest micro-USD (findings
|
||||
// F9/F10, 2026-09 review).
|
||||
//
|
||||
// F9: the previous version divided input and output separately
|
||||
// (tokens*priceMicrosPerM/tokensPerMillion, truncating each), which
|
||||
// silently billed 0 for any single request under ~50 tokens at a
|
||||
// realistic per-token price — up to ~2 micro-USD of truncation error per
|
||||
// request (one per direction), enough that account (limits.go) skipped
|
||||
// the cost-counter write entirely (`if costMicros != 0`) for genuinely
|
||||
// billable-but-cheap traffic. Rounding the COMBINED total once instead
|
||||
// halves the worst-case error and stops a nonzero combined cost from
|
||||
// truncating away to a written 0.
|
||||
//
|
||||
// F10: the previous unchecked int64 multiply overflowed for a token count
|
||||
// above roughly 1.2e11 at a $75/M price — entirely upstream-controlled,
|
||||
// since the token counts come from the provider's own reported usage
|
||||
// (limits.go's account, security audit run-1) — silently wrapping to a
|
||||
// negative or garbage value; a negative result was then clamped to 0 by
|
||||
// account, permanently masking the true cost. saturatingMulInt64/
|
||||
// saturatingAddInt64 (below) clamp each multiply and the final sum to
|
||||
// [math.MinInt64, math.MaxInt64] instead of wrapping, so an absurd
|
||||
// upstream-reported token count bills the maximum representable cost
|
||||
// rather than a wrapped-around near-zero or negative one.
|
||||
func computeCostMicros(price ModelPricing, u usage) int64 {
|
||||
inputMicrosPerM := int64(math.Round(price.InputPerM * tokensPerMillion))
|
||||
outputMicrosPerM := int64(math.Round(price.OutputPerM * tokensPerMillion))
|
||||
|
||||
inputCost := u.prompt * inputMicrosPerM / tokensPerMillion
|
||||
outputCost := u.completion * outputMicrosPerM / tokensPerMillion
|
||||
return inputCost + outputCost, true
|
||||
inputCost := saturatingMulInt64(u.prompt, inputMicrosPerM)
|
||||
outputCost := saturatingMulInt64(u.completion, outputMicrosPerM)
|
||||
total := saturatingAddInt64(inputCost, outputCost)
|
||||
|
||||
return roundedDivInt64(total, tokensPerMillion)
|
||||
}
|
||||
|
||||
// roundedDivInt64 divides total by denom, rounded to the nearest integer
|
||||
// (half away from zero) rather than truncated toward zero — see
|
||||
// computeCostMicros' own F9 discussion. denom is always tokensPerMillion
|
||||
// (a positive compile-time constant) at this function's one call site.
|
||||
//
|
||||
// Deliberately computed via quotient/remainder rather than the more
|
||||
// obvious "(total + denom/2) / denom": total can legitimately arrive at
|
||||
// math.MaxInt64 (computeCostMicros' own saturatingAddInt64 clamps there
|
||||
// on overflow, finding F10), and adding denom/2 to it before dividing
|
||||
// would itself silently overflow right back — exactly the wrap-around
|
||||
// bug this whole function exists to avoid one step earlier. Working from
|
||||
// the remainder instead never adds to total at all.
|
||||
func roundedDivInt64(total, denom int64) int64 {
|
||||
q := total / denom
|
||||
r := total % denom
|
||||
if r < 0 {
|
||||
r = -r
|
||||
}
|
||||
if r*2 >= denom {
|
||||
if total >= 0 {
|
||||
q++
|
||||
} else {
|
||||
q--
|
||||
}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
// saturatingMulInt64 returns a*b, clamped to math.MinInt64/math.MaxInt64
|
||||
// instead of wrapping on overflow (finding F10, 2026-09 review) — a
|
||||
// simple post-multiply check (divide back and compare), not math/bits:
|
||||
// this plugin runs interpreted under Yaegi (COMMON.md's own "stdlib only,
|
||||
// avoid reflection-heavy code" guidance), and a division-based overflow
|
||||
// check needs nothing beyond the four basic int64 operators.
|
||||
func saturatingMulInt64(a, b int64) int64 {
|
||||
if a == 0 || b == 0 {
|
||||
return 0
|
||||
}
|
||||
p := a * b
|
||||
if p/b != a {
|
||||
if (a > 0) == (b > 0) {
|
||||
return math.MaxInt64
|
||||
}
|
||||
return math.MinInt64
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// saturatingAddInt64 returns a+b, clamped to math.MinInt64/math.MaxInt64
|
||||
// instead of wrapping on overflow (finding F10, 2026-09 review) — the
|
||||
// standard "the sum moved the wrong way" overflow test for two's-
|
||||
// complement addition.
|
||||
func saturatingAddInt64(a, b int64) int64 {
|
||||
sum := a + b
|
||||
if b > 0 && sum < a {
|
||||
return math.MaxInt64
|
||||
}
|
||||
if b < 0 && sum > a {
|
||||
return math.MinInt64
|
||||
}
|
||||
return sum
|
||||
}
|
||||
|
||||
// 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.
|
||||
// zero by leaving a map entry nil), then builtinPricing, then
|
||||
// builtinModelMetaTable (pricing_data_gen.go's 236-entry LiteLLM-synced
|
||||
// table — review finding F2, 2026-09 review): builtinPricing alone only
|
||||
// ever carried 10 flagship models, so an operator serving anything else
|
||||
// the generated table already knows a cost for (gpt-4.1, o4-mini,
|
||||
// deepseek-chat, and 200+ others) billed 0 for it even though GET
|
||||
// /v1/models — which reads the identical table via resolveModelMeta,
|
||||
// modelmeta.go — advertised a real price for the exact same model,
|
||||
// letting AllowUnpricedWithCostBudget's cost guard refuse the request as
|
||||
// "unpriced" in the same breath the catalog called it priced. Converted
|
||||
// from the table's integer micro-USD-per-million-tokens unit to
|
||||
// ModelPricing's own float64-USD-per-million-tokens unit via
|
||||
// microUSDPerMTokToUSD (modelmeta.go) — the identical conversion GET
|
||||
// /v1/models' own pricing view already applies to this same table. A
|
||||
// zero-valued cost field in the table means "no data" (builtinModelMeta's
|
||||
// own doc comment), not "free" — at least one side must be positive
|
||||
// before an entry is trusted, the same rule resolveModelMeta's own cost
|
||||
// switch already applies to it. Precedence is unchanged by this addition:
|
||||
// an operator override always wins, builtinPricing (curated, more likely
|
||||
// to be current) wins over the generated table next, and the generated
|
||||
// table is only ever the last resort.
|
||||
func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPricing, bool) {
|
||||
if p, ok := overrides[model]; ok && p != nil {
|
||||
return *p, true
|
||||
@@ -126,25 +220,76 @@ func lookupPricing(model string, overrides map[string]*ModelPricing) (ModelPrici
|
||||
if p, ok := builtinPricing[model]; ok {
|
||||
return p, true
|
||||
}
|
||||
if m, ok := builtinModelMetaTable[model]; ok && (m.InputCostPerMTokMicroUSD > 0 || m.OutputCostPerMTokMicroUSD > 0) {
|
||||
return ModelPricing{
|
||||
InputPerM: microUSDPerMTokToUSD(m.InputCostPerMTokMicroUSD),
|
||||
OutputPerM: microUSDPerMTokToUSD(m.OutputCostPerMTokMicroUSD),
|
||||
}, true
|
||||
}
|
||||
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
|
||||
// validatePricing validates every entry in raw (Config.Pricing overrides)
|
||||
// and returns it unchanged — a constructor error rather than a silently
|
||||
// unenforceable cost budget (finding F4, 2026-09 review): costMicrosKnownFor
|
||||
// had no validation of its own, so a negative, NaN, or Inf price still
|
||||
// passed priceKnown/modelPriceKnown's own "is this model priced" check
|
||||
// (modelmeta.go) yet produced a meaningless or negative cost — permanently
|
||||
// defeating any costPerDayUSD/costPerMonthUSD budget configured for that
|
||||
// model (limits.go's account already clamps a negative cost to 0 and logs
|
||||
// it on every request, but by then the budget has already never fired,
|
||||
// silently, since construction). Mirrors LimitsConfig.validate's own
|
||||
// "non-nil, finite, >= 0" rule (limits.go) applied to InputPerM/OutputPerM
|
||||
// instead of a count/cost limit. raw is walked in sorted key order,
|
||||
// matching validateModelMeta's own convention (modelmeta.go), so a config
|
||||
// with more than one invalid entry always reports the same one first,
|
||||
// deterministically across repeated runs.
|
||||
func validatePricing(raw map[string]*ModelPricing) (map[string]*ModelPricing, error) {
|
||||
if len(raw) == 0 {
|
||||
return raw, nil
|
||||
}
|
||||
keys := make([]string, 0, len(raw))
|
||||
for k := range raw {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
|
||||
for _, k := range keys {
|
||||
p := raw[k]
|
||||
if p == nil {
|
||||
return nil, fmt.Errorf("llmgateway: pricing: entry %q must not be nil", k)
|
||||
}
|
||||
fields := []struct {
|
||||
name string
|
||||
v float64
|
||||
}{
|
||||
{"inputPerM", p.InputPerM},
|
||||
{"outputPerM", p.OutputPerM},
|
||||
}
|
||||
for _, f := range fields {
|
||||
if math.IsNaN(f.v) || math.IsInf(f.v, 0) {
|
||||
return nil, fmt.Errorf("llmgateway: pricing: entry %q: %s must be a finite number, got %v", k, f.name, f.v)
|
||||
}
|
||||
if f.v < 0 {
|
||||
return nil, fmt.Errorf("llmgateway: pricing: entry %q: %s must not be negative, got %v", k, f.name, f.v)
|
||||
}
|
||||
}
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
// warnUnknownModel calls pricingWarnFn for model, once per process
|
||||
// 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) {
|
||||
// warnUnknownModelFor calls warn for model, once per process lifetime, via
|
||||
// the process-global dedup set (warnedModels) — deliberately not scoped
|
||||
// per-caller, so a model already warned about under one Gateway
|
||||
// instance's warn func is not warned about again under a different
|
||||
// instance's (finding F11 / review-auth F8's own "keep the process-global
|
||||
// dedup cap if needed" allowance; costMicrosKnownFor passes its caller's
|
||||
// own warn func). 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 warnUnknownModelFor(model string, warn func(string)) {
|
||||
warnedModelsMu.Lock()
|
||||
defer warnedModelsMu.Unlock()
|
||||
if warnedModels[model] {
|
||||
@@ -153,10 +298,10 @@ func warnUnknownModel(model string) {
|
||||
if len(warnedModels) >= unknownModelWarnCap {
|
||||
if !warnCapNotified {
|
||||
warnCapNotified = true
|
||||
pricingWarnFn(warnCapMessage)
|
||||
warn(warnCapMessage)
|
||||
}
|
||||
return
|
||||
}
|
||||
warnedModels[model] = true
|
||||
pricingWarnFn(model)
|
||||
warn(model)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user