feat: ✨ Cost emission on event stream

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

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

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: