From be5ad32ad49ad4ad18993add2c4205f115fcb724 Mon Sep 17 00:00:00 2001 From: Aleksander Slominski Date: Wed, 26 Aug 2026 17:46:01 -0400 Subject: [PATCH] Fix: litellm-budget-track reads cost from response headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The plugin's OnResponse read pctx.Headers (the request headers) instead of pctx.ResponseHeaders, so it never observed the x-litellm-response-cost header and recorded $0 in every real deployment (both forward and reverse proxy listeners populate response headers on pctx.ResponseHeaders). No unit test covered the plugin, so the bug went unnoticed. This mirrors how the opa plugin reads response headers in OnResponse. Also fall back to X-Litellm-Response-Cost-Original when the bare header is absent: OpenAI /v1/chat/completions responses carry the effective (post-discount) header, but the Anthropic /v1/messages endpoint used by Claude Code — and newer LiteLLM releases — emit only the pre-discount -original variant. Add the plugin's first unit tests: response-header accumulation, the request-header regression guard, the -original fallback, bare-header precedence, invalid/missing cost handling, 429 budget enforcement, ledger persistence, config validation, stale-day ledger reset, and concurrent-response race safety. Assisted-By: Claude (Anthropic AI) Signed-off-by: Aleksander Slominski --- .../plugins/litellm_budgettrack/plugin.go | 21 +- .../litellm_budgettrack/plugin_test.go | 255 ++++++++++++++++++ authbridge/docs/litellm-budgettrack-plugin.md | 7 +- 3 files changed, 280 insertions(+), 3 deletions(-) create mode 100644 authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go index cfc883f0d..8a020c99f 100644 --- a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go @@ -1,4 +1,4 @@ -// Package litellm_budgettrack provides an inbound pipeline plugin that tracks +// Package litellm_budgettrack provides a pipeline plugin that tracks // per-request cost via the x-litellm-response-cost response header and // enforces a daily spending budget, rejecting requests with HTTP 429 // when the budget is exceeded. @@ -17,6 +17,19 @@ import ( "github.com/rossoctl/cortex/authbridge/authlib/plugins" ) +// Response cost headers emitted by LiteLLM. +// +// responseCostHeader is the effective (post-discount) cost and is present on +// OpenAI-style /v1/chat/completions responses. Newer LiteLLM releases — and the +// Anthropic /v1/messages endpoint that Claude Code uses — do not emit it, only +// the pre-discount "-original" variant, so we fall back to that when the bare +// header is absent. Without the fallback, budget tracking silently records $0 +// for Anthropic-format traffic. +const ( + responseCostHeader = "X-Litellm-Response-Cost" + responseCostOriginalHeader = "X-Litellm-Response-Cost-Original" +) + type budgetTrackConfig struct { SpendFile string `json:"spend_file" required:"true" description:"Path to the JSON spend ledger file."` MaxBudget float64 `json:"max_budget" required:"true" description:"Daily budget in USD."` @@ -80,7 +93,11 @@ func (p *BudgetTrack) OnRequest(_ context.Context, pctx *pipeline.Context) pipel // OnResponse reads x-litellm-response-cost and accumulates the spend. func (p *BudgetTrack) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { - costStr := pctx.Headers.Get("X-Litellm-Response-Cost") + costStr := pctx.ResponseHeaders.Get(responseCostHeader) + if costStr == "" { + // Anthropic /v1/messages (and newer LiteLLM) omit the bare header. + costStr = pctx.ResponseHeaders.Get(responseCostOriginalHeader) + } if costStr == "" { return pipeline.Action{Type: pipeline.Continue} } diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go new file mode 100644 index 000000000..ee1b1d22c --- /dev/null +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go @@ -0,0 +1,255 @@ +package litellm_budgettrack + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// configure builds a BudgetTrack with a temp-dir spend file and the given budget. +func configure(t *testing.T, maxBudget float64) *BudgetTrack { + t.Helper() + p := New() + cfg := budgetTrackConfig{ + SpendFile: filepath.Join(t.TempDir(), "spend.json"), + MaxBudget: maxBudget, + } + raw, _ := json.Marshal(cfg) + if err := p.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + return p +} + +// TestOnResponseReadsResponseHeader is the regression guard for the core fix: +// the cost must be read from ResponseHeaders, not the request Headers. +func TestOnResponseReadsResponseHeader(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ + ResponseHeaders: http.Header{responseCostHeader: {"0.0025"}}, + } + + if action := p.OnResponse(context.Background(), pctx); action.Type != pipeline.Continue { + t.Fatalf("OnResponse() = %v, want Continue", action.Type) + } + if p.ledger.TotalSpend != 0.0025 { + t.Errorf("TotalSpend = %v, want 0.0025", p.ledger.TotalSpend) + } + if p.ledger.TotalCalls != 1 { + t.Errorf("TotalCalls = %d, want 1", p.ledger.TotalCalls) + } +} + +// TestOnResponseIgnoresRequestHeader guards against the original bug: the cost +// header on the request side (pctx.Headers) must NOT be accumulated. +func TestOnResponseIgnoresRequestHeader(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ + Headers: http.Header{responseCostHeader: {"0.0025"}}, // wrong place; must be ignored + ResponseHeaders: http.Header{}, + } + + p.OnResponse(context.Background(), pctx) + if p.ledger.TotalSpend != 0 { + t.Errorf("TotalSpend = %v, want 0 (request-header cost must be ignored)", p.ledger.TotalSpend) + } +} + +// TestOnResponseFallsBackToOriginal covers the Anthropic /v1/messages case where +// only the pre-discount "-original" header is present. +func TestOnResponseFallsBackToOriginal(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ + ResponseHeaders: http.Header{responseCostOriginalHeader: {"2.204e-05"}}, + } + + p.OnResponse(context.Background(), pctx) + if p.ledger.TotalSpend != 2.204e-05 { + t.Errorf("TotalSpend = %v, want 2.204e-05 (fallback header)", p.ledger.TotalSpend) + } +} + +// TestOnResponseBareHeaderWins verifies the effective (post-discount) header +// takes precedence over "-original" when both are present. +func TestOnResponseBareHeaderWins(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ + ResponseHeaders: http.Header{ + responseCostHeader: {"0.001"}, + responseCostOriginalHeader: {"0.002"}, + }, + } + + p.OnResponse(context.Background(), pctx) + if p.ledger.TotalSpend != 0.001 { + t.Errorf("TotalSpend = %v, want 0.001 (bare header must win)", p.ledger.TotalSpend) + } +} + +// TestOnResponseIgnoresMissingOrInvalid verifies absent / non-positive / unparseable +// costs are skipped rather than corrupting the ledger. +func TestOnResponseIgnoresMissingOrInvalid(t *testing.T) { + for _, tc := range []struct { + name string + headers http.Header + }{ + {"missing", http.Header{}}, + {"zero", http.Header{responseCostHeader: {"0"}}}, + {"negative", http.Header{responseCostHeader: {"-1"}}}, + {"unparseable", http.Header{responseCostHeader: {"abc"}}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ResponseHeaders: tc.headers} + if action := p.OnResponse(context.Background(), pctx); action.Type != pipeline.Continue { + t.Fatalf("OnResponse() = %v, want Continue", action.Type) + } + if p.ledger.TotalSpend != 0 || p.ledger.TotalCalls != 0 { + t.Errorf("ledger mutated: spend=%v calls=%d, want 0/0", p.ledger.TotalSpend, p.ledger.TotalCalls) + } + }) + } +} + +// TestOnRequestEnforcesBudget verifies OnRequest denies with 429 once the +// accumulated spend reaches the daily budget, and allows before that. +func TestOnRequestEnforcesBudget(t *testing.T) { + p := configure(t, 0.001) + + // Under budget: allowed. + if action := p.OnRequest(context.Background(), &pipeline.Context{}); action.Type != pipeline.Continue { + t.Fatalf("OnRequest() under budget = %v, want Continue", action.Type) + } + + // Accumulate past the budget via a response. + p.OnResponse(context.Background(), &pipeline.Context{ + ResponseHeaders: http.Header{responseCostHeader: {"0.002"}}, + }) + + // Over budget: rejected with 429 / budget.exceeded. + action := p.OnRequest(context.Background(), &pipeline.Context{}) + if action.Type != pipeline.Reject { + t.Fatalf("OnRequest() over budget = %v, want Reject", action.Type) + } + if action.Violation == nil || action.Violation.Status != http.StatusTooManyRequests { + t.Errorf("Violation = %+v, want Status 429", action.Violation) + } + if action.Violation.Code != "budget.exceeded" { + t.Errorf("Violation.Code = %q, want budget.exceeded", action.Violation.Code) + } +} + +// TestLedgerPersistsAcrossInstances verifies the spend file is reloaded, so a +// restart on the same day resumes the accumulated total. +func TestLedgerPersistsAcrossInstances(t *testing.T) { + spendFile := filepath.Join(t.TempDir(), "spend.json") + raw, _ := json.Marshal(budgetTrackConfig{SpendFile: spendFile, MaxBudget: 5.00}) + + p1 := New() + if err := p1.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + p1.OnResponse(context.Background(), &pipeline.Context{ + ResponseHeaders: http.Header{responseCostHeader: {"0.01"}}, + }) + + p2 := New() + if err := p2.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + if p2.ledger.TotalSpend != 0.01 { + t.Errorf("reloaded TotalSpend = %v, want 0.01", p2.ledger.TotalSpend) + } +} + +// TestConfigureRejectsBadConfig verifies required-field and JSON validation. +func TestConfigureRejectsBadConfig(t *testing.T) { + spend := filepath.Join(t.TempDir(), "spend.json") + for _, tc := range []struct { + name string + raw string + }{ + {"empty spend_file", `{"max_budget": 5.0}`}, + {"zero max_budget", fmt.Sprintf(`{"spend_file": %q, "max_budget": 0}`, spend)}, + {"negative max_budget", fmt.Sprintf(`{"spend_file": %q, "max_budget": -1}`, spend)}, + {"invalid json", `{`}, + } { + t.Run(tc.name, func(t *testing.T) { + if err := New().Configure(json.RawMessage(tc.raw)); err == nil { + t.Errorf("Configure(%s) = nil, want error", tc.raw) + } + }) + } +} + +// TestLoadLedgerResetsStaleDay verifies a spend file left over from a previous +// day is discarded on Configure rather than counted against today's budget. +func TestLoadLedgerResetsStaleDay(t *testing.T) { + spend := filepath.Join(t.TempDir(), "spend.json") + stale := `{"date":"2000-01-01","total_spend":9.99,"total_calls":42}` + if err := os.WriteFile(spend, []byte(stale), 0o644); err != nil { + t.Fatalf("seed spend file: %v", err) + } + + p := New() + raw, _ := json.Marshal(budgetTrackConfig{SpendFile: spend, MaxBudget: 5.00}) + if err := p.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + + today := time.Now().UTC().Format("2006-01-02") + if p.ledger.Date != today { + t.Errorf("ledger.Date = %q, want %q", p.ledger.Date, today) + } + if p.ledger.TotalSpend != 0 || p.ledger.TotalCalls != 0 { + t.Errorf("stale ledger not reset: spend=%v calls=%d", p.ledger.TotalSpend, p.ledger.TotalCalls) + } + + // A same-day ledger, by contrast, is preserved. + sameDay := fmt.Sprintf(`{"date":%q,"total_spend":1.25,"total_calls":3}`, today) + if err := os.WriteFile(spend, []byte(sameDay), 0o644); err != nil { + t.Fatalf("seed same-day file: %v", err) + } + p2 := New() + if err := p2.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + if p2.ledger.TotalSpend != 1.25 || p2.ledger.TotalCalls != 3 { + t.Errorf("same-day ledger not preserved: spend=%v calls=%d", p2.ledger.TotalSpend, p2.ledger.TotalCalls) + } +} + +// TestConcurrentOnResponse exercises the mutex under concurrent responses. +// Run with -race to catch data races on the ledger. +func TestConcurrentOnResponse(t *testing.T) { + p := configure(t, 1000.0) // high budget so nothing is rejected + const goroutines = 50 + + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func() { + defer wg.Done() + p.OnResponse(context.Background(), &pipeline.Context{ + ResponseHeaders: http.Header{responseCostHeader: {"0.01"}}, + }) + }() + } + wg.Wait() + + if p.ledger.TotalCalls != goroutines { + t.Errorf("TotalCalls = %d, want %d", p.ledger.TotalCalls, goroutines) + } + // 50 × 0.01 = 0.50, within float tolerance. + if got := p.ledger.TotalSpend; got < 0.4999 || got > 0.5001 { + t.Errorf("TotalSpend = %v, want ~0.50", got) + } +} diff --git a/authbridge/docs/litellm-budgettrack-plugin.md b/authbridge/docs/litellm-budgettrack-plugin.md index 6e1e21ff9..da47bf177 100644 --- a/authbridge/docs/litellm-budgettrack-plugin.md +++ b/authbridge/docs/litellm-budgettrack-plugin.md @@ -90,7 +90,12 @@ The spend file (`spend-authbridge.json`) is a simple JSON object: ### OnResponse (cost accumulation) -1. Read `X-Litellm-Response-Cost` header from upstream response +1. Read the cost from the **response** headers (`pctx.ResponseHeaders`): + `X-Litellm-Response-Cost`, falling back to `X-Litellm-Response-Cost-Original` + when the bare header is absent. The bare (effective, post-discount) header is + present on OpenAI `/v1/chat/completions` responses; the Anthropic `/v1/messages` + endpoint used by Claude Code — and newer LiteLLM releases — emit only the + pre-discount `-original` variant. 2. If missing or non-positive → continue (no cost to track) 3. Lock mutex, reset if new day 4. Add cost to `total_spend`, increment `total_calls`