Feat: Track streaming (SSE) LiteLLM cost in litellm-budget-track

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")
+}

authbridge/authlib/plugins/litellm_budgettrack/plugin.go

@@ -1,13 +1,27 @@
 // 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"
 	"sync"
@@ -33,6 +47,25 @@ 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
 }
 
 type spendLedger struct {
@@ -59,7 +92,7 @@ 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.",
+		Description: "Track LLM cost (response header or streamed usage) and enforce a daily budget.",
 	}
 }
 
@@ -73,6 +106,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 +135,156 @@ 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}
+	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
+		}
 	}
-	cost, err := strconv.ParseFloat(costStr, 64)
-	if err != nil || cost <= 0 {
+	if !last {
 		return pipeline.Action{Type: pipeline.Continue}
 	}
 
+	// Terminal frame: settle the cost exactly once.
+	cost := headerCost(pctx)
+	if cost <= 0 {
+		if st := pipeline.GetState[usageState](pctx, stateKey); st != nil {
+			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 cost reported in the response headers, or 0 when
+// absent/zero/unparseable. Streamed responses report 0 here.
+func headerCost(pctx *pipeline.Context) float64 {
+	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
+	}
+	cost, err := strconv.ParseFloat(costStr, 64)
+	// strconv.ParseFloat accepts "NaN" / "Inf"; reject non-finite (and
+	// non-positive) so a garbage header falls through to the usage path
+	// rather than poisoning the ledger.
+	if err != nil || cost <= 0 || math.IsNaN(cost) || math.IsInf(cost, 0) {
+		return 0
+	}
+	return cost
+}
+
+// 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 +311,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)
 )