feat: ✨ Split token counters

authbridge/authlib/pipeline/extensions.go

@@ -157,11 +157,8 @@ type InferenceExtension struct {
 	ToolChoice  any                `json:"toolChoice,omitempty"` // "auto" | "none" | {type,function:{name}}
 
 	// Response fields (populated after OnResponse runs).
-	Completion       string `json:"completion,omitempty"`
-	FinishReason     string `json:"finishReason,omitempty"`
-	PromptTokens     int    `json:"promptTokens,omitempty"`
-	CompletionTokens int    `json:"completionTokens,omitempty"`
-	TotalTokens      int    `json:"totalTokens,omitempty"`
+	Completion   string `json:"completion,omitempty"`
+	FinishReason string `json:"finishReason,omitempty"`
 
 	// ToolCalls are the tool invocations the model requested. Populated on
 	// three of the four response paths — both non-streaming dialects and
@@ -176,22 +173,26 @@ type InferenceExtension struct {
 	// whose calls were never captured.
 	ToolCalls []InferenceToolCall `json:"toolCalls,omitempty"`
 
-	// CacheWriteTokens and CacheReadTokens split the cached portion of
-	// PromptTokens by how it was billed. PromptTokens is the whole prompt
-	// (uncached input + cache writes + cache reads), which is the right
-	// number for context-size questions but the wrong one for cost: a
-	// provider that prices prompt caching charges a premium to *write* an
-	// entry and a steep discount to *read* one, so two requests with an
-	// identical PromptTokens can differ by an order of magnitude in price.
-	// Both counts arrive in the same usage block the totals come from, so
-	// recording them separately costs nothing and is the only way a
-	// consumer can tell a cache-warming turn from a cache-hit turn.
-	//
-	// Zero means "not reported" — providers that don't price caching (and
-	// the OpenAI dialect, which reports cached tokens in a different shape)
-	// leave these unset while PromptTokens stays authoritative.
-	CacheWriteTokens int `json:"cacheWriteTokens,omitempty"`
-	CacheReadTokens  int `json:"cacheReadTokens,omitempty"`
+	// Legacy aggregates, derived from the split fields below via
+	// parsercommon.TokenUsage.Fill.
+	PromptTokens     int `json:"promptTokens,omitempty"`
+	CompletionTokens int `json:"completionTokens,omitempty"`
+	TotalTokens      int `json:"totalTokens,omitempty"`
+
+	// Split token counters — provider-neutral shape published by every
+	// inference parser via parsercommon.TokenUsage.Fill.
+	InputTokens      int `json:"inputTokens,omitempty"`      // uncached prompt tokens
+	CacheReadTokens  int `json:"cacheReadTokens,omitempty"`  // served from cache
+	CacheWriteTokens int `json:"cacheWriteTokens,omitempty"` // written to cache
+	OutputTokens     int `json:"outputTokens,omitempty"`     // generated tokens
+	ReasoningTokens  int `json:"reasoningTokens,omitempty"`  // reasoning-only output
+
+	// PresentKinds names which split sub-kinds the provider populated.
+	// Zero on a set bit means "reported zero"; zero on an unset bit
+	// means "not exposed." Bit layout matches parsercommon.Kind
+	// (Input=1, CacheRead=2, CacheWrite=4, Output=8, Reasoning=16);
+	// typed as uint8 to avoid importing parsercommon here.
+	PresentKinds uint8 `json:"presentKinds,omitempty"`
 
 	// Classification — see MCPExtension.IsAction.
 	IsAction bool `json:"isAction,omitempty"`

authbridge/authlib/pipeline/session_test.go

@@ -71,9 +71,9 @@ func TestSessionEvent_JSONRoundTrip(t *testing.T) {
 		MCP: &MCPExtension{Method: "tools/call", RPCID: "m-2"},
 		Inference: &InferenceExtension{
 			Model: "gpt-4", Stream: stream, MaxTokens: &maxTok,
-			Messages:  []InferenceMessage{{Role: "user", Content: "hi"}},
-			Tools:     []InferenceTool{{Name: "get_weather", Description: "d"}},
-			ToolCalls: []InferenceToolCall{{ID: "c1", Name: "get_weather", Arguments: `{"city":"NYC"}`}},
+			Messages:   []InferenceMessage{{Role: "user", Content: "hi"}},
+			Tools:      []InferenceTool{{Name: "get_weather", Description: "d"}},
+			ToolCalls:  []InferenceToolCall{{ID: "c1", Name: "get_weather", Arguments: `{"city":"NYC"}`}},
 			Completion: "Hello, world!", FinishReason: "stop", TotalTokens: 17,
 		},
 		Identity:   &EventIdentity{Subject: "alice", ClientID: "agent-1"},

authbridge/authlib/pipeline/snapshot_test.go

@@ -158,10 +158,10 @@ func TestSnapshotPlugins_Empty(t *testing.T) {
 
 func TestSnapshotPlugins_FilterAndStripSuffix(t *testing.T) {
 	custom := map[string]any{
-		"rate-limiter" + PluginEventSuffix:  map[string]any{"remaining": 42},
-		"audit" + PluginEventSuffix:         "logged",
-		"some-internal-key":                 "should-not-appear", // no /event suffix
-		"more.internal.state":               struct{ X int }{X: 1},
+		"rate-limiter" + PluginEventSuffix: map[string]any{"remaining": 42},
+		"audit" + PluginEventSuffix:        "logged",
+		"some-internal-key":                "should-not-appear", // no /event suffix
+		"more.internal.state":              struct{ X int }{X: 1},
 	}
 	got := SnapshotPlugins(custom)
 

authbridge/authlib/plugins/inferenceparser/anthropic.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/rossoctl/cortex/authbridge/authlib/pipeline"
+	"github.com/rossoctl/cortex/authbridge/authlib/plugins/internal/parsercommon"
 )
 
 // anthropicMessagesPath is the Anthropic Messages API endpoint. Clients
@@ -124,15 +125,36 @@ func parseAnthropicRequest(body []byte) *pipeline.InferenceExtension {
 // anthropicUsage mirrors the Messages API usage block. The true input size is
 // input_tokens + cache_creation_input_tokens + cache_read_input_tokens (cached
 // context still counts as input); promptTotal sums them.
+//
+// Cache fields are *int so an omitted field on the wire (e.g. message_start
+// on the ?beta=true path) stays absent in Present rather than being asserted
+// as a reported zero.
 type anthropicUsage 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"`
+	InputTokens              int  `json:"input_tokens"`
+	OutputTokens             int  `json:"output_tokens"`
+	CacheCreationInputTokens *int `json:"cache_creation_input_tokens"`
+	CacheReadInputTokens     *int `json:"cache_read_input_tokens"`
 }
 
-func (u anthropicUsage) promptTotal() int {
-	return u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens
+// toNeutral maps Anthropic's usage onto TokenUsage. Input and Output are
+// always emitted by the Messages API; cache sub-fields are observed via
+// their pointers so an absent field stays absent in Present. Reasoning is
+// not exposed by Anthropic.
+func (u anthropicUsage) toNeutral() parsercommon.TokenUsage {
+	n := parsercommon.TokenUsage{
+		Input:   u.InputTokens,
+		Output:  u.OutputTokens,
+		Present: parsercommon.KindInput | parsercommon.KindOutput,
+	}
+	if u.CacheReadInputTokens != nil {
+		n.CacheRead = *u.CacheReadInputTokens
+		n.Present |= parsercommon.KindCacheRead
+	}
+	if u.CacheCreationInputTokens != nil {
+		n.CacheWrite = *u.CacheCreationInputTokens
+		n.Present |= parsercommon.KindCacheWrite
+	}
+	return n
 }
 
 // --- non-streaming response ---
@@ -180,11 +202,7 @@ func parseAnthropicJSON(body []byte, ext *pipeline.InferenceExtension) {
 	if resp.StopReason != "" {
 		ext.FinishReason = resp.StopReason
 	}
-	ext.PromptTokens = resp.Usage.promptTotal()
-	ext.CompletionTokens = resp.Usage.OutputTokens
-	ext.TotalTokens = ext.PromptTokens + ext.CompletionTokens
-	ext.CacheWriteTokens = resp.Usage.CacheCreationInputTokens
-	ext.CacheReadTokens = resp.Usage.CacheReadInputTokens
+	resp.Usage.toNeutral().Fill(ext)
 }
 
 // --- streaming ---
@@ -269,17 +287,6 @@ func (s *inferenceStreamState) closeAnthropicTool() {
 	s.openTool = nil
 }
 
-// totalAnthropicUsage derives the running total from the parts it was given.
-// It runs after every usage update rather than only when output tokens arrive,
-// because TotalTokens is the gate finalize uses to decide whether any count is
-// worth recording — so leaving it at zero discards a prompt size already known.
-// Two streams hit that: a terminal message_delta reporting the prompt with
-// output_tokens == 0, and a turn the caller interrupted after message_start,
-// which never reaches a message_delta at all. Both were billed for the prompt.
-func (s *inferenceStreamState) totalAnthropicUsage() {
-	s.usage.TotalTokens = s.usage.PromptTokens + s.usage.CompletionTokens
-}
-
 // foldAnthropicFrame folds one Messages SSE event into the running stream state.
 // The prompt size is taken as the largest total seen, because different Messages
 // API paths report it on different events: message_start on the plain path,
@@ -295,10 +302,8 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline
 	switch ev.Type {
 	case "message_start":
 		if ev.Message != nil {
-			state.usage.PromptTokens = ev.Message.Usage.promptTotal()
-			state.usage.CacheWriteTokens = ev.Message.Usage.CacheCreationInputTokens
-			state.usage.CacheReadTokens = ev.Message.Usage.CacheReadInputTokens
-			state.totalAnthropicUsage()
+			mergeAnthropicPromptMaxSeen(state, ev.Message.Usage.toNeutral())
+			state.hasUsage = true
 		}
 	case "content_block_start":
 		// A tool call opens here and is populated by later deltas. Text
@@ -325,33 +330,38 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline
 			ext.FinishReason = ev.Delta.StopReason
 		}
 		if ev.Usage != nil {
-			// The prompt side can arrive here rather than in message_start.
-			// Clients using the ?beta=true Messages path (Claude Code sends
-			// anthropic-beta: claude-code-*) get a message_start carrying only
-			// input_tokens, with cache_creation_input_tokens and
-			// cache_read_input_tokens deferred to message_delta — so reading
-			// the prompt size from message_start alone undercounts a cached
-			// agent request by orders of magnitude (a 33k-token turn recorded
-			// as 9). Take the larger value: on the non-beta path message_delta
-			// carries no input counts, and assigning unconditionally would
-			// clobber the correct message_start total with zero.
-			if p := ev.Usage.promptTotal(); p > state.usage.PromptTokens {
-				state.usage.PromptTokens = p
-				// Keep the split consistent with whichever usage block
-				// won the total, so the parts always sum into it.
-				state.usage.CacheWriteTokens = ev.Usage.CacheCreationInputTokens
-				state.usage.CacheReadTokens = ev.Usage.CacheReadInputTokens
+			// ?beta=true path defers cache counts from message_start to
+			// message_delta; non-beta path carries no input counts here.
+			// Max-seen per sub-field handles both without clobbering.
+			neutral := ev.Usage.toNeutral()
+			mergeAnthropicPromptMaxSeen(state, neutral)
+			if neutral.Output > 0 {
+				state.usage.Output = neutral.Output // cumulative
 			}
-			if ev.Usage.OutputTokens > 0 {
-				// usage.output_tokens in message_delta is cumulative — take
-				// the latest rather than accumulating.
-				state.usage.CompletionTokens = ev.Usage.OutputTokens
-			}
-			state.totalAnthropicUsage()
+			state.hasUsage = true
 		}
 	}
 }
 
+// mergeAnthropicPromptMaxSeen updates prompt-side sub-fields with
+// max-seen semantics so a later event carrying zero cannot clobber an
+// earlier real count. See foldAnthropicFrame for why both events need
+// this.
+func mergeAnthropicPromptMaxSeen(state *inferenceStreamState, incoming parsercommon.TokenUsage) {
+	// Presence is a union across events: once a sub-field is observed on
+	// the wire, later events that omit it must not clear the bit.
+	state.usage.Present |= incoming.Present
+	if incoming.Input > state.usage.Input {
+		state.usage.Input = incoming.Input
+	}
+	if incoming.CacheRead > state.usage.CacheRead {
+		state.usage.CacheRead = incoming.CacheRead
+	}
+	if incoming.CacheWrite > state.usage.CacheWrite {
+		state.usage.CacheWrite = incoming.CacheWrite
+	}
+}
+
 // parseAnthropicSSE folds a fully-buffered Messages SSE body. Mirrors
 // parseInferenceSSE for the legacy OnResponse path; the live listener uses
 // foldAnthropicFrame via OnResponseFrame instead.