diff --git a/authbridge/authlib/plugins/litellm_budgettrack/forwardproxy_integration_test.go b/authbridge/authlib/plugins/litellm_budgettrack/forwardproxy_integration_test.go new file mode 100644 index 00000000..0ba9e1ab --- /dev/null +++ b/authbridge/authlib/plugins/litellm_budgettrack/forwardproxy_integration_test.go @@ -0,0 +1,105 @@ +package litellm_budgettrack + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/url" + "path/filepath" + "testing" + "time" + + fwd "github.com/rossoctl/cortex/authbridge/authlib/listener/forwardproxy" + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" + "github.com/rossoctl/cortex/authbridge/authlib/session" +) + +// TestForwardProxyStreamedSSEUpdatesLedger is the listener-level test the PR #815 +// review asked for: stand up the real forward proxy with a streamed +// (text/event-stream) upstream and BudgetTrack in the outbound pipeline, drive a +// request through the proxy, and assert the ledger moved. +// +// This exercises the path the direct-call unit tests structurally cannot: whether +// the listener actually dispatches response frames to the plugin. On the +// header-only version (before this branch), a streamed response never reaches the +// plugin's cost accounting, so this test would fail — which is exactly the gap the +// reviewer flagged. +func TestForwardProxyStreamedSSEUpdatesLedger(t *testing.T) { + // Upstream emits Anthropic-style streamed usage and NO cost header — as + // LiteLLM does for streamed responses — so the plugin must price it from the + // parsed token usage. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + f, _ := w.(http.Flusher) + io.WriteString(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":100,\"output_tokens\":1}}}\n\n") + if f != nil { + f.Flush() + } + io.WriteString(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":40}}\n\n") + if f != nil { + f.Flush() + } + })) + t.Cleanup(upstream.Close) + + spend := filepath.Join(t.TempDir(), "spend.json") + p := New() + raw, _ := json.Marshal(budgetTrackConfig{ + SpendFile: spend, MaxBudget: 5, InputCostPerToken: 1e-6, OutputCostPerToken: 5e-6, + }) + if err := p.Configure(raw); err != nil { + t.Fatalf("Configure: %v", err) + } + // WrapConfigured is what the real build applies; it preserves StreamingResponder. + wrapped := pipeline.WrapConfigured(p, raw) + + pipe, err := pipeline.New([]pipeline.Plugin{wrapped}) + if err != nil { + t.Fatalf("pipeline.New: %v", err) + } + if !pipe.HasStreamingResponders() { + t.Fatal("pipeline does not recognize BudgetTrack as a StreamingResponder") + } + store := session.New(5*time.Minute, 100, 0) + t.Cleanup(store.Close) + srv, err := fwd.NewServer(pipeline.NewHolder(pipe), store, nil) + if err != nil { + t.Fatalf("NewServer: %v", err) + } + proxy := httptest.NewServer(srv.Handler()) + t.Cleanup(proxy.Close) + + pu, _ := url.Parse(proxy.URL) + client := &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(pu)}} + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, upstream.URL+"/v1/messages", nil) + if err != nil { + t.Fatalf("new request: %v", err) + } + resp, err := client.Do(req) + if err != nil { + t.Fatalf("request via proxy: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + if err := resp.Body.Close(); err != nil { + t.Errorf("close body: %v", err) + } + + // The terminal last=true dispatch runs in the handler's defer after the + // stream is forwarded, so poll briefly for the ledger to settle. + want := 100*1e-6 + 40*5e-6 // 0.0003 + for i := 0; i < 200; i++ { + p.mu.Lock() + got, calls := p.ledger.TotalSpend, p.ledger.TotalCalls + p.mu.Unlock() + if calls > 0 { + if got < want-1e-12 || got > want+1e-12 { + t.Fatalf("ledger TotalSpend = %v, want %v", got, want) + } + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("ledger never updated from a streamed SSE response — the forward proxy did not dispatch frames to the plugin") +} diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go index 8a020c99..c7b57a46 100644 --- a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go @@ -1,15 +1,30 @@ // 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. +// per-request cost and enforces a daily spending budget, rejecting requests +// with HTTP 429 when the budget is exceeded. +// +// Cost is resolved in two ways: +// +// - Non-streaming responses carry the cost in a response header +// (x-litellm-response-cost, or the pre-discount -original variant), read +// on the terminal frame. +// - Streaming responses (text/event-stream — what Claude Code's +// /v1/messages uses) report cost 0 in the header because the total is not +// known when the headers are sent. For these, the plugin parses the token +// usage out of the terminal SSE events (Anthropic message_delta / +// message_stop, or OpenAI's final chunk usage) and prices it from the +// configured per-token rates. Streaming cost tracking is therefore active +// only when input_cost_per_token / output_cost_per_token are configured. package litellm_budgettrack import ( + "bytes" "context" "encoding/json" "fmt" + "math" "os" "strconv" + "strings" "sync" "time" @@ -33,6 +48,26 @@ const ( 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."` + // InputCostPerToken / OutputCostPerToken price streamed responses whose + // header cost is 0 (the total is unknown when streaming headers are sent). + // USD per token; optional. When both are zero, streamed responses cannot be + // priced and contribute 0 to the ledger. + InputCostPerToken float64 `json:"input_cost_per_token" description:"USD per input/prompt token, for pricing streamed responses."` + OutputCostPerToken float64 `json:"output_cost_per_token" description:"USD per output/completion token, for pricing streamed responses."` +} + +// stateKey names the per-request scratch holding token usage accumulated across +// streaming frames until the terminal frame prices it. +const stateKey = "litellm-budget-track" + +// usageState accumulates the largest token counts seen across a stream's +// frames. Anthropic reports input_tokens in message_start and the cumulative +// output_tokens in the final message_delta, so taking the max of each yields +// the final totals; OpenAI reports both together in its terminal usage chunk. +type usageState struct { + inputTokens int + outputTokens int + settled bool // terminal frame already priced this request (exactly-once) } type spendLedger struct { @@ -59,7 +94,15 @@ func (p *BudgetTrack) Name() string { return "litellm-budget-track" } func (p *BudgetTrack) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{ - Description: "Track x-litellm-response-cost and enforce daily budget limit.", + // ReadsBody: the plugin parses the response body (streamed usage). It + // makes Pipeline.NeedsBody() true so the extproc (envoy-sidecar) listener + // buffers the response body and takes its body-phase branch; without it + // that listener dispatches a single header-only RunResponseFrame and the + // streamed accounting silently records nothing (or double-charges if + // Envoy is statically configured BUFFERED). The proxy listeners gate on + // HasStreamingResponders() and are unaffected. Mirrors inference-parser. + ReadsBody: true, + Description: "Track LLM cost (response header or streamed usage) and enforce a daily budget.", } } @@ -73,6 +116,17 @@ func (p *BudgetTrack) Configure(raw json.RawMessage) error { if p.cfg.MaxBudget <= 0 { return fmt.Errorf("litellm-budget-track: max_budget must be > 0") } + // Per-token rates must be finite and non-negative. A negative rate would make + // a streamed request's cost negative, which accumulate() drops — so the request + // would silently neither charge budget nor record a call. Reject at config time. + for name, rate := range map[string]float64{ + "input_cost_per_token": p.cfg.InputCostPerToken, + "output_cost_per_token": p.cfg.OutputCostPerToken, + } { + if rate < 0 || math.IsNaN(rate) || math.IsInf(rate, 0) { + return fmt.Errorf("litellm-budget-track: %s must be finite and >= 0", name) + } + } p.loadLedger() return nil } @@ -91,31 +145,188 @@ func (p *BudgetTrack) OnRequest(_ context.Context, pctx *pipeline.Context) pipel return pipeline.Action{Type: pipeline.Continue} } -// OnResponse reads x-litellm-response-cost and accumulates the spend. +// OnResponse handles the buffered path on listeners that do not route through +// OnResponseFrame. On the proxy listeners this plugin is a StreamingResponder, +// so pipeline.RunResponse skips it and OnResponseFrame drives accumulation +// instead; this remains for listeners that only call OnResponse. func (p *BudgetTrack) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { - costStr := pctx.ResponseHeaders.Get(responseCostHeader) - if costStr == "" { - // Anthropic /v1/messages (and newer LiteLLM) omit the bare header. - costStr = pctx.ResponseHeaders.Get(responseCostOriginalHeader) + if cost, _ := headerCost(pctx); cost > 0 { + p.accumulate(cost) } - if costStr == "" { + return pipeline.Action{Type: pipeline.Continue} +} + +// OnResponseFrame observes each response frame. It parses token usage out of +// streamed SSE frames and, on the terminal frame, prices the request: the +// response-header cost when present (non-streaming), otherwise the parsed +// usage times the configured per-token rates (streaming). +func (p *BudgetTrack) OnResponseFrame(_ context.Context, pctx *pipeline.Context, frame []byte, last bool) pipeline.Action { + if in, out, ok := parseFrameUsage(frame); ok { + st := pipeline.GetState[usageState](pctx, stateKey) + if st == nil { + st = &usageState{} + pipeline.SetState(pctx, stateKey, st) + } + if in > st.inputTokens { + st.inputTokens = in + } + if out > st.outputTokens { + st.outputTokens = out + } + } + if !last { return pipeline.Action{Type: pipeline.Continue} } - cost, err := strconv.ParseFloat(costStr, 64) - if err != nil || cost <= 0 { + + // Terminal frame: settle the cost exactly once. Materialize the scratch + // unconditionally (a header-only response never allocated it above) so the + // guard also covers that path — a listener that dispatches last=true twice + // (e.g. extproc header + buffered-body phases) must not double-charge. + st := pipeline.GetState[usageState](pctx, stateKey) + if st == nil { + st = &usageState{} + pipeline.SetState(pctx, stateKey, st) + } + if st.settled { return pipeline.Action{Type: pipeline.Continue} } + st.settled = true + cost, present := headerCost(pctx) + if cost <= 0 { + // Fall back to per-token pricing only when there is no authoritative + // header cost: the header is absent, or this is a streamed response + // (where LiteLLM always reports 0). A present "0" on a non-streamed + // response is a genuine free call (cache hit / error) — charge nothing, + // don't invent a cost from the usage block. + if !present || isEventStream(pctx) { + cost = float64(st.inputTokens)*p.cfg.InputCostPerToken + + float64(st.outputTokens)*p.cfg.OutputCostPerToken + } + } + if cost > 0 { + p.accumulate(cost) + } + return pipeline.Action{Type: pipeline.Continue} +} + +// accumulate adds one priced call to today's ledger and persists it. A +// non-finite or non-positive cost is ignored: NaN/±Inf would poison +// TotalSpend (making the budget check meaningless) and break the JSON +// marshal, so this is the single chokepoint that guarantees the ledger +// only ever holds finite money. +func (p *BudgetTrack) accumulate(cost float64) { + if cost <= 0 || math.IsNaN(cost) || math.IsInf(cost, 0) { + return + } p.mu.Lock() p.resetIfNewDay() p.ledger.TotalSpend += cost p.ledger.TotalCalls++ p.saveLedger() p.mu.Unlock() +} - return pipeline.Action{Type: pipeline.Continue} +// headerCost returns the usable positive cost reported in the response headers +// and whether a cost header was present at all. present distinguishes "no +// header" (fall back to usage pricing) from "header says 0" (a genuine free +// call — cache hit / error — that must NOT be re-priced from usage). A present +// but non-positive/non-finite header yields (0, true). +func headerCost(pctx *pipeline.Context) (cost float64, present bool) { + 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 0, false + } + c, err := strconv.ParseFloat(costStr, 64) + // strconv.ParseFloat accepts "NaN" / "Inf"; reject non-finite (and + // non-positive) so a garbage or zero header does not poison the ledger. The + // header was still present, so report that. + if err != nil || c <= 0 || math.IsNaN(c) || math.IsInf(c, 0) { + return 0, true + } + return c, true +} + +// isEventStream reports whether the response is a text/event-stream (SSE) — the +// streamed shape where LiteLLM reports cost 0 in the header, so usage-based +// pricing is the intended fallback. +func isEventStream(pctx *pipeline.Context) bool { + ct := pctx.ResponseHeaders.Get("Content-Type") + if i := strings.IndexByte(ct, ';'); i >= 0 { + ct = ct[:i] + } + return strings.EqualFold(strings.TrimSpace(ct), "text/event-stream") } +// parseFrameUsage extracts token usage from a response frame, covering +// Anthropic (usage, or message.usage in message_start) and OpenAI +// (usage.prompt_tokens / completion_tokens). Returns the largest input/output +// token counts found. +// +// The listener's sseframe reader strips the "data:" prefix and returns the +// bare payload, so a streamed frame arrives as raw JSON. The buffered +// application/json path also delivers the whole body as one raw-JSON frame. +// We therefore try the frame as JSON directly, and also scan any "data:" +// lines for the case a frame still carries SSE framing. +func parseFrameUsage(frame []byte) (in, out int, found bool) { + consider := func(b []byte) { + b = bytes.TrimSpace(b) + if len(b) == 0 || b[0] != '{' { + return + } + var ev struct { + Usage *usageJSON `json:"usage"` + Message *struct { + Usage *usageJSON `json:"usage"` + } `json:"message"` + } + if json.Unmarshal(b, &ev) != nil { + return + } + u := ev.Usage + if u == nil && ev.Message != nil { + u = ev.Message.Usage // Anthropic message_start nests usage + } + if u == nil { + return + } + if i := u.inputTotal(); i > in { + in, found = i, true + } + if o := u.outputTotal(); o > out { + out, found = o, true + } + } + + consider(frame) // bare-JSON frame (sseframe payload, or buffered body) + for _, line := range bytes.Split(frame, []byte("\n")) { + if line = bytes.TrimSpace(line); bytes.HasPrefix(line, []byte("data:")) { + consider(bytes.TrimPrefix(line, []byte("data:"))) + } + } + return in, out, found +} + +// usageJSON accepts both Anthropic and OpenAI usage shapes. +type usageJSON struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` +} + +func (u usageJSON) inputTotal() int { + return u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens + u.PromptTokens +} + +func (u usageJSON) outputTotal() int { return u.OutputTokens + u.CompletionTokens } + func (p *BudgetTrack) todayUTC() string { return time.Now().UTC().Format("2006-01-02") } @@ -142,11 +353,18 @@ func (p *BudgetTrack) loadLedger() { } func (p *BudgetTrack) saveLedger() { - data, _ := json.MarshalIndent(p.ledger, "", " ") + data, err := json.MarshalIndent(p.ledger, "", " ") + if err != nil { + // Never overwrite a good ledger with a failed marshal (e.g. a + // non-finite TotalSpend that slipped through). accumulate already + // rejects non-finite costs; this is the belt-and-suspenders guard. + return + } _ = os.WriteFile(p.cfg.SpendFile, data, 0644) } var ( - _ pipeline.Plugin = (*BudgetTrack)(nil) - _ pipeline.Configurable = (*BudgetTrack)(nil) + _ pipeline.Plugin = (*BudgetTrack)(nil) + _ pipeline.Configurable = (*BudgetTrack)(nil) + _ pipeline.StreamingResponder = (*BudgetTrack)(nil) ) diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go index ee1b1d22..d44d5c66 100644 --- a/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go @@ -139,8 +139,12 @@ func TestOnRequestEnforcesBudget(t *testing.T) { 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) + // Stop before dereferencing: a nil Violation must not panic the next lines. + if action.Violation == nil { + t.Fatal("Violation is nil, want 429 budget.exceeded") + } + if action.Violation.Status != http.StatusTooManyRequests { + t.Errorf("Violation.Status = %d, want 429", action.Violation.Status) } if action.Violation.Code != "budget.exceeded" { t.Errorf("Violation.Code = %q, want budget.exceeded", action.Violation.Code) @@ -180,6 +184,8 @@ func TestConfigureRejectsBadConfig(t *testing.T) { {"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)}, + {"negative input rate", fmt.Sprintf(`{"spend_file": %q, "max_budget": 5, "input_cost_per_token": -0.001}`, spend)}, + {"negative output rate", fmt.Sprintf(`{"spend_file": %q, "max_budget": 5, "output_cost_per_token": -0.001}`, spend)}, {"invalid json", `{`}, } { t.Run(tc.name, func(t *testing.T) { @@ -253,3 +259,209 @@ func TestConcurrentOnResponse(t *testing.T) { t.Errorf("TotalSpend = %v, want ~0.50", got) } } + +// --- streaming (SSE usage) tests --- + +// configurePriced builds a plugin with per-token streaming prices set. +func configurePriced(t *testing.T, maxBudget, inPer, outPer float64) *BudgetTrack { + t.Helper() + p := New() + raw, _ := json.Marshal(budgetTrackConfig{ + SpendFile: filepath.Join(t.TempDir(), "spend.json"), + MaxBudget: maxBudget, + InputCostPerToken: inPer, + OutputCostPerToken: outPer, + }) + if err := p.Configure(raw); err != nil { + t.Fatalf("Configure() error = %v", err) + } + return p +} + +// Anthropic-style streamed /v1/messages frames: input in message_start, +// cumulative output in message_delta, both in message_stop. +const ( + frameMessageStart = "event: message_start\n" + + `data: {"type":"message_start","message":{"usage":{"input_tokens":100,"output_tokens":1}}}` + "\n" + frameContentDelta = "event: content_block_delta\n" + + `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}}` + "\n" + frameMessageDelta = "event: message_delta\n" + + `data: {"type":"message_delta","usage":{"input_tokens":0,"output_tokens":40}}` + "\n" +) + +// TestStreamingPricesFromUsage: header cost is absent/0 (streaming), so cost is +// computed from the parsed usage and the configured per-token rates. +func TestStreamingPricesFromUsage(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) // $1/1M in, $5/1M out + ctx := context.Background() + pctx := &pipeline.Context{ResponseHeaders: http.Header{}} // streamed: no cost header + + p.OnResponseFrame(ctx, pctx, []byte(frameMessageStart), false) + p.OnResponseFrame(ctx, pctx, []byte(frameContentDelta), false) + p.OnResponseFrame(ctx, pctx, []byte(frameMessageDelta), false) + p.OnResponseFrame(ctx, pctx, nil, true) // terminal frame settles cost + + want := 100*1e-6 + 40*5e-6 // 0.0001 + 0.0002 = 0.0003 + if got := p.ledger.TotalSpend; got < want-1e-12 || got > want+1e-12 { + t.Errorf("TotalSpend = %v, want %v", got, want) + } + if p.ledger.TotalCalls != 1 { + t.Errorf("TotalCalls = %d, want 1", p.ledger.TotalCalls) + } +} + +// TestStreamingWithoutPricesRecordsZero: no per-token rates configured -> a +// streamed response cannot be priced and must not corrupt the ledger. +func TestStreamingWithoutPricesRecordsZero(t *testing.T) { + p := configure(t, 5.00) // no prices + ctx := context.Background() + pctx := &pipeline.Context{ResponseHeaders: http.Header{}} + p.OnResponseFrame(ctx, pctx, []byte(frameMessageStart), false) + p.OnResponseFrame(ctx, pctx, []byte(frameMessageDelta), false) + p.OnResponseFrame(ctx, pctx, nil, true) + if p.ledger.TotalSpend != 0 || p.ledger.TotalCalls != 0 { + t.Errorf("ledger mutated without prices: spend=%v calls=%d", p.ledger.TotalSpend, p.ledger.TotalCalls) + } +} + +// TestHeaderCostWinsOverUsage: when the terminal frame has a real header cost +// (non-streaming buffered path delivered as a single frame), it is used and the +// per-token pricing is ignored. +func TestHeaderCostWinsOverUsage(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + pctx := &pipeline.Context{ResponseHeaders: http.Header{responseCostHeader: {"0.02"}}} + // single-frame buffered json also carries usage, which must be ignored + body := []byte(`data: {"usage":{"prompt_tokens":100,"completion_tokens":40}}`) + p.OnResponseFrame(context.Background(), pctx, body, true) + if got := p.ledger.TotalSpend; got != 0.02 { + t.Errorf("TotalSpend = %v, want 0.02 (header cost must win)", got) + } +} + +// TestOnResponseFrameOriginalFallback: streamed header 0 but non-streaming +// -original present on the terminal frame is still honored. +func TestOnResponseFrameOriginalFallback(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + pctx := &pipeline.Context{ResponseHeaders: http.Header{responseCostOriginalHeader: {"6.688e-05"}}} + p.OnResponseFrame(context.Background(), pctx, nil, true) + if got := p.ledger.TotalSpend; got < 6.687e-05 || got > 6.689e-05 { + t.Errorf("TotalSpend = %v, want 6.688e-05", got) + } +} + +// TestParseFrameUsageOpenAI covers the OpenAI terminal usage chunk shape. +func TestParseFrameUsageOpenAI(t *testing.T) { + frame := []byte(`data: {"choices":[],"usage":{"prompt_tokens":30,"completion_tokens":12}}`) + in, out, ok := parseFrameUsage(frame) + if !ok || in != 30 || out != 12 { + t.Errorf("parseFrameUsage = (%d,%d,%v), want (30,12,true)", in, out, ok) + } +} + +// TestStreamingBareFrames reflects reality: the sseframe reader strips the +// "data:" prefix, so OnResponseFrame receives bare JSON payloads. +func TestStreamingBareFrames(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + ctx := context.Background() + pctx := &pipeline.Context{ResponseHeaders: http.Header{}} + // bare-JSON frames (no "data:" prefix), as ReadFrame returns them + p.OnResponseFrame(ctx, pctx, []byte(`{"type":"message_start","message":{"usage":{"input_tokens":100,"output_tokens":1}}}`), false) + p.OnResponseFrame(ctx, pctx, []byte(`{"type":"message_delta","usage":{"output_tokens":40}}`), false) + p.OnResponseFrame(ctx, pctx, nil, true) + want := 100*1e-6 + 40*5e-6 + if got := p.ledger.TotalSpend; got < want-1e-12 || got > want+1e-12 { + t.Errorf("TotalSpend = %v, want %v (bare-JSON frames)", got, want) + } +} + +// TestParseFrameUsageBareJSON unit-checks the bare-payload path directly. +func TestParseFrameUsageBareJSON(t *testing.T) { + in, out, ok := parseFrameUsage([]byte(`{"type":"message_delta","usage":{"input_tokens":14,"output_tokens":8}}`)) + if !ok || in != 14 || out != 8 { + t.Errorf("parseFrameUsage(bare) = (%d,%d,%v), want (14,8,true)", in, out, ok) + } +} + +// TestNonFiniteCostRejected guards the data-integrity fix from PR #815 review: +// strconv.ParseFloat accepts "NaN"/"Inf", both slip past a bare `cost <= 0` +// check, poison TotalSpend, and break the JSON marshal. The ledger must stay +// clean and its file must not be overwritten with garbage. +func TestNonFiniteCostRejected(t *testing.T) { + for _, hdr := range []string{"NaN", "Inf", "+Inf", "-Inf"} { + t.Run(hdr, func(t *testing.T) { + p := configure(t, 5.00) + p.OnResponse(context.Background(), &pipeline.Context{ + ResponseHeaders: http.Header{responseCostHeader: {hdr}}, + }) + if p.ledger.TotalSpend != 0 || p.ledger.TotalCalls != 0 { + t.Errorf("%s: ledger mutated: spend=%v calls=%d", hdr, p.ledger.TotalSpend, p.ledger.TotalCalls) + } + // The spend file must remain valid JSON (not overwritten with garbage). + if data, err := os.ReadFile(p.cfg.SpendFile); err == nil && len(data) > 0 { + var l spendLedger + if json.Unmarshal(data, &l) != nil { + t.Errorf("%s: spend file corrupted: %s", hdr, data) + } + } + }) + } +} + +// TestCapabilitiesDeclaresReadsBody is the must-fix from PR #816 review: the +// plugin parses the response body, so it must declare ReadsBody or the extproc +// listener won't buffer the body (streamed accounting records nothing). +func TestCapabilitiesDeclaresReadsBody(t *testing.T) { + if !New().Capabilities().ReadsBody { + t.Error("Capabilities().ReadsBody = false; extproc will not buffer the body and streamed cost is lost") + } +} + +// TestOnResponseFrameSettlesOnce guards the exactly-once contract: a second +// terminal dispatch (e.g. extproc header + buffered-body phases) must not +// double-charge the ledger. +func TestOnResponseFrameSettlesOnce(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + ctx := context.Background() + pctx := &pipeline.Context{ResponseHeaders: http.Header{}} + p.OnResponseFrame(ctx, pctx, []byte(`{"type":"message_start","message":{"usage":{"input_tokens":100,"output_tokens":1}}}`), false) + p.OnResponseFrame(ctx, pctx, []byte(`{"type":"message_delta","usage":{"output_tokens":40}}`), false) + p.OnResponseFrame(ctx, pctx, nil, true) // first terminal — charges + p.OnResponseFrame(ctx, pctx, nil, true) // second terminal — must be a no-op + want := 100*1e-6 + 40*5e-6 + if p.ledger.TotalCalls != 1 { + t.Errorf("TotalCalls = %d, want 1 (double terminal dispatch must not double-charge)", p.ledger.TotalCalls) + } + if got := p.ledger.TotalSpend; got < want-1e-12 || got > want+1e-12 { + t.Errorf("TotalSpend = %v, want %v", got, want) + } +} + +// TestZeroCostHeaderNonStreamedNotRepriced: a genuine free call (cost header +// "0", non-streamed) must be charged 0, not re-priced from its usage block. +func TestZeroCostHeaderNonStreamedNotRepriced(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + pctx := &pipeline.Context{ResponseHeaders: http.Header{ + responseCostHeader: {"0"}, + "Content-Type": {"application/json"}, + }} + p.OnResponseFrame(context.Background(), pctx, []byte(`{"usage":{"input_tokens":100,"output_tokens":40}}`), true) + if p.ledger.TotalSpend != 0 || p.ledger.TotalCalls != 0 { + t.Errorf("free non-streamed call re-priced from usage: spend=%v calls=%d", p.ledger.TotalSpend, p.ledger.TotalCalls) + } +} + +// TestZeroCostHeaderStreamedPricesFromUsage: streamed responses always report a +// 0 cost header, so the usage fallback must still apply for text/event-stream. +func TestZeroCostHeaderStreamedPricesFromUsage(t *testing.T) { + p := configurePriced(t, 5.00, 1e-6, 5e-6) + pctx := &pipeline.Context{ResponseHeaders: http.Header{ + responseCostHeader: {"0"}, + "Content-Type": {"text/event-stream; charset=utf-8"}, + }} + p.OnResponseFrame(context.Background(), pctx, []byte(`{"type":"message_start","message":{"usage":{"input_tokens":100,"output_tokens":40}}}`), false) + p.OnResponseFrame(context.Background(), pctx, nil, true) + want := 100*1e-6 + 40*5e-6 + if got := p.ledger.TotalSpend; got < want-1e-12 || got > want+1e-12 { + t.Errorf("streamed zero-header call not priced from usage: got %v want %v", got, want) + } +} diff --git a/authbridge/authlib/plugins/litellm_budgettrack/streaming_integration_test.go b/authbridge/authlib/plugins/litellm_budgettrack/streaming_integration_test.go new file mode 100644 index 00000000..c2c1f071 --- /dev/null +++ b/authbridge/authlib/plugins/litellm_budgettrack/streaming_integration_test.go @@ -0,0 +1,30 @@ +package litellm_budgettrack + +import ( + "encoding/json" + "testing" + + "github.com/rossoctl/cortex/authbridge/authlib/pipeline" +) + +// Verifies the plugin is recognized as a StreamingResponder through the same +// wrapping the real build applies — the gap that unit tests calling +// OnResponseFrame directly cannot catch. +func TestPipelineDetectsStreamingResponder(t *testing.T) { + raw, _ := json.Marshal(budgetTrackConfig{SpendFile: t.TempDir() + "/s.json", MaxBudget: 5, InputCostPerToken: 1e-6}) + p := New() + if err := p.Configure(raw); err != nil { + t.Fatal(err) + } + wrapped := pipeline.WrapConfigured(p, raw) + if _, ok := wrapped.(pipeline.StreamingResponder); !ok { + t.Fatal("wrapped plugin is NOT a StreamingResponder") + } + pl, err := pipeline.New([]pipeline.Plugin{wrapped}) + if err != nil { + t.Fatal(err) + } + if !pl.HasStreamingResponders() { + t.Fatal("pipeline.HasStreamingResponders() = false; forward proxy will use the buffered path and never call OnResponseFrame") + } +} diff --git a/authbridge/docs/litellm-budgettrack-plugin.md b/authbridge/docs/litellm-budgettrack-plugin.md index da47bf17..01b9241c 100644 --- a/authbridge/docs/litellm-budgettrack-plugin.md +++ b/authbridge/docs/litellm-budgettrack-plugin.md @@ -30,9 +30,28 @@ Agent → cortex.py → AuthBridge (litellm-budget-track) → LiteLLM upstream └── spend-authbridge.json (daily ledger) ``` -The plugin runs in the **inbound** pipeline direction: -- `OnRequest` — pre-flight budget check (reject if over limit) -- `OnResponse` — post-flight cost accumulation (read header, update ledger) +The plugin hooks: +- `OnRequest` — pre-flight budget check (reject if over limit). +- `OnResponseFrame` — post-flight cost accounting for **every** response. Because the + plugin is a `StreamingResponder`, in-tree listeners route all responses through this + hook — a buffered `application/json` body as a single terminal frame, a streamed + `text/event-stream` body frame-by-frame — and `pipeline.RunResponse` skips the plugin + unconditionally. (`OnResponse` remains only as a fallback for a hypothetical listener + that calls it but never `OnResponseFrame`; no in-tree listener does.) + +### Cost source (what the terminal frame charges) + +The cost is settled **once**, on the terminal frame, from one of two sources: + +- **Response header** — `x-litellm-response-cost`, falling back to the pre-discount + `-original` variant. Used whenever the header carries a usable positive cost. A + header of `0` on a non-streamed response is a genuine free call (cache hit / error) + and is charged `0` — it is **not** re-priced from usage. +- **Parsed token usage × configured rates** — used only when the cost header is + **absent**, or the response is `text/event-stream` (LiteLLM always reports `0` in the + header for streams, e.g. Claude Code's `/v1/messages`). The plugin sums the token + usage from the terminal SSE events and multiplies by `input_cost_per_token` / + `output_cost_per_token`. Without those rates a streamed response contributes `0`. ## Files @@ -59,6 +78,8 @@ pipeline: |-------|------|----------|-------------| | `spend_file` | string | yes | Path to the JSON ledger file (created if missing) | | `max_budget` | float | yes | Daily budget in USD (must be > 0) | +| `input_cost_per_token` | float | no | USD per input/prompt token; prices streamed responses (whose header cost is 0) from parsed usage | +| `output_cost_per_token` | float | no | USD per output/completion token; prices streamed responses from parsed usage | ## Ledger Format