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:
+171
-10
@@ -1,11 +1,25 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strconv"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// costMicrosOnly is costMicrosKnown for a test that only needs the price
|
||||
// itself, mirroring the removed package-level costMicros helper (dead
|
||||
// code removal, review finding, 2026-09 review): costMicros had no
|
||||
// production caller left once every real caller moved to
|
||||
// costMicrosKnown/unifiedCostMicrosKnown and their own -ok-returning
|
||||
// variants, but this file's existing table tests only ever asserted on
|
||||
// the price, never on the ok bool, so they keep that exact shape here
|
||||
// instead of threading a discarded ok through every case.
|
||||
func costMicrosOnly(model string, u usage, overrides map[string]*ModelPricing) int64 {
|
||||
got, _ := costMicrosKnown(model, u, overrides)
|
||||
return got
|
||||
}
|
||||
|
||||
// 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
|
||||
@@ -27,8 +41,8 @@ func TestCostMicrosBuiltin(t *testing.T) {
|
||||
}
|
||||
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)
|
||||
if got := costMicrosOnly(c.model, c.u, nil); got != c.want {
|
||||
t.Errorf("costMicrosOnly(%q, %+v, nil) = %d, want %d", c.model, c.u, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -44,10 +58,10 @@ func TestCostMicrosOverride(t *testing.T) {
|
||||
}
|
||||
u := usage{prompt: 1_000_000, completion: 1_000_000}
|
||||
|
||||
if got, want := costMicros("gpt-5", u, overrides), int64(2_000_000); got != want {
|
||||
if got, want := costMicrosOnly("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 {
|
||||
if got, want := costMicrosOnly("in-house-model", u, overrides), int64(3_000_000); got != want {
|
||||
t.Errorf("override for unknown model: costMicros = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
@@ -58,11 +72,31 @@ func TestCostMicrosOverride(t *testing.T) {
|
||||
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 {
|
||||
if got, want := costMicrosOnly("gpt-5", u, overrides), int64(11_250_000); got != want {
|
||||
t.Errorf("nil override entry: costMicros = %d, want %d (builtin price)", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// pricingWarnFn is the test-only warn sink costMicrosKnown passes to
|
||||
// costMicrosKnownFor; production callers pass their own g.pricingWarn.
|
||||
// Guarded by warnedModelsMu so tests stay race-clean.
|
||||
var pricingWarnFn = func(string) {}
|
||||
|
||||
func setPricingWarnFn(fn func(string)) {
|
||||
warnedModelsMu.Lock()
|
||||
defer warnedModelsMu.Unlock()
|
||||
pricingWarnFn = fn
|
||||
}
|
||||
|
||||
// costMicrosKnown is a test convenience wrapper around costMicrosKnownFor
|
||||
// using the test-controlled pricingWarnFn.
|
||||
func costMicrosKnown(model string, u usage, overrides map[string]*ModelPricing) (int64, bool) {
|
||||
warnedModelsMu.Lock()
|
||||
warn := pricingWarnFn
|
||||
warnedModelsMu.Unlock()
|
||||
return costMicrosKnownFor(model, u, overrides, warn)
|
||||
}
|
||||
|
||||
// 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
|
||||
@@ -101,11 +135,11 @@ func TestCostMicrosUnknownModelWarnsOnce(t *testing.T) {
|
||||
warned := resetPricingWarnState(t)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
if got := costMicros("totally-unpriced-model", usage{prompt: 100, completion: 100}, nil); got != 0 {
|
||||
if got := costMicrosOnly("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 {
|
||||
if got := costMicrosOnly("another-unpriced-model", usage{prompt: 1}, nil); got != 0 {
|
||||
t.Errorf("costMicros = %d, want 0", got)
|
||||
}
|
||||
|
||||
@@ -132,8 +166,8 @@ func TestCostMicrosUnknownModelWarnCap(t *testing.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)
|
||||
if got := costMicrosOnly(model, usage{prompt: 1}, nil); got != 0 {
|
||||
t.Fatalf("costMicrosOnly(%q, ...) = %d, want 0", model, got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -162,7 +196,7 @@ func TestCostMicrosUnknownModelWarnCap(t *testing.T) {
|
||||
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 {
|
||||
if got, want := costMicrosOnly("round-test", u, overrides), int64(16_400_000); got != want {
|
||||
t.Errorf("costMicros with an 8.2 price = %d, want %d", got, want)
|
||||
}
|
||||
}
|
||||
@@ -190,3 +224,130 @@ func TestBuiltinPricingSeedComplete(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestLookupPricing_FallsBackToGeneratedTable (finding F2, 2026-09
|
||||
// review) asserts a model present ONLY in builtinModelMetaTable — not in
|
||||
// the small curated builtinPricing list — still resolves a real price
|
||||
// through lookupPricing/costMicrosKnown, matching the price GET
|
||||
// /v1/models already advertises for the same model via resolveModelMeta.
|
||||
// Before this fix, such a model billed 0 even though the catalog called
|
||||
// it priced. "deepseek-chat" is one of 200+ such models
|
||||
// (pricing_data_gen.go).
|
||||
func TestLookupPricing_FallsBackToGeneratedTable(t *testing.T) {
|
||||
if _, ok := builtinPricing["deepseek-chat"]; ok {
|
||||
t.Fatal("test assumption broken: deepseek-chat must not be in the curated builtinPricing table")
|
||||
}
|
||||
meta, ok := builtinModelMetaTable["deepseek-chat"]
|
||||
if !ok || meta.InputCostPerMTokMicroUSD <= 0 || meta.OutputCostPerMTokMicroUSD <= 0 {
|
||||
t.Fatal("test assumption broken: deepseek-chat must be in builtinModelMetaTable with a positive cost on both sides")
|
||||
}
|
||||
|
||||
price, ok := lookupPricing("deepseek-chat", nil)
|
||||
if !ok {
|
||||
t.Fatal(`lookupPricing("deepseek-chat", nil) ok = false, want a price resolved from builtinModelMetaTable`)
|
||||
}
|
||||
wantIn := microUSDPerMTokToUSD(meta.InputCostPerMTokMicroUSD)
|
||||
wantOut := microUSDPerMTokToUSD(meta.OutputCostPerMTokMicroUSD)
|
||||
if price.InputPerM != wantIn || price.OutputPerM != wantOut {
|
||||
t.Errorf(`lookupPricing("deepseek-chat", nil) = %+v, want {InputPerM:%v OutputPerM:%v}`, price, wantIn, wantOut)
|
||||
}
|
||||
|
||||
got, ok := costMicrosKnown("deepseek-chat", usage{prompt: 1_000_000, completion: 1_000_000}, nil)
|
||||
if !ok {
|
||||
t.Fatal(`costMicrosKnown("deepseek-chat", ...) ok = false, want true`)
|
||||
}
|
||||
want := meta.InputCostPerMTokMicroUSD + meta.OutputCostPerMTokMicroUSD // 1M tokens each side == exactly the per-M micro-USD price
|
||||
if got != want {
|
||||
t.Errorf(`costMicrosKnown("deepseek-chat", 1M prompt + 1M completion) = %d, want %d`, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLookupPricing_OverrideWinsOverGeneratedTable (finding F2) pins that
|
||||
// an operator override still takes priority over the generated-table
|
||||
// fallback added by this fix, for a model that has no builtinPricing
|
||||
// entry at all — the fallback must never shadow a configured override.
|
||||
func TestLookupPricing_OverrideWinsOverGeneratedTable(t *testing.T) {
|
||||
overrides := map[string]*ModelPricing{"deepseek-chat": {InputPerM: 9, OutputPerM: 9}}
|
||||
price, ok := lookupPricing("deepseek-chat", overrides)
|
||||
if !ok || price.InputPerM != 9 || price.OutputPerM != 9 {
|
||||
t.Errorf(`lookupPricing("deepseek-chat", override) = %+v, %v, want the override {9 9}`, price, ok)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCostMicrosKnown_RoundsCombinedTotal (finding F9, 2026-09 review)
|
||||
// asserts the input and output totals are summed BEFORE dividing back
|
||||
// down to micro-USD, rather than each direction truncated separately: at
|
||||
// $0.02/M, 30 prompt tokens alone truncate to 0 (30*20_000/1_000_000
|
||||
// floors to 0) and so do 30 completion tokens alone, but their COMBINED
|
||||
// total (1,200,000 micro-USD-tokens) divides to a nonzero 1 micro-USD —
|
||||
// the exact "small cheap request silently bills 0" gap the finding
|
||||
// reports, which also meant account (limits.go) skipped the cost-counter
|
||||
// write entirely for that traffic.
|
||||
func TestCostMicrosKnown_RoundsCombinedTotal(t *testing.T) {
|
||||
overrides := map[string]*ModelPricing{"cheap-model": {InputPerM: 0.02, OutputPerM: 0.02}}
|
||||
got, ok := costMicrosKnown("cheap-model", usage{prompt: 30, completion: 30}, overrides)
|
||||
if !ok {
|
||||
t.Fatal("costMicrosKnown ok = false, want true")
|
||||
}
|
||||
if got != 1 {
|
||||
t.Errorf("costMicrosKnown(prompt=30, completion=30 @ $0.02/M) = %d, want 1 (summed-then-divided, not truncated-to-0 per direction)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCostMicrosKnown_OverflowSaturates (finding F10, 2026-09 review)
|
||||
// asserts an absurd upstream-reported token count saturates at
|
||||
// math.MaxInt64 instead of silently wrapping to a negative or garbage
|
||||
// value through an unchecked int64 multiply — upstream-controlled per
|
||||
// limits.go's account (security audit run-1), since the token counts come
|
||||
// from the provider's own reported usage.
|
||||
func TestCostMicrosKnown_OverflowSaturates(t *testing.T) {
|
||||
overrides := map[string]*ModelPricing{"expensive-model": {InputPerM: 75, OutputPerM: 75}}
|
||||
got, ok := costMicrosKnown("expensive-model", usage{prompt: math.MaxInt64 / 2, completion: math.MaxInt64 / 2}, overrides)
|
||||
if !ok {
|
||||
t.Fatal("costMicrosKnown ok = false, want true")
|
||||
}
|
||||
// Both inputCost and outputCost individually overflow the int64
|
||||
// multiply and saturate at math.MaxInt64 (saturatingMulInt64), and
|
||||
// their sum saturates there too (saturatingAddInt64) rather than
|
||||
// wrapping to a negative/garbage value — but the value costMicrosKnown
|
||||
// actually RETURNS is that saturated total divided back down to
|
||||
// micro-USD (roundedDivInt64, computeCostMicros' own final step), not
|
||||
// raw math.MaxInt64 itself.
|
||||
want := roundedDivInt64(math.MaxInt64, tokensPerMillion)
|
||||
if got != want {
|
||||
t.Errorf("costMicrosKnown with an absurd upstream-reported token count = %d, want %d (the saturated total divided back to micro-USD)", got, want)
|
||||
}
|
||||
if got <= 0 {
|
||||
t.Errorf("costMicrosKnown with an absurd upstream-reported token count = %d, want a large positive value, not a wrapped negative/garbage one", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidatePricing_RejectsBadEntries (finding F4, 2026-09 review)
|
||||
// table-tests validatePricing's construction-time checks: nil, NaN, Inf,
|
||||
// and negative entries must all fail with a config error; a valid entry
|
||||
// must pass unchanged.
|
||||
func TestValidatePricing_RejectsBadEntries(t *testing.T) {
|
||||
cases := []struct {
|
||||
pricing map[string]*ModelPricing
|
||||
name string
|
||||
wantErr bool
|
||||
}{
|
||||
{nil, "nil map", false},
|
||||
{map[string]*ModelPricing{}, "empty map", false},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: 2}}, "valid entry", false},
|
||||
{map[string]*ModelPricing{"m": nil}, "nil entry", true},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: -1, OutputPerM: 1}}, "negative input", true},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: -1}}, "negative output", true},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: math.NaN(), OutputPerM: 1}}, "NaN input", true},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: 1, OutputPerM: math.Inf(1)}}, "+Inf output", true},
|
||||
{map[string]*ModelPricing{"m": {InputPerM: math.Inf(-1), OutputPerM: 1}}, "-Inf input", true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
_, err := validatePricing(c.pricing)
|
||||
if (err != nil) != c.wantErr {
|
||||
t.Errorf("validatePricing(%+v) err = %v, wantErr %v", c.pricing, err, c.wantErr)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user