diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go index 22c8d886e..5e82a4b51 100644 --- a/authbridge/authlib/plugins/litellm_budgettrack/plugin.go +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin.go @@ -83,6 +83,28 @@ func (c budgetTrackConfig) cacheReadRate() float64 { // streaming frames until the terminal frame prices it. const stateKey = "litellm-budget-track" +// Wire values for costEvent.Source — consumers branch on these. +const ( + sourceGatewayHeader = "gateway-header" + sourceUsageFallback = "usage-fallback" +) + +// costEvent surfaces one priced response. Emitted per response when +// cost > 0; consumers see it as SessionEvent.Plugins["litellm-budget-track"]. +type costEvent struct { + CostUSD float64 `json:"cost_usd"` + + // Source is sourceGatewayHeader (authoritative x-litellm-response-cost) + // or sourceUsageFallback (priced from token counters, used for streamed + // responses whose header always reports 0). + Source string `json:"source"` + + // DailyTotalUSD is the ledger total after this response was added. + // DailyMaxUSD is the configured cap. + DailyTotalUSD float64 `json:"daily_total_usd"` + DailyMaxUSD float64 `json:"daily_max_usd"` +} + // 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 @@ -178,7 +200,9 @@ func (p *BudgetTrack) OnRequest(_ context.Context, pctx *pipeline.Context) pipel // instead; this remains for listeners that only call OnResponse. func (p *BudgetTrack) OnResponse(_ context.Context, pctx *pipeline.Context) pipeline.Action { if cost, _ := headerCost(pctx); cost > 0 { - p.accumulate(cost) + if total, ok := p.accumulate(cost); ok { + p.emitCost(pctx, cost, sourceGatewayHeader, total) + } } return pipeline.Action{Type: pipeline.Continue} } @@ -228,6 +252,7 @@ func (p *BudgetTrack) OnResponseFrame(_ context.Context, pctx *pipeline.Context, st.settled = true cost, present := headerCost(pctx) + source := sourceGatewayHeader 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 @@ -242,29 +267,49 @@ func (p *BudgetTrack) OnResponseFrame(_ context.Context, pctx *pipeline.Context, float64(st.cacheWriteTokens)*p.cfg.cacheWriteRate() + float64(st.cacheReadTokens)*p.cfg.cacheReadRate() + float64(st.outputTokens)*p.cfg.OutputCostPerToken + source = sourceUsageFallback } } if cost > 0 { - p.accumulate(cost) + if total, ok := p.accumulate(cost); ok { + p.emitCost(pctx, cost, source, total) + } } 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) { +// accumulate adds one priced call to today's ledger and persists it, +// returning the post-add TotalSpend and whether the cost was recorded. +// A non-finite or non-positive cost is ignored (added=false): 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) (dailyTotal float64, added bool) { if cost <= 0 || math.IsNaN(cost) || math.IsInf(cost, 0) { - return + return 0, false } p.mu.Lock() p.resetIfNewDay() p.ledger.TotalSpend += cost p.ledger.TotalCalls++ + total := p.ledger.TotalSpend p.saveLedger() p.mu.Unlock() + return total, true +} + +// emitCost writes the costEvent to pctx.Extensions.Custom; the listener +// forwards it to SessionEvent.Plugins under the plugin name. +func (p *BudgetTrack) emitCost(pctx *pipeline.Context, cost float64, source string, dailyTotal float64) { + if pctx.Extensions.Custom == nil { + pctx.Extensions.Custom = map[string]any{} + } + pctx.Extensions.Custom[p.Name()+pipeline.PluginEventSuffix] = costEvent{ + CostUSD: cost, + Source: source, + DailyTotalUSD: dailyTotal, + DailyMaxUSD: p.cfg.MaxBudget, + } } // headerCost returns the usable positive cost reported in the response headers diff --git a/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go index 45718ac37..5fcb5ed9e 100644 --- a/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go +++ b/authbridge/authlib/plugins/litellm_budgettrack/plugin_test.go @@ -521,3 +521,114 @@ func TestZeroCostHeaderStreamedPricesFromUsage(t *testing.T) { t.Errorf("streamed zero-header call not priced from usage: got %v want %v", got, want) } } + +// getCostEvent pulls the emitted costEvent out of pctx.Extensions.Custom +// under the same key the listener would read. Nil when nothing was emitted. +func getCostEvent(t *testing.T, pctx *pipeline.Context) *costEvent { + t.Helper() + if pctx.Extensions.Custom == nil { + return nil + } + v, ok := pctx.Extensions.Custom["litellm-budget-track"+pipeline.PluginEventSuffix].(costEvent) + if !ok { + return nil + } + return &v +} + +// TestEmitCost_HeaderPath: OnResponse (buffered) prices from the header +// and emits a costEvent tagged gateway-header, with DailyTotalUSD matching +// the ledger's post-add state and DailyMaxUSD from config. +func TestEmitCost_HeaderPath(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ResponseHeaders: http.Header{responseCostHeader: {"0.0025"}}} + p.OnResponse(context.Background(), pctx) + + ev := getCostEvent(t, pctx) + if ev == nil { + t.Fatal("no costEvent emitted") + } + if ev.CostUSD != 0.0025 { + t.Errorf("CostUSD = %v, want 0.0025", ev.CostUSD) + } + if ev.Source != sourceGatewayHeader { + t.Errorf("Source = %q, want %s", ev.Source, sourceGatewayHeader) + } + if ev.DailyTotalUSD != 0.0025 { + t.Errorf("DailyTotalUSD = %v, want 0.0025", ev.DailyTotalUSD) + } + if ev.DailyMaxUSD != 5.00 { + t.Errorf("DailyMaxUSD = %v, want 5.00", ev.DailyMaxUSD) + } + + // A second priced response on the same plugin must show + // DailyTotalUSD accumulating while CostUSD stays per-response. + // Without this a bug that emitted per-call cost as the daily + // total would pass every other assertion here. + pctx2 := &pipeline.Context{ResponseHeaders: http.Header{responseCostHeader: {"0.0025"}}} + p.OnResponse(context.Background(), pctx2) + ev2 := getCostEvent(t, pctx2) + if ev2 == nil { + t.Fatal("no costEvent emitted on second response") + } + if ev2.CostUSD != 0.0025 { + t.Errorf("second CostUSD = %v, want 0.0025 (per-response, not cumulative)", ev2.CostUSD) + } + if ev2.DailyTotalUSD != 0.0050 { + t.Errorf("second DailyTotalUSD = %v, want 0.0050 (accumulated across both responses)", ev2.DailyTotalUSD) + } +} + +// TestEmitCost_UsageFallback: streamed responses (0-cost header) get priced +// from the token counters and the emitted event names the fallback source. +func TestEmitCost_UsageFallback(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) + + ev := getCostEvent(t, pctx) + if ev == nil { + t.Fatal("no costEvent emitted") + } + if ev.Source != sourceUsageFallback { + t.Errorf("Source = %q, want %s", ev.Source, sourceUsageFallback) + } + want := 100*1e-6 + 40*5e-6 + if ev.CostUSD < want-1e-12 || ev.CostUSD > want+1e-12 { + t.Errorf("CostUSD = %v, want %v", ev.CostUSD, want) + } + // Daily fields populated on the streaming path too. + if ev.DailyTotalUSD < want-1e-12 || ev.DailyTotalUSD > want+1e-12 { + t.Errorf("DailyTotalUSD = %v, want %v (matches first-response cost)", ev.DailyTotalUSD, want) + } + if ev.DailyMaxUSD != 5.00 { + t.Errorf("DailyMaxUSD = %v, want 5.00", ev.DailyMaxUSD) + } +} + +// TestEmitCost_NoEmitWhenUnpriced: an OnResponse call whose header is missing +// or invalid must NOT emit a cost event (matches the ledger-untouched +// invariant asserted by TestOnResponseIgnoresMissingOrInvalid). +func TestEmitCost_NoEmitWhenUnpriced(t *testing.T) { + for _, tc := range []struct { + name string + headers http.Header + }{ + {"missing", http.Header{}}, + {"zero", http.Header{responseCostHeader: {"0"}}}, + {"unparseable", http.Header{responseCostHeader: {"abc"}}}, + } { + t.Run(tc.name, func(t *testing.T) { + p := configure(t, 5.00) + pctx := &pipeline.Context{ResponseHeaders: tc.headers} + p.OnResponse(context.Background(), pctx) + if ev := getCostEvent(t, pctx); ev != nil { + t.Errorf("costEvent emitted on unpriced response: %+v", *ev) + } + }) + } +} diff --git a/authbridge/docs/litellm-budgettrack-plugin.md b/authbridge/docs/litellm-budgettrack-plugin.md index c027b6a39..d63a73db0 100644 --- a/authbridge/docs/litellm-budgettrack-plugin.md +++ b/authbridge/docs/litellm-budgettrack-plugin.md @@ -152,6 +152,25 @@ The spend file (`spend-authbridge.json`) is a simple JSON object: 5. Write ledger to disk 6. Continue pipeline +## Observability + +Each priced response surfaces on the session-event stream at +`SessionEvent.Plugins["litellm-budget-track"]`, so consumers can read +per-response cost without duplicating the pricing math or reading the file +ledger. Unpriced responses (missing header + no per-token rates configured, or a +zero-cost cache hit) produce no event. + +| Field | Type | Meaning | +|-------|------|---------| +| `cost_usd` | float64 | Cost of this single response, in dollars. | +| `source` | string | `"gateway-header"` (authoritative — LiteLLM stamped the header) or `"usage-fallback"` (priced from token counters, used for streamed responses whose header always reports 0). | +| `daily_total_usd` | float64 | Ledger total after this response was added. | +| `daily_max_usd` | float64 | Configured daily cap (`max_budget`). | + +See [`plugin-reference.md#emitting-session-events`](./plugin-reference.md#emitting-session-events) +for how the listener promotes `pctx.Extensions.Custom` entries to +`SessionEvent.Plugins`. + ## Build The plugin is included by default in `authbridge-proxy` builds. To exclude: