From f64448ee8966dc0c121ba0e76b0d1e3da9e3ba77 Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Fri, 28 Aug 2026 16:59:18 -0600
Subject: [PATCH 1/7] :sparkles::white_check_mark: Split token counters

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 authbridge/authlib/pipeline/extensions.go     |  36 ++-
 .../plugins/inferenceparser/anthropic.go      |  84 +++---
 .../authlib/plugins/inferenceparser/plugin.go |  74 +++---
 .../plugins/inferenceparser/plugin_test.go    |  30 ---
 .../inferenceparser/splittokens_test.go       | 165 ++++++++++++
 .../internal/parsercommon/tokenusage.go       |  38 +++
 .../authlib/plugins/sessionbudget/plugin.go   | 248 +++++++++++++-----
 .../plugins/sessionbudget/plugin_test.go      |   8 +-
 .../plugins/sessionbudget/split_test.go       | 163 ++++++++++++
 9 files changed, 649 insertions(+), 197 deletions(-)
 create mode 100644 authbridge/authlib/plugins/inferenceparser/splittokens_test.go
 create mode 100644 authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
 create mode 100644 authbridge/authlib/plugins/sessionbudget/split_test.go

diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go
index 83bc17ef4..b7333e7c3 100644
--- a/authbridge/authlib/pipeline/extensions.go
+++ b/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,19 @@ 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
 
 	// Classification — see MCPExtension.IsAction.
 	IsAction bool `json:"isAction,omitempty"`
diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go
index 2d078d732..5656c2c6a 100644
--- a/authbridge/authlib/plugins/inferenceparser/anthropic.go
+++ b/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
@@ -132,7 +133,19 @@ type anthropicUsage struct {
 }
 
 func (u anthropicUsage) promptTotal() int {
-	return u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens
+	return u.toNeutral().PromptTotal()
+}
+
+// toNeutral maps Anthropic's usage block onto the neutral TokenUsage.
+// Straight rename — Anthropic already splits input / cache-read /
+// cache-write on the wire.
+func (u anthropicUsage) toNeutral() parsercommon.TokenUsage {
+	return parsercommon.TokenUsage{
+		Input:      u.InputTokens,
+		CacheRead:  u.CacheReadInputTokens,
+		CacheWrite: u.CacheCreationInputTokens,
+		Output:     u.OutputTokens,
+	}
 }
 
 // --- non-streaming response ---
@@ -180,11 +193,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 +278,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 +293,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 +321,35 @@ 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
-			}
-			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
+			// ?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
 			}
-			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) {
+	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.
diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index 1e0a7ac5b..936e77b75 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -161,14 +161,14 @@ func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context)
 	return pipeline.Action{Type: pipeline.Continue}
 }
 
-// inferenceStreamState is the scratch state kept on the extension for
-// the duration of a streaming response. Lives in pctx.Extensions.Custom
-// under a private key — kept off the public InferenceExtension shape so
-// the API stays clean. The struct accumulates the in-progress
-// completion until last=true triggers finalization.
+// inferenceStreamState is the scratch state kept on pctx.Extensions.Custom
+// for the duration of a streaming response. Provider-specific fold
+// functions normalize their wire format into the neutral usage field;
+// hasUsage flags whether any event carried usage counts (some providers
+// omit the block unless the client opts in).
 //
-// A streamed tool call is spread over many frames — id and name on the
-// opening frame, arguments as fragments after it — so it has to be
+// A streamed Anthropic tool call is spread over many frames — id and name
+// on the opening frame, arguments as fragments after it — so it has to be
 // assembled here rather than read off any single frame. toolCalls keeps
 // emission order; toolsByIndex resolves a fragment to its call, since
 // interleaved blocks (a text block and two tool calls) are only
@@ -176,7 +176,8 @@ func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context)
 // openTool is the fallback for a provider that omits the index.
 type inferenceStreamState struct {
 	completion strings.Builder
-	usage      inferenceUsage
+	usage      parsercommon.TokenUsage
+	hasUsage   bool
 
 	toolCalls    []*anthropicToolCallState
 	toolsByIndex map[int]*anthropicToolCallState
@@ -189,12 +190,8 @@ type inferenceStreamState struct {
 // after a streaming pass) is a no-op instead of a double-count.
 func (s *inferenceStreamState) finalize(ext *pipeline.InferenceExtension) {
 	ext.Completion = s.completion.String()
-	if s.usage.TotalTokens > 0 {
-		ext.PromptTokens = s.usage.PromptTokens
-		ext.CompletionTokens = s.usage.CompletionTokens
-		ext.TotalTokens = s.usage.TotalTokens
-		ext.CacheWriteTokens = s.usage.CacheWriteTokens
-		ext.CacheReadTokens = s.usage.CacheReadTokens
+	if s.hasUsage {
+		s.usage.Fill(ext)
 	}
 	if len(s.toolCalls) == 0 {
 		return
@@ -312,9 +309,8 @@ func foldOpenAIFrame(frame []byte, state *inferenceStreamState, ext *pipeline.In
 	// the provider reports it, and it may legitimately differ from
 	// prompt+completion.
 	if chunk.Usage.TotalTokens > 0 {
-		state.usage.PromptTokens = chunk.Usage.PromptTokens
-		state.usage.CompletionTokens = chunk.Usage.CompletionTokens
-		state.usage.TotalTokens = chunk.Usage.TotalTokens
+		state.usage = chunk.Usage.toNeutral()
+		state.hasUsage = true
 	}
 }
 
@@ -359,9 +355,7 @@ func parseInferenceJSON(body []byte, ext *pipeline.InferenceExtension) {
 			})
 		}
 	}
-	ext.PromptTokens = resp.Usage.PromptTokens
-	ext.CompletionTokens = resp.Usage.CompletionTokens
-	ext.TotalTokens = resp.Usage.TotalTokens
+	resp.Usage.toNeutral().Fill(ext)
 }
 
 // parseInferenceSSE concatenates content deltas across SSE events and captures
@@ -392,9 +386,7 @@ func parseInferenceSSE(body []byte, ext *pipeline.InferenceExtension) {
 			}
 		}
 		if chunk.Usage.TotalTokens > 0 {
-			ext.PromptTokens = chunk.Usage.PromptTokens
-			ext.CompletionTokens = chunk.Usage.CompletionTokens
-			ext.TotalTokens = chunk.Usage.TotalTokens
+			chunk.Usage.toNeutral().Fill(ext)
 		}
 	}
 	ext.Completion = completion.String()
@@ -446,21 +438,37 @@ type inferenceDelta struct {
 	Content string `json:"content"`
 }
 
-// inferenceUsage decodes the OpenAI usage block and doubles as the
-// dialect-neutral accumulator for a streaming response's token counts.
-//
-// The two cache fields are json:"-" because nothing on the wire fills them
-// in this shape: the Anthropic path sets them from its own usage struct
-// (see anthropicUsage), and the OpenAI dialect reports cached tokens under
-// a different key entirely. They live here so the shared finalize has one
-// place to read every count from.
+// inferenceUsage decodes the OpenAI usage block. See toNeutral for the
+// inclusive-prompt normalization.
 type inferenceUsage struct {
 	PromptTokens     int `json:"prompt_tokens"`
 	CompletionTokens int `json:"completion_tokens"`
 	TotalTokens      int `json:"total_tokens"`
 
-	CacheWriteTokens int `json:"-"`
-	CacheReadTokens  int `json:"-"`
+	PromptTokensDetails struct {
+		CachedTokens int `json:"cached_tokens"`
+	} `json:"prompt_tokens_details"`
+	CompletionTokensDetails struct {
+		ReasoningTokens int `json:"reasoning_tokens"`
+	} `json:"completion_tokens_details"`
+}
+
+// toNeutral maps OpenAI's usage onto TokenUsage. prompt_tokens includes
+// cached_tokens on the wire — subtract to get uncached input and clamp
+// at 0 for malformed responses. CacheWrite stays 0 (OpenAI bills cache
+// writes as ordinary input).
+func (u inferenceUsage) toNeutral() parsercommon.TokenUsage {
+	cached := u.PromptTokensDetails.CachedTokens
+	input := u.PromptTokens - cached
+	if input < 0 {
+		input = 0
+	}
+	return parsercommon.TokenUsage{
+		Input:     input,
+		CacheRead: cached,
+		Output:    u.CompletionTokens,
+		Reasoning: u.CompletionTokensDetails.ReasoningTokens,
+	}
 }
 
 type inferenceRequest struct {
diff --git a/authbridge/authlib/plugins/inferenceparser/plugin_test.go b/authbridge/authlib/plugins/inferenceparser/plugin_test.go
index 50a29ca6a..ec091b88d 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin_test.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin_test.go
@@ -470,36 +470,6 @@ func TestInferenceParser_OnResponse_SSE(t *testing.T) {
 	}
 }
 
-// TestFoldOpenAIFrame_PreservesCacheAccumulator pins the field-by-field copy in
-// foldOpenAIFrame. inferenceUsage is both the OpenAI wire shape and the
-// dialect-neutral accumulator, and its two cache fields are json:"-" — so they
-// are always zero in a decoded chunk, and assigning the whole struct would
-// clear whatever had been accumulated.
-//
-// Nothing on the OpenAI path fills those fields today, so this guards a
-// refactor rather than a live bug: the failure mode is silent, and it arrives
-// the moment someone wires up prompt_tokens_details.cached_tokens.
-func TestFoldOpenAIFrame_PreservesCacheAccumulator(t *testing.T) {
-	state := &inferenceStreamState{}
-	state.usage.CacheWriteTokens = 111
-	state.usage.CacheReadTokens = 222
-	ext := &pipeline.InferenceExtension{Model: "gpt-4", Stream: true}
-
-	frame := []byte(`{"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`)
-	foldOpenAIFrame(frame, state, ext)
-
-	if state.usage.CacheWriteTokens != 111 || state.usage.CacheReadTokens != 222 {
-		t.Errorf("cache counts = %d/%d, want 111/222 preserved",
-			state.usage.CacheWriteTokens, state.usage.CacheReadTokens)
-	}
-	// The wire-backed counts still land, TotalTokens taken as reported rather
-	// than recomputed from prompt+completion.
-	if state.usage.PromptTokens != 5 || state.usage.CompletionTokens != 2 || state.usage.TotalTokens != 7 {
-		t.Errorf("usage = %d/%d/%d, want 5/2/7",
-			state.usage.PromptTokens, state.usage.CompletionTokens, state.usage.TotalTokens)
-	}
-}
-
 func TestInferenceParser_OnResponse_InvalidJSON(t *testing.T) {
 	p := NewInferenceParser()
 	pctx := &pipeline.Context{
diff --git a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
new file mode 100644
index 000000000..bd2acec3f
--- /dev/null
+++ b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
@@ -0,0 +1,165 @@
+package inferenceparser
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/rossoctl/cortex/authbridge/authlib/pipeline"
+)
+
+// checkSplit surfaces every split field in one error message so a
+// mismatched counter is obvious.
+func checkSplit(t *testing.T, ext *pipeline.InferenceExtension, want pipeline.InferenceExtension) {
+	t.Helper()
+	got := fmt.Sprintf("input=%d cache_read=%d cache_write=%d output=%d reasoning=%d prompt=%d completion=%d total=%d",
+		ext.InputTokens, ext.CacheReadTokens, ext.CacheWriteTokens, ext.OutputTokens, ext.ReasoningTokens,
+		ext.PromptTokens, ext.CompletionTokens, ext.TotalTokens)
+	expected := fmt.Sprintf("input=%d cache_read=%d cache_write=%d output=%d reasoning=%d prompt=%d completion=%d total=%d",
+		want.InputTokens, want.CacheReadTokens, want.CacheWriteTokens, want.OutputTokens, want.ReasoningTokens,
+		want.PromptTokens, want.CompletionTokens, want.TotalTokens)
+	if got != expected {
+		t.Errorf("split token counters mismatch\n got: %s\nwant: %s", got, expected)
+	}
+}
+
+// Non-streaming Anthropic response with both cache_read and cache_write.
+func TestSplitTokens_AnthropicJSON(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-opus-4-8"}
+	body := []byte(`{
+		"content": [{"type": "text", "text": "ok"}],
+		"stop_reason": "end_turn",
+		"usage": {
+			"input_tokens": 50,
+			"cache_creation_input_tokens": 100,
+			"cache_read_input_tokens": 200,
+			"output_tokens": 25
+		}
+	}`)
+	parseAnthropicJSON(body, ext)
+
+	checkSplit(t, ext, pipeline.InferenceExtension{
+		InputTokens:      50,
+		CacheReadTokens:  200,
+		CacheWriteTokens: 100,
+		OutputTokens:     25,
+		ReasoningTokens:  0,
+		PromptTokens:     350, // 50 + 200 + 100
+		CompletionTokens: 25,
+		TotalTokens:      375,
+	})
+}
+
+// Non-beta SSE: prompt counts on message_start, output on message_delta.
+func TestSplitTokens_AnthropicSSE_NonBeta(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-opus-4-8"}
+	body := []byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":40,\"cache_creation_input_tokens\":10,\"cache_read_input_tokens\":30,\"output_tokens\":0}}}\n" +
+		"data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n" +
+		"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":12}}\n")
+	parseAnthropicSSE(body, ext)
+
+	checkSplit(t, ext, pipeline.InferenceExtension{
+		InputTokens:      40,
+		CacheReadTokens:  30,
+		CacheWriteTokens: 10,
+		OutputTokens:     12,
+		ReasoningTokens:  0,
+		PromptTokens:     80,
+		CompletionTokens: 12,
+		TotalTokens:      92,
+	})
+}
+
+// ?beta=true SSE: message_start carries only input_tokens; cache counts
+// arrive on message_delta. Exercises mergeAnthropicPromptMaxSeen.
+func TestSplitTokens_AnthropicSSE_Beta(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-opus-4-8"}
+	body := []byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":9,\"output_tokens\":0}}}\n" +
+		"data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n" +
+		"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"input_tokens\":9,\"cache_creation_input_tokens\":1234,\"cache_read_input_tokens\":32529,\"output_tokens\":399}}\n")
+	parseAnthropicSSE(body, ext)
+
+	checkSplit(t, ext, pipeline.InferenceExtension{
+		InputTokens:      9,
+		CacheReadTokens:  32529,
+		CacheWriteTokens: 1234,
+		OutputTokens:     399,
+		ReasoningTokens:  0,
+		PromptTokens:     33772,
+		CompletionTokens: 399,
+		TotalTokens:      34171,
+	})
+}
+
+// OpenAI JSON: inclusive-prompt normalization (Input = prompt - cached)
+// and reasoning pulled from completion_tokens_details.
+func TestSplitTokens_OpenAIJSON(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	body := []byte(`{
+		"choices": [{"message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
+		"usage": {
+			"prompt_tokens": 1500,
+			"completion_tokens": 400,
+			"total_tokens": 1900,
+			"prompt_tokens_details": {"cached_tokens": 200},
+			"completion_tokens_details": {"reasoning_tokens": 50}
+		}
+	}`)
+	parseInferenceJSON(body, ext)
+
+	checkSplit(t, ext, pipeline.InferenceExtension{
+		InputTokens:      1300,
+		CacheReadTokens:  200,
+		CacheWriteTokens: 0,
+		OutputTokens:     400,
+		ReasoningTokens:  50,
+		PromptTokens:     1500,
+		CompletionTokens: 400,
+		TotalTokens:      1900,
+	})
+}
+
+// Malformed OpenAI response where cached_tokens > prompt_tokens must
+// clamp InputTokens to 0, not go negative.
+func TestSplitTokens_OpenAIJSON_ClampNegativeInput(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	body := []byte(`{
+		"choices": [{"message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+		"usage": {
+			"prompt_tokens": 100,
+			"completion_tokens": 20,
+			"total_tokens": 120,
+			"prompt_tokens_details": {"cached_tokens": 250}
+		}
+	}`)
+	parseInferenceJSON(body, ext)
+
+	if ext.InputTokens != 0 {
+		t.Errorf("InputTokens = %d, want 0 (clamped)", ext.InputTokens)
+	}
+	if ext.CacheReadTokens != 250 {
+		t.Errorf("CacheReadTokens = %d, want 250", ext.CacheReadTokens)
+	}
+}
+
+// TestSplitTokens_OpenAISSE covers the OpenAI streaming shape with
+// stream_options.include_usage — the usage block arrives on the final
+// chunk and must map onto the neutral shape via the same normalization
+// as the JSON path.
+func TestSplitTokens_OpenAISSE(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o", Stream: true}
+	body := []byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n" +
+		"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":800,\"completion_tokens\":150,\"total_tokens\":950,\"prompt_tokens_details\":{\"cached_tokens\":600},\"completion_tokens_details\":{\"reasoning_tokens\":25}}}\n" +
+		"data: [DONE]\n")
+	parseInferenceSSE(body, ext)
+
+	checkSplit(t, ext, pipeline.InferenceExtension{
+		InputTokens:      200,
+		CacheReadTokens:  600,
+		CacheWriteTokens: 0,
+		OutputTokens:     150,
+		ReasoningTokens:  25,
+		PromptTokens:     800,
+		CompletionTokens: 150,
+		TotalTokens:      950,
+	})
+}
diff --git a/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
new file mode 100644
index 000000000..2852f29b0
--- /dev/null
+++ b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
@@ -0,0 +1,38 @@
+package parsercommon
+
+import "github.com/rossoctl/cortex/authbridge/authlib/pipeline"
+
+// TokenUsage is the provider-neutral token accounting shape. Provider
+// parsers normalize their wire format (e.g. OpenAI's inclusive
+// prompt_tokens) into these fields before publishing via Fill.
+type TokenUsage struct {
+	Input      int // uncached prompt tokens
+	CacheRead  int // prompt tokens served from cache
+	CacheWrite int // prompt tokens written to cache
+	Output     int // generated completion tokens
+	Reasoning  int // reasoning-only output (subset of Output)
+}
+
+// PromptTotal is the sum of all prompt-side sub-kinds.
+func (u TokenUsage) PromptTotal() int {
+	return u.Input + u.CacheRead + u.CacheWrite
+}
+
+// Total is PromptTotal + Output. Reasoning is a subset of Output and
+// intentionally not added.
+func (u TokenUsage) Total() int {
+	return u.PromptTotal() + u.Output
+}
+
+// Fill writes the split counters and derived legacy aggregates onto ext.
+func (u TokenUsage) Fill(ext *pipeline.InferenceExtension) {
+	ext.InputTokens = u.Input
+	ext.CacheReadTokens = u.CacheRead
+	ext.CacheWriteTokens = u.CacheWrite
+	ext.OutputTokens = u.Output
+	ext.ReasoningTokens = u.Reasoning
+
+	ext.PromptTokens = u.PromptTotal()
+	ext.CompletionTokens = u.Output
+	ext.TotalTokens = u.Total()
+}
diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go
index 0cd5cfcb0..9df47f75e 100644
--- a/authbridge/authlib/plugins/sessionbudget/plugin.go
+++ b/authbridge/authlib/plugins/sessionbudget/plugin.go
@@ -26,6 +26,11 @@ import (
 type config struct {
 	RedisURL               string `json:"redis_url" required:"true" description:"Redis/Valkey connection URL."`
 	MaxTokens              int64  `json:"max_tokens" description:"Cumulative token ceiling per session. 0 = no limit."`
+	MaxInputTokens         int64  `json:"max_input_tokens" description:"Per-kind ceiling for uncached prompt tokens. 0 = no limit."`
+	MaxCacheReadTokens     int64  `json:"max_cache_read_tokens" description:"Per-kind ceiling for prompt tokens served from cache. 0 = no limit."`
+	MaxCacheWriteTokens    int64  `json:"max_cache_write_tokens" description:"Per-kind ceiling for prompt tokens written to cache. 0 = no limit."`
+	MaxOutputTokens        int64  `json:"max_output_tokens" description:"Per-kind ceiling for generated completion tokens. 0 = no limit."`
+	MaxReasoningTokens     int64  `json:"max_reasoning_tokens" description:"Per-kind ceiling for reasoning-only output tokens (subset of output). 0 = no limit."`
 	MaxCalls               int64  `json:"max_calls" description:"Max LLM/inference calls per session. Only inference-parser output increments this counter; MCP tool calls and other outbound traffic do not. Once the limit is reached, all subsequent outbound requests (including MCP tool calls) are blocked until the session resets. 0 = no limit."`
 	MaxDurationSeconds     int64  `json:"max_duration_seconds" description:"Wall-clock session lifetime in seconds. 0 = no limit."`
 	OnExceed               string `json:"on_exceed" description:"Action on breach: deny, observe (shadow), or pause (HITL webhook approval)." default:"deny" enum:"deny,observe,pause"`
@@ -53,10 +58,18 @@ type approvalFlight struct {
 }
 
 type counters struct {
-	tokens         int64
-	calls          int64
-	startedAt      time.Time
-	lastApprovedAt time.Time
+	tokens int64
+	// Per-kind sub-counters. Stored as separate fields (rather than derived
+	// from tokens) so a future weighted-total policy can multiply them
+	// without reshaping the counter first.
+	inputTokens      int64
+	cacheReadTokens  int64
+	cacheWriteTokens int64
+	outputTokens     int64
+	reasoningTokens  int64
+	calls            int64
+	startedAt        time.Time
+	lastApprovedAt   time.Time
 	// pendingApproval is non-nil while a webhook call for this session is in
 	// flight. Concurrent breaches wait on flight.done; the leader publishes
 	// flight.approved before closing done, then clears this field.
@@ -122,8 +135,10 @@ func (p *SessionBudget) Configure(raw json.RawMessage) error {
 	if p.cfg.RedisURL == "" {
 		return fmt.Errorf("session-budget: redis_url is required")
 	}
-	if p.cfg.MaxTokens <= 0 && p.cfg.MaxCalls <= 0 && p.cfg.MaxDurationSeconds <= 0 {
-		return fmt.Errorf("session-budget: at least one limit (max_tokens, max_calls, max_duration_seconds) must be > 0")
+	if p.cfg.MaxTokens <= 0 && p.cfg.MaxCalls <= 0 && p.cfg.MaxDurationSeconds <= 0 &&
+		p.cfg.MaxInputTokens <= 0 && p.cfg.MaxCacheReadTokens <= 0 && p.cfg.MaxCacheWriteTokens <= 0 &&
+		p.cfg.MaxOutputTokens <= 0 && p.cfg.MaxReasoningTokens <= 0 {
+		return fmt.Errorf("session-budget: at least one limit (max_tokens, max_input_tokens, max_cache_read_tokens, max_cache_write_tokens, max_output_tokens, max_reasoning_tokens, max_calls, max_duration_seconds) must be > 0")
 	}
 	if p.cfg.SessionTTLSeconds < 0 {
 		return fmt.Errorf("session-budget: session_ttl_seconds must be > 0 (got %d)", p.cfg.SessionTTLSeconds)
@@ -381,7 +396,14 @@ func (p *SessionBudget) OnResponseFrame(_ context.Context, pctx *pipeline.Contex
 		return pipeline.Action{Type: pipeline.Continue}
 	}
 
-	tokens := int64(inf.TotalTokens)
+	delta := tokenDelta{
+		total:      int64(inf.TotalTokens),
+		input:      int64(inf.InputTokens),
+		cacheRead:  int64(inf.CacheReadTokens),
+		cacheWrite: int64(inf.CacheWriteTokens),
+		output:     int64(inf.OutputTokens),
+		reasoning:  int64(inf.ReasoningTokens),
+	}
 
 	p.mu.Lock()
 	c, ok := p.cache[sessionID]
@@ -389,22 +411,37 @@ func (p *SessionBudget) OnResponseFrame(_ context.Context, pctx *pipeline.Contex
 		c = &counters{startedAt: time.Now()}
 		p.cache[sessionID] = c
 	}
-	c.tokens += tokens
+	c.tokens += delta.total
+	c.inputTokens += delta.input
+	c.cacheReadTokens += delta.cacheRead
+	c.cacheWriteTokens += delta.cacheWrite
+	c.outputTokens += delta.output
+	c.reasoningTokens += delta.reasoning
 	c.calls++
 	c.pendingWrites++
 	p.mu.Unlock()
 
-	go p.accumulate(sessionID, tokens)
+	go p.accumulate(sessionID, delta)
 
 	return pipeline.Action{Type: pipeline.Continue}
 }
 
 func (p *SessionBudget) buildDetails(snap *counters) map[string]any {
 	details := map[string]any{
-		"spent_tokens": snap.tokens,
-		"spent_calls":  snap.calls,
-		"token_limit":  p.cfg.MaxTokens,
-		"call_limit":   p.cfg.MaxCalls,
+		"spent_tokens":             snap.tokens,
+		"spent_input_tokens":       snap.inputTokens,
+		"spent_cache_read_tokens":  snap.cacheReadTokens,
+		"spent_cache_write_tokens": snap.cacheWriteTokens,
+		"spent_output_tokens":      snap.outputTokens,
+		"spent_reasoning_tokens":   snap.reasoningTokens,
+		"spent_calls":              snap.calls,
+		"token_limit":              p.cfg.MaxTokens,
+		"input_token_limit":        p.cfg.MaxInputTokens,
+		"cache_read_token_limit":   p.cfg.MaxCacheReadTokens,
+		"cache_write_token_limit":  p.cfg.MaxCacheWriteTokens,
+		"output_token_limit":       p.cfg.MaxOutputTokens,
+		"reasoning_token_limit":    p.cfg.MaxReasoningTokens,
+		"call_limit":               p.cfg.MaxCalls,
 	}
 	if p.cfg.MaxDurationSeconds > 0 && !snap.startedAt.IsZero() {
 		details["duration_seconds"] = int64(time.Since(snap.startedAt).Seconds())
@@ -414,14 +451,24 @@ func (p *SessionBudget) buildDetails(snap *counters) map[string]any {
 }
 
 type pauseRequest struct {
-	SessionID       string `json:"session_id"`
-	Reason          string `json:"reason"`
-	SpentTokens     int64  `json:"spent_tokens"`
-	SpentCalls      int64  `json:"spent_calls"`
-	TokenLimit      int64  `json:"token_limit"`
-	CallLimit       int64  `json:"call_limit"`
-	DurationSeconds int64  `json:"duration_seconds,omitempty"`
-	DurationLimit   int64  `json:"duration_limit,omitempty"`
+	SessionID             string `json:"session_id"`
+	Reason                string `json:"reason"`
+	SpentTokens           int64  `json:"spent_tokens"`
+	SpentInputTokens      int64  `json:"spent_input_tokens,omitempty"`
+	SpentCacheReadTokens  int64  `json:"spent_cache_read_tokens,omitempty"`
+	SpentCacheWriteTokens int64  `json:"spent_cache_write_tokens,omitempty"`
+	SpentOutputTokens     int64  `json:"spent_output_tokens,omitempty"`
+	SpentReasoningTokens  int64  `json:"spent_reasoning_tokens,omitempty"`
+	SpentCalls            int64  `json:"spent_calls"`
+	TokenLimit            int64  `json:"token_limit"`
+	InputTokenLimit       int64  `json:"input_token_limit,omitempty"`
+	CacheReadTokenLimit   int64  `json:"cache_read_token_limit,omitempty"`
+	CacheWriteTokenLimit  int64  `json:"cache_write_token_limit,omitempty"`
+	OutputTokenLimit      int64  `json:"output_token_limit,omitempty"`
+	ReasoningTokenLimit   int64  `json:"reasoning_token_limit,omitempty"`
+	CallLimit             int64  `json:"call_limit"`
+	DurationSeconds       int64  `json:"duration_seconds,omitempty"`
+	DurationLimit         int64  `json:"duration_limit,omitempty"`
 }
 
 type pauseResponse struct {
@@ -435,12 +482,22 @@ func (p *SessionBudget) callPauseWebhook(sessionID, reason string, snap *counter
 	defer cancel()
 
 	body := pauseRequest{
-		SessionID:   sessionID,
-		Reason:      reason,
-		SpentTokens: snap.tokens,
-		SpentCalls:  snap.calls,
-		TokenLimit:  p.cfg.MaxTokens,
-		CallLimit:   p.cfg.MaxCalls,
+		SessionID:             sessionID,
+		Reason:                reason,
+		SpentTokens:           snap.tokens,
+		SpentInputTokens:      snap.inputTokens,
+		SpentCacheReadTokens:  snap.cacheReadTokens,
+		SpentCacheWriteTokens: snap.cacheWriteTokens,
+		SpentOutputTokens:     snap.outputTokens,
+		SpentReasoningTokens:  snap.reasoningTokens,
+		SpentCalls:            snap.calls,
+		TokenLimit:            p.cfg.MaxTokens,
+		InputTokenLimit:       p.cfg.MaxInputTokens,
+		CacheReadTokenLimit:   p.cfg.MaxCacheReadTokens,
+		CacheWriteTokenLimit:  p.cfg.MaxCacheWriteTokens,
+		OutputTokenLimit:      p.cfg.MaxOutputTokens,
+		ReasoningTokenLimit:   p.cfg.MaxReasoningTokens,
+		CallLimit:             p.cfg.MaxCalls,
 	}
 	if p.cfg.MaxDurationSeconds > 0 && !snap.startedAt.IsZero() {
 		body.DurationSeconds = int64(time.Since(snap.startedAt).Seconds())
@@ -488,6 +545,21 @@ func (p *SessionBudget) evaluate(c *counters) string {
 	if p.cfg.MaxTokens > 0 && c.tokens >= p.cfg.MaxTokens {
 		return fmt.Sprintf("token limit reached: %d/%d", c.tokens, p.cfg.MaxTokens)
 	}
+	if p.cfg.MaxInputTokens > 0 && c.inputTokens >= p.cfg.MaxInputTokens {
+		return fmt.Sprintf("input token limit reached: %d/%d", c.inputTokens, p.cfg.MaxInputTokens)
+	}
+	if p.cfg.MaxCacheReadTokens > 0 && c.cacheReadTokens >= p.cfg.MaxCacheReadTokens {
+		return fmt.Sprintf("cache-read token limit reached: %d/%d", c.cacheReadTokens, p.cfg.MaxCacheReadTokens)
+	}
+	if p.cfg.MaxCacheWriteTokens > 0 && c.cacheWriteTokens >= p.cfg.MaxCacheWriteTokens {
+		return fmt.Sprintf("cache-write token limit reached: %d/%d", c.cacheWriteTokens, p.cfg.MaxCacheWriteTokens)
+	}
+	if p.cfg.MaxOutputTokens > 0 && c.outputTokens >= p.cfg.MaxOutputTokens {
+		return fmt.Sprintf("output token limit reached: %d/%d", c.outputTokens, p.cfg.MaxOutputTokens)
+	}
+	if p.cfg.MaxReasoningTokens > 0 && c.reasoningTokens >= p.cfg.MaxReasoningTokens {
+		return fmt.Sprintf("reasoning token limit reached: %d/%d", c.reasoningTokens, p.cfg.MaxReasoningTokens)
+	}
 	if p.cfg.MaxCalls > 0 && c.calls >= p.cfg.MaxCalls {
 		return fmt.Sprintf("call limit reached: %d/%d", c.calls, p.cfg.MaxCalls)
 	}
@@ -500,8 +572,20 @@ func (p *SessionBudget) evaluate(c *counters) string {
 	return ""
 }
 
+// tokenDelta is one response's contribution to the counters, both the
+// aggregate `total` and each split sub-kind. Held as one value so the
+// async accumulate goroutine gets everything in one shot.
+type tokenDelta struct {
+	total      int64
+	input      int64
+	cacheRead  int64
+	cacheWrite int64
+	output     int64
+	reasoning  int64
+}
+
 // accumulate writes counters to Redis. On failure, writes are dropped (fail-open).
-func (p *SessionBudget) accumulate(sessionID string, tokens int64) {
+func (p *SessionBudget) accumulate(sessionID string, delta tokenDelta) {
 	defer func() {
 		p.mu.Lock()
 		if cc, ok := p.cache[sessionID]; ok && cc.pendingWrites > 0 {
@@ -516,9 +600,25 @@ func (p *SessionBudget) accumulate(sessionID string, tokens int64) {
 	key := p.redisKey(sessionID)
 	ttl := time.Duration(p.cfg.SessionTTLSeconds) * time.Second
 
-	if tokens > 0 {
-		if _, err := p.store.HashIncr(ctx, key, "tokens", tokens); err != nil {
-			p.log.Warn("redis HashIncr tokens failed", "session", sessionID, "err", err)
+	// One HashIncr per non-zero sub-kind. Zero deltas are skipped so old
+	// sessions whose per-kind counters were never written stay absent in
+	// Redis rather than accumulating a stream of no-op fields.
+	for _, kv := range []struct {
+		field string
+		v     int64
+	}{
+		{"tokens", delta.total},
+		{"input_tokens", delta.input},
+		{"cache_read_tokens", delta.cacheRead},
+		{"cache_write_tokens", delta.cacheWrite},
+		{"output_tokens", delta.output},
+		{"reasoning_tokens", delta.reasoning},
+	} {
+		if kv.v <= 0 {
+			continue
+		}
+		if _, err := p.store.HashIncr(ctx, key, kv.field, kv.v); err != nil {
+			p.log.Warn("redis HashIncr failed", "session", sessionID, "field", kv.field, "err", err)
 		}
 	}
 
@@ -571,17 +671,12 @@ func (p *SessionBudget) hydrateCache(sessionID string) bool {
 		if len(fields) == 0 {
 			return false, nil
 		}
-		tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64)
-		calls, _ := strconv.ParseInt(fields["calls"], 10, 64)
-		var startedAt time.Time
-		if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil {
-			startedAt = time.Unix(ts, 0)
-		}
+		parsed := parseCountersFromFields(fields)
 		p.mu.Lock()
 		// Do not overwrite: OnResponseFrame may have seeded an entry between
 		// our HashGet and this lock. Its counters are fresher than Redis.
 		if _, exists := p.cache[sessionID]; !exists {
-			p.cache[sessionID] = &counters{tokens: tokens, calls: calls, startedAt: startedAt}
+			p.cache[sessionID] = parsed
 		}
 		p.mu.Unlock()
 		return true, nil
@@ -623,42 +718,44 @@ func (p *SessionBudget) refreshCache() {
 			continue
 		}
 
-		tokens, _ := strconv.ParseInt(fields["tokens"], 10, 64)
-		calls, _ := strconv.ParseInt(fields["calls"], 10, 64)
-		var startedAt time.Time
-		if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil {
-			startedAt = time.Unix(ts, 0)
-		}
+		parsed := parseCountersFromFields(fields)
 
 		p.mu.Lock()
-		var lastApprovedAt time.Time
-		var pendingApproval *approvalFlight
-		var pendingWrites int
 		if existing, ok := p.cache[sessionID]; ok {
-			// Take the max of local and Redis to avoid regressing counters when
-			// in-flight accumulate goroutines haven't committed to Redis yet.
-			if tokens < existing.tokens {
-				tokens = existing.tokens
+			// Take the max of local and Redis on every counter to avoid
+			// regressing when in-flight accumulate goroutines haven't
+			// committed yet. Applied uniformly across the aggregate and
+			// each per-kind sub-counter.
+			if parsed.tokens < existing.tokens {
+				parsed.tokens = existing.tokens
+			}
+			if parsed.inputTokens < existing.inputTokens {
+				parsed.inputTokens = existing.inputTokens
+			}
+			if parsed.cacheReadTokens < existing.cacheReadTokens {
+				parsed.cacheReadTokens = existing.cacheReadTokens
+			}
+			if parsed.cacheWriteTokens < existing.cacheWriteTokens {
+				parsed.cacheWriteTokens = existing.cacheWriteTokens
 			}
-			if calls < existing.calls {
-				calls = existing.calls
+			if parsed.outputTokens < existing.outputTokens {
+				parsed.outputTokens = existing.outputTokens
 			}
-			if startedAt.IsZero() && !existing.startedAt.IsZero() {
-				startedAt = existing.startedAt
+			if parsed.reasoningTokens < existing.reasoningTokens {
+				parsed.reasoningTokens = existing.reasoningTokens
 			}
-			lastApprovedAt = existing.lastApprovedAt
+			if parsed.calls < existing.calls {
+				parsed.calls = existing.calls
+			}
+			if parsed.startedAt.IsZero() && !existing.startedAt.IsZero() {
+				parsed.startedAt = existing.startedAt
+			}
+			parsed.lastApprovedAt = existing.lastApprovedAt
 			// Preserve mid-webhook: dropping would let a concurrent breach fire a duplicate.
-			pendingApproval = existing.pendingApproval
-			pendingWrites = existing.pendingWrites
-		}
-		p.cache[sessionID] = &counters{
-			tokens:          tokens,
-			calls:           calls,
-			startedAt:       startedAt,
-			lastApprovedAt:  lastApprovedAt,
-			pendingApproval: pendingApproval,
-			pendingWrites:   pendingWrites,
+			parsed.pendingApproval = existing.pendingApproval
+			parsed.pendingWrites = existing.pendingWrites
 		}
+		p.cache[sessionID] = parsed
 		p.mu.Unlock()
 	}
 }
@@ -680,6 +777,25 @@ func (p *SessionBudget) redisKey(sessionID string) string {
 	return "session-budget:" + sessionID
 }
 
+// parseCountersFromFields turns a Redis hash into a counters value. Missing
+// per-kind fields parse as 0, which is what legacy sessions written before
+// the split-token migration look like — they keep enforcing max_tokens and
+// start their per-kind counters from zero.
+func parseCountersFromFields(fields map[string]string) *counters {
+	c := &counters{}
+	c.tokens, _ = strconv.ParseInt(fields["tokens"], 10, 64)
+	c.inputTokens, _ = strconv.ParseInt(fields["input_tokens"], 10, 64)
+	c.cacheReadTokens, _ = strconv.ParseInt(fields["cache_read_tokens"], 10, 64)
+	c.cacheWriteTokens, _ = strconv.ParseInt(fields["cache_write_tokens"], 10, 64)
+	c.outputTokens, _ = strconv.ParseInt(fields["output_tokens"], 10, 64)
+	c.reasoningTokens, _ = strconv.ParseInt(fields["reasoning_tokens"], 10, 64)
+	c.calls, _ = strconv.ParseInt(fields["calls"], 10, 64)
+	if ts, err := strconv.ParseInt(fields["started_at"], 10, 64); err == nil {
+		c.startedAt = time.Unix(ts, 0)
+	}
+	return c
+}
+
 var (
 	_ pipeline.Plugin             = (*SessionBudget)(nil)
 	_ pipeline.Configurable       = (*SessionBudget)(nil)
diff --git a/authbridge/authlib/plugins/sessionbudget/plugin_test.go b/authbridge/authlib/plugins/sessionbudget/plugin_test.go
index cfb3a6d34..b44caa880 100644
--- a/authbridge/authlib/plugins/sessionbudget/plugin_test.go
+++ b/authbridge/authlib/plugins/sessionbudget/plugin_test.go
@@ -336,7 +336,7 @@ func TestAccumulate_WritesToStore(t *testing.T) {
 	store := newMemStore()
 	p.store = store
 
-	p.accumulate("sess-1", 100)
+	p.accumulate("sess-1", tokenDelta{total: 100})
 
 	fields, _ := store.HashGet(context.Background(), "session-budget:sess-1")
 	if fields["tokens"] != "100" {
@@ -355,7 +355,7 @@ func TestAccumulate_ZeroTokens(t *testing.T) {
 	store := newMemStore()
 	p.store = store
 
-	p.accumulate("sess-1", 0)
+	p.accumulate("sess-1", tokenDelta{})
 
 	fields, _ := store.HashGet(context.Background(), "session-budget:sess-1")
 	if fields["tokens"] != "" {
@@ -394,7 +394,7 @@ func TestAccumulate_ExpireSelfHealsAfterFailure(t *testing.T) {
 	p.store = spy
 
 	// Call #1: HashSetNX succeeds (started_at recorded), Expire fails.
-	p.accumulate("sess", 10)
+	p.accumulate("sess", tokenDelta{total: 10})
 	inner.mu.Lock()
 	_, hasTTL := inner.ttls["session-budget:sess"]
 	inner.mu.Unlock()
@@ -405,7 +405,7 @@ func TestAccumulate_ExpireSelfHealsAfterFailure(t *testing.T) {
 	// Call #2: HashSetNX returns false (started_at already set), but Expire
 	// must still run and restore the TTL. Pre-fix code gated Expire on
 	// HashSetNX-returned-true and would leave the key TTL-less forever.
-	p.accumulate("sess", 10)
+	p.accumulate("sess", tokenDelta{total: 10})
 	inner.mu.Lock()
 	ttl, hasTTL := inner.ttls["session-budget:sess"]
 	inner.mu.Unlock()
diff --git a/authbridge/authlib/plugins/sessionbudget/split_test.go b/authbridge/authlib/plugins/sessionbudget/split_test.go
new file mode 100644
index 000000000..3a78b3f20
--- /dev/null
+++ b/authbridge/authlib/plugins/sessionbudget/split_test.go
@@ -0,0 +1,163 @@
+package sessionbudget
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/rossoctl/cortex/authbridge/authlib/pipeline"
+)
+
+// newSplitPlugin builds a plugin with one per-kind limit set and no
+// aggregate token limit, so a test can isolate the per-kind path from
+// the existing MaxTokens enforcement.
+func newSplitPlugin(t *testing.T, field string, limit int64) *SessionBudget {
+	t.Helper()
+	p := New()
+	cfg := fmt.Sprintf(`{
+		"redis_url": "mem://test",
+		%q: %d,
+		"refresh_interval": "100ms"
+	}`, field, limit)
+	if err := p.Configure(json.RawMessage(cfg)); err != nil {
+		t.Fatalf("Configure: %v", err)
+	}
+	p.store = newMemStore()
+	return p
+}
+
+// makePctxSplit is makePctx with the per-kind fields populated. The
+// aggregate TotalTokens is left as the sum of the sub-kinds so the
+// existing MaxTokens plumbing keeps seeing the same value it did
+// before this feature landed.
+func makePctxSplit(sessionID string, input, cacheRead, cacheWrite, output, reasoning int) *pipeline.Context {
+	total := input + cacheRead + cacheWrite + output
+	return &pipeline.Context{
+		Direction: pipeline.Outbound,
+		Headers:   http.Header{},
+		Session:   &pipeline.SessionView{ID: sessionID},
+		Extensions: pipeline.Extensions{
+			Inference: &pipeline.InferenceExtension{
+				TotalTokens:      total,
+				InputTokens:      input,
+				CacheReadTokens:  cacheRead,
+				CacheWriteTokens: cacheWrite,
+				OutputTokens:     output,
+				ReasoningTokens:  reasoning,
+			},
+		},
+	}
+}
+
+// TestOnResponseFrame_AccumulatesSplit confirms every sub-counter
+// increments from a single response frame. This is the base contract
+// the per-kind limits build on — if a sub-counter is stuck at zero,
+// its limit can never trip.
+func TestOnResponseFrame_AccumulatesSplit(t *testing.T) {
+	p := newTestPlugin(0, 1, 0)
+	pctx := makePctxSplit("sess-1", 100, 50, 25, 200, 75)
+
+	if act := p.OnResponseFrame(context.Background(), pctx, nil, true); act.Type != pipeline.Continue {
+		t.Fatalf("frame: got %v want Continue", act.Type)
+	}
+
+	p.mu.RLock()
+	c := p.cache["sess-1"]
+	p.mu.RUnlock()
+	if c == nil {
+		t.Fatal("no cache entry")
+	}
+
+	checks := []struct {
+		name string
+		got  int64
+		want int64
+	}{
+		{"tokens", c.tokens, 375},
+		{"inputTokens", c.inputTokens, 100},
+		{"cacheReadTokens", c.cacheReadTokens, 50},
+		{"cacheWriteTokens", c.cacheWriteTokens, 25},
+		{"outputTokens", c.outputTokens, 200},
+		{"reasoningTokens", c.reasoningTokens, 75},
+	}
+	for _, ch := range checks {
+		if ch.got != ch.want {
+			t.Errorf("%s = %d, want %d", ch.name, ch.got, ch.want)
+		}
+	}
+}
+
+// TestOnRequest_RejectsAtSplitLimit runs each per-kind limit through
+// the same at-limit assertion so a regression in any one branch of
+// evaluate() surfaces as its own failing row.
+func TestOnRequest_RejectsAtSplitLimit(t *testing.T) {
+	cases := []struct {
+		name     string
+		field    string
+		limit    int64
+		seed     func(*counters)
+		wantWord string
+	}{
+		{"input", "max_input_tokens", 100, func(c *counters) { c.inputTokens = 100 }, "input"},
+		{"cache_read", "max_cache_read_tokens", 100, func(c *counters) { c.cacheReadTokens = 100 }, "cache-read"},
+		{"cache_write", "max_cache_write_tokens", 100, func(c *counters) { c.cacheWriteTokens = 100 }, "cache-write"},
+		{"output", "max_output_tokens", 100, func(c *counters) { c.outputTokens = 100 }, "output"},
+		{"reasoning", "max_reasoning_tokens", 100, func(c *counters) { c.reasoningTokens = 100 }, "reasoning"},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			p := newSplitPlugin(t, tc.field, tc.limit)
+			c := &counters{}
+			tc.seed(c)
+			p.mu.Lock()
+			p.cache["sess-1"] = c
+			p.mu.Unlock()
+
+			action := p.OnRequest(context.Background(), makePctx("sess-1", 0))
+			if action.Type != pipeline.Reject {
+				t.Fatalf("expected Reject, got %v", action.Type)
+			}
+		})
+	}
+}
+
+// TestEvaluate_AggregateStillEnforced pins that MaxTokens keeps
+// working unchanged when only the aggregate limit is set — the
+// per-kind fields are optional additions, not a replacement.
+func TestEvaluate_AggregateStillEnforced(t *testing.T) {
+	p := newTestPlugin(500, 0, 0)
+	p.mu.Lock()
+	p.cache["sess-1"] = &counters{tokens: 500, inputTokens: 200, outputTokens: 300}
+	p.mu.Unlock()
+
+	action := p.OnRequest(context.Background(), makePctx("sess-1", 0))
+	if action.Type != pipeline.Reject {
+		t.Fatalf("aggregate limit: expected Reject, got %v", action.Type)
+	}
+}
+
+// TestHydrate_LegacyKeyHasNoSplitFields covers the migration
+// scenario: a Redis hash written before per-kind counters existed
+// only carries `tokens`, `calls`, `started_at`. Hydrating it must
+// succeed and set every per-kind counter to zero — no ParseInt
+// error leaks out, no field goes to a garbage value.
+func TestHydrate_LegacyKeyHasNoSplitFields(t *testing.T) {
+	fields := map[string]string{
+		"tokens":     "1234",
+		"calls":      "7",
+		"started_at": "1700000000",
+	}
+	c := parseCountersFromFields(fields)
+	if c.tokens != 1234 || c.calls != 7 {
+		t.Fatalf("aggregate parse: tokens=%d calls=%d", c.tokens, c.calls)
+	}
+	if c.inputTokens != 0 || c.cacheReadTokens != 0 || c.cacheWriteTokens != 0 ||
+		c.outputTokens != 0 || c.reasoningTokens != 0 {
+		t.Errorf("per-kind counters not zero on legacy key: %+v", c)
+	}
+	if c.startedAt.IsZero() {
+		t.Error("startedAt lost")
+	}
+}

From df43d5952f94067255f5507841362723cb65072d Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Mon, 31 Aug 2026 08:44:24 -0700
Subject: [PATCH 2/7] :art: Format for tests

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 authbridge/authlib/pipeline/session_test.go  | 6 +++---
 authbridge/authlib/pipeline/snapshot_test.go | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/authbridge/authlib/pipeline/session_test.go b/authbridge/authlib/pipeline/session_test.go
index 932ca3fb2..8097e7463 100644
--- a/authbridge/authlib/pipeline/session_test.go
+++ b/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"},
diff --git a/authbridge/authlib/pipeline/snapshot_test.go b/authbridge/authlib/pipeline/snapshot_test.go
index 60dae7145..739e5c768 100644
--- a/authbridge/authlib/pipeline/snapshot_test.go
+++ b/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)
 

From 41e3cd91f4d9c041cd1aa47c04f45d0ac700ab28 Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Mon, 31 Aug 2026 08:48:13 -0700
Subject: [PATCH 3/7] :loud_sound: Update logs with token details

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 authbridge/authlib/plugins/inferenceparser/plugin.go | 7 +++++++
 authbridge/authlib/plugins/sessionbudget/plugin.go   | 6 +++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index 936e77b75..f3ef66306 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -327,11 +327,18 @@ func getOrCreateStreamState(pctx *pipeline.Context) *inferenceStreamState {
 // once a response is finalized — shared by OnResponse and
 // OnResponseFrame so streaming and buffered finalize identically.
 func logInferenceFinalized(ext *pipeline.InferenceExtension) {
+	// Split counters log unconditionally: zero means the provider did
+	// not report that sub-kind, and stable field names help dashboards.
 	slog.Info("inference-parser: response",
 		"model", ext.Model,
 		"finishReason", ext.FinishReason,
 		"promptTokens", ext.PromptTokens,
 		"completionTokens", ext.CompletionTokens,
+		"inputTokens", ext.InputTokens,
+		"cacheReadTokens", ext.CacheReadTokens,
+		"cacheWriteTokens", ext.CacheWriteTokens,
+		"outputTokens", ext.OutputTokens,
+		"reasoningTokens", ext.ReasoningTokens,
 	)
 	slog.Debug("inference-parser: completion", "text", parsercommon.Truncate(ext.Completion, parsercommon.DebugBodyMax))
 }
diff --git a/authbridge/authlib/plugins/sessionbudget/plugin.go b/authbridge/authlib/plugins/sessionbudget/plugin.go
index 9df47f75e..57d4b1d87 100644
--- a/authbridge/authlib/plugins/sessionbudget/plugin.go
+++ b/authbridge/authlib/plugins/sessionbudget/plugin.go
@@ -600,9 +600,9 @@ func (p *SessionBudget) accumulate(sessionID string, delta tokenDelta) {
 	key := p.redisKey(sessionID)
 	ttl := time.Duration(p.cfg.SessionTTLSeconds) * time.Second
 
-	// One HashIncr per non-zero sub-kind. Zero deltas are skipped so old
-	// sessions whose per-kind counters were never written stay absent in
-	// Redis rather than accumulating a stream of no-op fields.
+	// One HashIncr per positive sub-kind. Non-positive deltas are skipped
+	// so per-kind fields stay absent on legacy sessions that never wrote
+	// them. Counters only grow, so <= 0 also swallows any stray negative.
 	for _, kv := range []struct {
 		field string
 		v     int64

From f20a8bf11aed038d82aa4cabc01167d47b1b788e Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Mon, 31 Aug 2026 11:57:26 -0700
Subject: [PATCH 4/7] :wrench::white_check_mark: Include presence

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 authbridge/authlib/pipeline/extensions.go     |  7 ++
 .../plugins/inferenceparser/anthropic.go      |  7 +-
 .../authlib/plugins/inferenceparser/plugin.go | 65 ++++++++++++-------
 .../inferenceparser/splittokens_test.go       | 55 ++++++++++++++++
 .../internal/parsercommon/tokenusage.go       | 34 +++++++---
 5 files changed, 133 insertions(+), 35 deletions(-)

diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go
index b7333e7c3..0892fa0f1 100644
--- a/authbridge/authlib/pipeline/extensions.go
+++ b/authbridge/authlib/pipeline/extensions.go
@@ -187,6 +187,13 @@ type InferenceExtension struct {
 	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"`
 }
diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go
index 5656c2c6a..353d89f57 100644
--- a/authbridge/authlib/plugins/inferenceparser/anthropic.go
+++ b/authbridge/authlib/plugins/inferenceparser/anthropic.go
@@ -136,15 +136,16 @@ func (u anthropicUsage) promptTotal() int {
 	return u.toNeutral().PromptTotal()
 }
 
-// toNeutral maps Anthropic's usage block onto the neutral TokenUsage.
-// Straight rename — Anthropic already splits input / cache-read /
-// cache-write on the wire.
+// toNeutral maps Anthropic's usage onto TokenUsage. The Messages API
+// unconditionally emits all four input/cache/output counters, so
+// Present names all four. Reasoning is not exposed by Anthropic.
 func (u anthropicUsage) toNeutral() parsercommon.TokenUsage {
 	return parsercommon.TokenUsage{
 		Input:      u.InputTokens,
 		CacheRead:  u.CacheReadInputTokens,
 		CacheWrite: u.CacheCreationInputTokens,
 		Output:     u.OutputTokens,
+		Present:    parsercommon.KindInput | parsercommon.KindCacheRead | parsercommon.KindCacheWrite | parsercommon.KindOutput,
 	}
 }
 
diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index f3ef66306..4b65a0ed4 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -323,22 +323,27 @@ func getOrCreateStreamState(pctx *pipeline.Context) *inferenceStreamState {
 	return s
 }
 
-// logInferenceFinalized emits the operator-facing INFO log + Observe
-// once a response is finalized — shared by OnResponse and
-// OnResponseFrame so streaming and buffered finalize identically.
+// logInferenceFinalized emits the operator-facing INFO log once a
+// response is finalized; shared by the buffered and streaming paths.
+// Split counters render -1 when ext.PresentKinds says the provider
+// did not expose that sub-kind, distinct from a reported 0.
 func logInferenceFinalized(ext *pipeline.InferenceExtension) {
-	// Split counters log unconditionally: zero means the provider did
-	// not report that sub-kind, and stable field names help dashboards.
+	tok := func(bit parsercommon.Kind, v int) int {
+		if ext.PresentKinds&uint8(bit) == 0 {
+			return -1
+		}
+		return v
+	}
 	slog.Info("inference-parser: response",
 		"model", ext.Model,
 		"finishReason", ext.FinishReason,
 		"promptTokens", ext.PromptTokens,
 		"completionTokens", ext.CompletionTokens,
-		"inputTokens", ext.InputTokens,
-		"cacheReadTokens", ext.CacheReadTokens,
-		"cacheWriteTokens", ext.CacheWriteTokens,
-		"outputTokens", ext.OutputTokens,
-		"reasoningTokens", ext.ReasoningTokens,
+		"inputTokens", tok(parsercommon.KindInput, ext.InputTokens),
+		"cacheReadTokens", tok(parsercommon.KindCacheRead, ext.CacheReadTokens),
+		"cacheWriteTokens", tok(parsercommon.KindCacheWrite, ext.CacheWriteTokens),
+		"outputTokens", tok(parsercommon.KindOutput, ext.OutputTokens),
+		"reasoningTokens", tok(parsercommon.KindReasoning, ext.ReasoningTokens),
 	)
 	slog.Debug("inference-parser: completion", "text", parsercommon.Truncate(ext.Completion, parsercommon.DebugBodyMax))
 }
@@ -447,35 +452,49 @@ type inferenceDelta struct {
 
 // inferenceUsage decodes the OpenAI usage block. See toNeutral for the
 // inclusive-prompt normalization.
+//
+// PromptTokensDetails and CompletionTokensDetails are pointers so the
+// parser can tell "block absent" (older API or non-reasoning model)
+// from "block present with cached_tokens/reasoning_tokens = 0."
 type inferenceUsage struct {
 	PromptTokens     int `json:"prompt_tokens"`
 	CompletionTokens int `json:"completion_tokens"`
 	TotalTokens      int `json:"total_tokens"`
 
-	PromptTokensDetails struct {
+	PromptTokensDetails *struct {
 		CachedTokens int `json:"cached_tokens"`
 	} `json:"prompt_tokens_details"`
-	CompletionTokensDetails struct {
+	CompletionTokensDetails *struct {
 		ReasoningTokens int `json:"reasoning_tokens"`
 	} `json:"completion_tokens_details"`
 }
 
 // toNeutral maps OpenAI's usage onto TokenUsage. prompt_tokens includes
 // cached_tokens on the wire — subtract to get uncached input and clamp
-// at 0 for malformed responses. CacheWrite stays 0 (OpenAI bills cache
-// writes as ordinary input).
+// at 0 for malformed responses. CacheWrite stays absent (OpenAI bills
+// cache writes as ordinary input). Input/Output are always present;
+// CacheRead and Reasoning are present only when their _details block
+// is on the wire (hence the pointer fields on inferenceUsage).
 func (u inferenceUsage) toNeutral() parsercommon.TokenUsage {
-	cached := u.PromptTokensDetails.CachedTokens
-	input := u.PromptTokens - cached
-	if input < 0 {
-		input = 0
+	usage := parsercommon.TokenUsage{
+		Input:   u.PromptTokens,
+		Output:  u.CompletionTokens,
+		Present: parsercommon.KindInput | parsercommon.KindOutput,
+	}
+	if u.PromptTokensDetails != nil {
+		cached := u.PromptTokensDetails.CachedTokens
+		usage.Input = u.PromptTokens - cached
+		if usage.Input < 0 {
+			usage.Input = 0
+		}
+		usage.CacheRead = cached
+		usage.Present |= parsercommon.KindCacheRead
 	}
-	return parsercommon.TokenUsage{
-		Input:     input,
-		CacheRead: cached,
-		Output:    u.CompletionTokens,
-		Reasoning: u.CompletionTokensDetails.ReasoningTokens,
+	if u.CompletionTokensDetails != nil {
+		usage.Reasoning = u.CompletionTokensDetails.ReasoningTokens
+		usage.Present |= parsercommon.KindReasoning
 	}
+	return usage
 }
 
 type inferenceRequest struct {
diff --git a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
index bd2acec3f..199f7a7e8 100644
--- a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
+++ b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/rossoctl/cortex/authbridge/authlib/pipeline"
+	"github.com/rossoctl/cortex/authbridge/authlib/plugins/internal/parsercommon"
 )
 
 // checkSplit surfaces every split field in one error message so a
@@ -163,3 +164,57 @@ func TestSplitTokens_OpenAISSE(t *testing.T) {
 		TotalTokens:      950,
 	})
 }
+
+// OpenAI response without the optional _details blocks: only Input
+// and Output are on the wire, so no cache-read or reasoning bit is set.
+func TestPresentKinds_OpenAI_NoDetailsBlocks(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "llama3.1"}
+	parseInferenceJSON([]byte(`{
+		"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],
+		"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}
+	}`), ext)
+
+	want := uint8(parsercommon.KindInput | parsercommon.KindOutput)
+	if ext.PresentKinds != want {
+		t.Errorf("PresentKinds = %b, want %b (Input|Output only)", ext.PresentKinds, want)
+	}
+}
+
+// OpenAI response WITH both _details blocks, both carrying zero:
+// CacheRead and Reasoning bits set — "reported zero," not "absent."
+// Distinguishability against the previous test is the whole point.
+func TestPresentKinds_OpenAI_WithDetailsBlocks(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	parseInferenceJSON([]byte(`{
+		"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],
+		"usage":{
+			"prompt_tokens":12,"completion_tokens":3,"total_tokens":15,
+			"prompt_tokens_details":{"cached_tokens":0},
+			"completion_tokens_details":{"reasoning_tokens":0}
+		}
+	}`), ext)
+
+	want := uint8(parsercommon.KindInput | parsercommon.KindOutput | parsercommon.KindCacheRead | parsercommon.KindReasoning)
+	if ext.PresentKinds != want {
+		t.Errorf("PresentKinds = %b, want %b (all four)", ext.PresentKinds, want)
+	}
+}
+
+// Anthropic response: Input/CacheRead/CacheWrite/Output always
+// present (Messages API emits all four); Reasoning never present.
+func TestPresentKinds_Anthropic(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-haiku-4-5"}
+	parseAnthropicJSON([]byte(`{
+		"content":[{"type":"text","text":"ok"}],
+		"stop_reason":"end_turn",
+		"usage":{
+			"input_tokens":8,"output_tokens":2,
+			"cache_creation_input_tokens":0,"cache_read_input_tokens":0
+		}
+	}`), ext)
+
+	want := uint8(parsercommon.KindInput | parsercommon.KindCacheRead | parsercommon.KindCacheWrite | parsercommon.KindOutput)
+	if ext.PresentKinds != want {
+		t.Errorf("PresentKinds = %b, want %b (four kinds, reasoning absent)", ext.PresentKinds, want)
+	}
+}
diff --git a/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
index 2852f29b0..f76939e70 100644
--- a/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
+++ b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
@@ -2,15 +2,29 @@ package parsercommon
 
 import "github.com/rossoctl/cortex/authbridge/authlib/pipeline"
 
-// TokenUsage is the provider-neutral token accounting shape. Provider
-// parsers normalize their wire format (e.g. OpenAI's inclusive
-// prompt_tokens) into these fields before publishing via Fill.
+// Kind is a bitmask naming which sub-kinds the provider populated.
+// A zero on a set bit is "reported zero"; a zero on an unset bit is
+// "not exposed."
+type Kind uint8
+
+const (
+	KindInput Kind = 1 << iota
+	KindCacheRead
+	KindCacheWrite
+	KindOutput
+	KindReasoning
+)
+
+// TokenUsage is the provider-neutral token accounting shape. Parsers
+// normalize their wire format into these fields before publishing via
+// Fill, and set Present to name which sub-kinds the wire carried.
 type TokenUsage struct {
-	Input      int // uncached prompt tokens
-	CacheRead  int // prompt tokens served from cache
-	CacheWrite int // prompt tokens written to cache
-	Output     int // generated completion tokens
-	Reasoning  int // reasoning-only output (subset of Output)
+	Input      int  // uncached prompt tokens
+	CacheRead  int  // prompt tokens served from cache
+	CacheWrite int  // prompt tokens written to cache
+	Output     int  // generated completion tokens
+	Reasoning  int  // reasoning-only output (subset of Output)
+	Present    Kind // which sub-kinds the provider reported
 }
 
 // PromptTotal is the sum of all prompt-side sub-kinds.
@@ -24,13 +38,15 @@ func (u TokenUsage) Total() int {
 	return u.PromptTotal() + u.Output
 }
 
-// Fill writes the split counters and derived legacy aggregates onto ext.
+// Fill writes the split counters, the Present bitmask, and the derived
+// legacy aggregates onto ext.
 func (u TokenUsage) Fill(ext *pipeline.InferenceExtension) {
 	ext.InputTokens = u.Input
 	ext.CacheReadTokens = u.CacheRead
 	ext.CacheWriteTokens = u.CacheWrite
 	ext.OutputTokens = u.Output
 	ext.ReasoningTokens = u.Reasoning
+	ext.PresentKinds = uint8(u.Present)
 
 	ext.PromptTokens = u.PromptTotal()
 	ext.CompletionTokens = u.Output

From 44db46f43017ef41c49d77b66a173225d638ee5f Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Mon, 31 Aug 2026 15:05:01 -0700
Subject: [PATCH 5/7] :goal_net::white_check_mark: Handle unconditional fill

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 .../authlib/plugins/inferenceparser/plugin.go       |  5 ++++-
 .../plugins/inferenceparser/splittokens_test.go     | 13 +++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index 4b65a0ed4..50e895c91 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -367,7 +367,10 @@ func parseInferenceJSON(body []byte, ext *pipeline.InferenceExtension) {
 			})
 		}
 	}
-	resp.Usage.toNeutral().Fill(ext)
+	// No usage block: leave PresentKinds at 0 (matches SSE path).
+	if resp.Usage.TotalTokens > 0 {
+		resp.Usage.toNeutral().Fill(ext)
+	}
 }
 
 // parseInferenceSSE concatenates content deltas across SSE events and captures
diff --git a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
index 199f7a7e8..83d8ae089 100644
--- a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
+++ b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
@@ -200,6 +200,19 @@ func TestPresentKinds_OpenAI_WithDetailsBlocks(t *testing.T) {
 	}
 }
 
+// No usage block on the wire: PresentKinds must stay 0 so the
+// log renders -1 (not exposed) rather than 0 (reported zero).
+func TestPresentKinds_OpenAI_NoUsageBlock(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	parseInferenceJSON([]byte(`{
+		"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]
+	}`), ext)
+
+	if ext.PresentKinds != 0 {
+		t.Errorf("PresentKinds = %b, want 0", ext.PresentKinds)
+	}
+}
+
 // Anthropic response: Input/CacheRead/CacheWrite/Output always
 // present (Messages API emits all four); Reasoning never present.
 func TestPresentKinds_Anthropic(t *testing.T) {

From 03b020d0350ec2d5a39cbadfcef4d903b4f6cb82 Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Tue, 1 Sep 2026 09:27:21 -0700
Subject: [PATCH 6/7] :goal_net::white_check_mark: Preserve presence

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 .../plugins/inferenceparser/anthropic.go      | 41 ++++++++----
 .../authlib/plugins/inferenceparser/plugin.go | 34 ++++++----
 .../inferenceparser/splittokens_test.go       | 67 ++++++++++++++++++-
 .../internal/parsercommon/tokenusage.go       | 23 ++++---
 .../internal/parsercommon/tokenusage_test.go  | 25 +++++++
 5 files changed, 153 insertions(+), 37 deletions(-)
 create mode 100644 authbridge/authlib/plugins/internal/parsercommon/tokenusage_test.go

diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go
index 353d89f57..c27298e00 100644
--- a/authbridge/authlib/plugins/inferenceparser/anthropic.go
+++ b/authbridge/authlib/plugins/inferenceparser/anthropic.go
@@ -125,28 +125,40 @@ 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.toNeutral().PromptTotal()
 }
 
-// toNeutral maps Anthropic's usage onto TokenUsage. The Messages API
-// unconditionally emits all four input/cache/output counters, so
-// Present names all four. Reasoning is not exposed by Anthropic.
+// 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 {
-	return parsercommon.TokenUsage{
-		Input:      u.InputTokens,
-		CacheRead:  u.CacheReadInputTokens,
-		CacheWrite: u.CacheCreationInputTokens,
-		Output:     u.OutputTokens,
-		Present:    parsercommon.KindInput | parsercommon.KindCacheRead | parsercommon.KindCacheWrite | parsercommon.KindOutput,
+	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 ---
@@ -340,6 +352,9 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline
 // 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
 	}
diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index 50e895c91..904d1a624 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -299,15 +299,11 @@ func foldOpenAIFrame(frame []byte, state *inferenceStreamState, ext *pipeline.In
 			ext.FinishReason = c.FinishReason
 		}
 	}
-	// Copy the three wire-backed counts field by field rather than assigning
-	// the whole struct. inferenceUsage doubles as the dialect-neutral
-	// accumulator and its two cache fields are json:"-", so they are always
-	// zero in a freshly decoded chunk — a whole-struct assignment would clear
-	// whatever the accumulator held. Nothing on the OpenAI path fills them
-	// today, which is precisely why that clobber would be silent when
-	// something does. TotalTokens is taken off the wire rather than recomputed:
-	// the provider reports it, and it may legitimately differ from
-	// prompt+completion.
+	// OpenAI streams cumulative usage: each usage-bearing chunk restates
+	// the full totals, so replacing state.usage with the latest chunk's
+	// neutral form is correct. Gate on TotalTokens > 0 so chunks with no
+	// usage block (every non-final chunk) don't clear an accumulator that
+	// a prior chunk populated.
 	if chunk.Usage.TotalTokens > 0 {
 		state.usage = chunk.Usage.toNeutral()
 		state.hasUsage = true
@@ -376,8 +372,15 @@ func parseInferenceJSON(body []byte, ext *pipeline.InferenceExtension) {
 // parseInferenceSSE concatenates content deltas across SSE events and captures
 // the last finish_reason and usage block (sent when stream_options.include_usage
 // is set). The stream terminates with a "data: [DONE]" marker which is skipped.
+//
+// OpenAI streams cumulative usage: each usage-bearing chunk restates the full
+// totals, so the latest chunk's neutral form is authoritative. Accumulate into
+// a local TokenUsage and Fill once at the end — matching foldOpenAIFrame's
+// contract, so PresentKinds and ReportedTotal reflect only the final chunk.
 func parseInferenceSSE(body []byte, ext *pipeline.InferenceExtension) {
 	var completion strings.Builder
+	var usage parsercommon.TokenUsage
+	var hasUsage bool
 	for _, line := range bytes.Split(body, []byte("\n")) {
 		line = bytes.TrimSpace(line)
 		if !bytes.HasPrefix(line, []byte("data:")) {
@@ -401,10 +404,14 @@ func parseInferenceSSE(body []byte, ext *pipeline.InferenceExtension) {
 			}
 		}
 		if chunk.Usage.TotalTokens > 0 {
-			chunk.Usage.toNeutral().Fill(ext)
+			usage = chunk.Usage.toNeutral()
+			hasUsage = true
 		}
 	}
 	ext.Completion = completion.String()
+	if hasUsage {
+		usage.Fill(ext)
+	}
 }
 
 type inferenceResponse struct {
@@ -480,9 +487,10 @@ type inferenceUsage struct {
 // is on the wire (hence the pointer fields on inferenceUsage).
 func (u inferenceUsage) toNeutral() parsercommon.TokenUsage {
 	usage := parsercommon.TokenUsage{
-		Input:   u.PromptTokens,
-		Output:  u.CompletionTokens,
-		Present: parsercommon.KindInput | parsercommon.KindOutput,
+		Input:         u.PromptTokens,
+		Output:        u.CompletionTokens,
+		ReportedTotal: u.TotalTokens,
+		Present:       parsercommon.KindInput | parsercommon.KindOutput,
 	}
 	if u.PromptTokensDetails != nil {
 		cached := u.PromptTokensDetails.CachedTokens
diff --git a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
index 83d8ae089..43e204dd1 100644
--- a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
+++ b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
@@ -119,6 +119,20 @@ func TestSplitTokens_OpenAIJSON(t *testing.T) {
 	})
 }
 
+// Gateway reports only total_tokens (prompt/completion zero): the wire
+// total must be preserved rather than recomputed as 0 from the sub-fields.
+func TestSplitTokens_OpenAIJSON_PreservesReportedTotal(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	parseInferenceJSON([]byte(`{
+		"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],
+		"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":950}
+	}`), ext)
+
+	if ext.TotalTokens != 950 {
+		t.Errorf("TotalTokens = %d, want 950 (reported total preserved)", ext.TotalTokens)
+	}
+}
+
 // Malformed OpenAI response where cached_tokens > prompt_tokens must
 // clamp InputTokens to 0, not go negative.
 func TestSplitTokens_OpenAIJSON_ClampNegativeInput(t *testing.T) {
@@ -180,9 +194,8 @@ func TestPresentKinds_OpenAI_NoDetailsBlocks(t *testing.T) {
 	}
 }
 
-// OpenAI response WITH both _details blocks, both carrying zero:
-// CacheRead and Reasoning bits set — "reported zero," not "absent."
-// Distinguishability against the previous test is the whole point.
+// _details blocks present but zero: bits set ("reported zero"), not
+// absent — the distinction from the previous test.
 func TestPresentKinds_OpenAI_WithDetailsBlocks(t *testing.T) {
 	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
 	parseInferenceJSON([]byte(`{
@@ -231,3 +244,51 @@ func TestPresentKinds_Anthropic(t *testing.T) {
 		t.Errorf("PresentKinds = %b, want %b (four kinds, reasoning absent)", ext.PresentKinds, want)
 	}
 }
+
+// Anthropic non-streaming response WITHOUT the cache_* fields on the wire:
+// only Input and Output bits should be set.
+func TestPresentKinds_Anthropic_NoCacheFields(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-haiku-4-5"}
+	parseAnthropicJSON([]byte(`{
+		"content":[{"type":"text","text":"ok"}],
+		"stop_reason":"end_turn",
+		"usage":{"input_tokens":8,"output_tokens":2}
+	}`), ext)
+
+	want := uint8(parsercommon.KindInput | parsercommon.KindOutput)
+	if ext.PresentKinds != want {
+		t.Errorf("PresentKinds = %b, want %b (Input|Output only)", ext.PresentKinds, want)
+	}
+}
+
+// Anthropic ?beta=true SSE interrupted after message_start: no cache fields
+// ever land on the wire, so their bits must stay cleared.
+func TestPresentKinds_AnthropicSSE_Beta_MessageStartOnly(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "claude-opus-4-8"}
+	body := []byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":9,\"output_tokens\":0}}}\n")
+	parseAnthropicSSE(body, ext)
+
+	want := uint8(parsercommon.KindInput | parsercommon.KindOutput)
+	if ext.PresentKinds != want {
+		t.Errorf("PresentKinds = %b, want %b (no cache fields observed)", ext.PresentKinds, want)
+	}
+}
+
+// OpenAI streaming: usage is cumulative, so state.usage is replaced on
+// each usage-bearing chunk. A later chunk that omits _details must
+// therefore drop CacheRead — this pins that behavior against a future
+// refactor that turns usage merging back into a merge.
+func TestFoldOpenAIFrame_UsageIsCumulative(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o", Stream: true}
+	body := []byte("data: {\"choices\":[{\"delta\":{},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":800,\"completion_tokens\":100,\"total_tokens\":900,\"prompt_tokens_details\":{\"cached_tokens\":600}}}\n" +
+		"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":800,\"completion_tokens\":150,\"total_tokens\":950}}\n" +
+		"data: [DONE]\n")
+	parseInferenceSSE(body, ext)
+
+	if ext.CacheReadTokens != 0 {
+		t.Errorf("CacheReadTokens = %d, want 0 (later cumulative chunk without _details)", ext.CacheReadTokens)
+	}
+	if ext.OutputTokens != 150 {
+		t.Errorf("OutputTokens = %d, want 150 (last chunk wins)", ext.OutputTokens)
+	}
+}
diff --git a/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
index f76939e70..1a1be9737 100644
--- a/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
+++ b/authbridge/authlib/plugins/internal/parsercommon/tokenusage.go
@@ -19,12 +19,13 @@ const (
 // normalize their wire format into these fields before publishing via
 // Fill, and set Present to name which sub-kinds the wire carried.
 type TokenUsage struct {
-	Input      int  // uncached prompt tokens
-	CacheRead  int  // prompt tokens served from cache
-	CacheWrite int  // prompt tokens written to cache
-	Output     int  // generated completion tokens
-	Reasoning  int  // reasoning-only output (subset of Output)
-	Present    Kind // which sub-kinds the provider reported
+	Input         int  // uncached prompt tokens
+	CacheRead     int  // prompt tokens served from cache
+	CacheWrite    int  // prompt tokens written to cache
+	Output        int  // generated completion tokens
+	Reasoning     int  // reasoning-only output (subset of Output)
+	ReportedTotal int  // provider's own total_tokens if reported, else 0
+	Present       Kind // which sub-kinds the provider reported
 }
 
 // PromptTotal is the sum of all prompt-side sub-kinds.
@@ -39,7 +40,9 @@ func (u TokenUsage) Total() int {
 }
 
 // Fill writes the split counters, the Present bitmask, and the derived
-// legacy aggregates onto ext.
+// legacy aggregates onto ext. TotalTokens prefers the provider's own
+// reported total when present — a gateway that reports only total_tokens
+// (with prompt/completion zero) would otherwise record 0 here.
 func (u TokenUsage) Fill(ext *pipeline.InferenceExtension) {
 	ext.InputTokens = u.Input
 	ext.CacheReadTokens = u.CacheRead
@@ -50,5 +53,9 @@ func (u TokenUsage) Fill(ext *pipeline.InferenceExtension) {
 
 	ext.PromptTokens = u.PromptTotal()
 	ext.CompletionTokens = u.Output
-	ext.TotalTokens = u.Total()
+	if u.ReportedTotal > 0 {
+		ext.TotalTokens = u.ReportedTotal
+	} else {
+		ext.TotalTokens = u.Total()
+	}
 }
diff --git a/authbridge/authlib/plugins/internal/parsercommon/tokenusage_test.go b/authbridge/authlib/plugins/internal/parsercommon/tokenusage_test.go
new file mode 100644
index 000000000..cc8600411
--- /dev/null
+++ b/authbridge/authlib/plugins/internal/parsercommon/tokenusage_test.go
@@ -0,0 +1,25 @@
+package parsercommon
+
+import "testing"
+
+// Kind values are serialized into session events (via ext.PresentKinds)
+// and consumed by the log renderer's -1 sentinel logic, so their numeric
+// layout is a wire contract, not just an internal iota.
+func TestKindBitLayout(t *testing.T) {
+	cases := []struct {
+		name string
+		got  Kind
+		want Kind
+	}{
+		{"KindInput", KindInput, 1},
+		{"KindCacheRead", KindCacheRead, 2},
+		{"KindCacheWrite", KindCacheWrite, 4},
+		{"KindOutput", KindOutput, 8},
+		{"KindReasoning", KindReasoning, 16},
+	}
+	for _, c := range cases {
+		if c.got != c.want {
+			t.Errorf("%s = %d, want %d", c.name, c.got, c.want)
+		}
+	}
+}

From 33b65c00535f712741f9159ed3b29938b0be2e22 Mon Sep 17 00:00:00 2001
From: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Date: Tue, 1 Sep 2026 09:48:27 -0700
Subject: [PATCH 7/7] :goal_net: Update OpenAI usage response parsing

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
---
 .../plugins/inferenceparser/anthropic.go      |  4 --
 .../authlib/plugins/inferenceparser/plugin.go | 59 +++++++++++--------
 .../inferenceparser/splittokens_test.go       | 21 +++++++
 3 files changed, 56 insertions(+), 28 deletions(-)

diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go
index c27298e00..ce0adeb15 100644
--- a/authbridge/authlib/plugins/inferenceparser/anthropic.go
+++ b/authbridge/authlib/plugins/inferenceparser/anthropic.go
@@ -136,10 +136,6 @@ type anthropicUsage struct {
 	CacheReadInputTokens     *int `json:"cache_read_input_tokens"`
 }
 
-func (u anthropicUsage) promptTotal() int {
-	return u.toNeutral().PromptTotal()
-}
-
 // 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
diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go
index 904d1a624..3643edb54 100644
--- a/authbridge/authlib/plugins/inferenceparser/plugin.go
+++ b/authbridge/authlib/plugins/inferenceparser/plugin.go
@@ -301,10 +301,10 @@ func foldOpenAIFrame(frame []byte, state *inferenceStreamState, ext *pipeline.In
 	}
 	// OpenAI streams cumulative usage: each usage-bearing chunk restates
 	// the full totals, so replacing state.usage with the latest chunk's
-	// neutral form is correct. Gate on TotalTokens > 0 so chunks with no
-	// usage block (every non-final chunk) don't clear an accumulator that
-	// a prior chunk populated.
-	if chunk.Usage.TotalTokens > 0 {
+	// neutral form is correct. Gate on hasAny so chunks with no usage
+	// block (every non-final chunk) don't clear an accumulator that a
+	// prior chunk populated.
+	if chunk.Usage.hasAny() {
 		state.usage = chunk.Usage.toNeutral()
 		state.hasUsage = true
 	}
@@ -364,7 +364,7 @@ func parseInferenceJSON(body []byte, ext *pipeline.InferenceExtension) {
 		}
 	}
 	// No usage block: leave PresentKinds at 0 (matches SSE path).
-	if resp.Usage.TotalTokens > 0 {
+	if resp.Usage.hasAny() {
 		resp.Usage.toNeutral().Fill(ext)
 	}
 }
@@ -403,7 +403,7 @@ func parseInferenceSSE(body []byte, ext *pipeline.InferenceExtension) {
 				ext.FinishReason = c.FinishReason
 			}
 		}
-		if chunk.Usage.TotalTokens > 0 {
+		if chunk.Usage.hasAny() {
 			usage = chunk.Usage.toNeutral()
 			hasUsage = true
 		}
@@ -460,16 +460,13 @@ type inferenceDelta struct {
 	Content string `json:"content"`
 }
 
-// inferenceUsage decodes the OpenAI usage block. See toNeutral for the
-// inclusive-prompt normalization.
-//
-// PromptTokensDetails and CompletionTokensDetails are pointers so the
-// parser can tell "block absent" (older API or non-reasoning model)
-// from "block present with cached_tokens/reasoning_tokens = 0."
+// inferenceUsage decodes the OpenAI usage block. All fields are pointers
+// so "key absent" is distinguishable from "key present with value 0" —
+// a total-only response must not assert KindInput/KindOutput.
 type inferenceUsage struct {
-	PromptTokens     int `json:"prompt_tokens"`
-	CompletionTokens int `json:"completion_tokens"`
-	TotalTokens      int `json:"total_tokens"`
+	PromptTokens     *int `json:"prompt_tokens"`
+	CompletionTokens *int `json:"completion_tokens"`
+	TotalTokens      *int `json:"total_tokens"`
 
 	PromptTokensDetails *struct {
 		CachedTokens int `json:"cached_tokens"`
@@ -479,22 +476,36 @@ type inferenceUsage struct {
 	} `json:"completion_tokens_details"`
 }
 
+// hasAny reports whether any recognized field was on the wire. Call
+// sites gate toNeutral on this so an omitted or empty usage block
+// doesn't overwrite a prior chunk's state.
+func (u inferenceUsage) hasAny() bool {
+	return u.PromptTokens != nil || u.CompletionTokens != nil || u.TotalTokens != nil ||
+		u.PromptTokensDetails != nil || u.CompletionTokensDetails != nil
+}
+
 // toNeutral maps OpenAI's usage onto TokenUsage. prompt_tokens includes
 // cached_tokens on the wire — subtract to get uncached input and clamp
 // at 0 for malformed responses. CacheWrite stays absent (OpenAI bills
-// cache writes as ordinary input). Input/Output are always present;
-// CacheRead and Reasoning are present only when their _details block
-// is on the wire (hence the pointer fields on inferenceUsage).
+// cache writes as ordinary input). Each Present bit is set only when
+// its key was on the wire, so an absent key stays "not exposed"
+// (-1 sentinel) rather than "reported zero."
 func (u inferenceUsage) toNeutral() parsercommon.TokenUsage {
-	usage := parsercommon.TokenUsage{
-		Input:         u.PromptTokens,
-		Output:        u.CompletionTokens,
-		ReportedTotal: u.TotalTokens,
-		Present:       parsercommon.KindInput | parsercommon.KindOutput,
+	usage := parsercommon.TokenUsage{}
+	if u.TotalTokens != nil {
+		usage.ReportedTotal = *u.TotalTokens
+	}
+	if u.PromptTokens != nil {
+		usage.Input = *u.PromptTokens
+		usage.Present |= parsercommon.KindInput
+	}
+	if u.CompletionTokens != nil {
+		usage.Output = *u.CompletionTokens
+		usage.Present |= parsercommon.KindOutput
 	}
 	if u.PromptTokensDetails != nil {
 		cached := u.PromptTokensDetails.CachedTokens
-		usage.Input = u.PromptTokens - cached
+		usage.Input -= cached
 		if usage.Input < 0 {
 			usage.Input = 0
 		}
diff --git a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
index 43e204dd1..adae98489 100644
--- a/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
+++ b/authbridge/authlib/plugins/inferenceparser/splittokens_test.go
@@ -133,6 +133,27 @@ func TestSplitTokens_OpenAIJSON_PreservesReportedTotal(t *testing.T) {
 	}
 }
 
+// prompt_tokens/completion_tokens keys absent (not zero): both bits
+// must stay cleared. Contrast with PreservesReportedTotal, where the
+// keys are present with value 0 and the bits stay set.
+func TestPresentKinds_OpenAI_TotalOnly(t *testing.T) {
+	ext := &pipeline.InferenceExtension{Model: "gpt-4o"}
+	parseInferenceJSON([]byte(`{
+		"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],
+		"usage":{"total_tokens":950}
+	}`), ext)
+
+	if ext.PresentKinds&uint8(parsercommon.KindInput) != 0 {
+		t.Errorf("KindInput set, want cleared (prompt_tokens absent from wire)")
+	}
+	if ext.PresentKinds&uint8(parsercommon.KindOutput) != 0 {
+		t.Errorf("KindOutput set, want cleared (completion_tokens absent from wire)")
+	}
+	if ext.TotalTokens != 950 {
+		t.Errorf("TotalTokens = %d, want 950", ext.TotalTokens)
+	}
+}
+
 // Malformed OpenAI response where cached_tokens > prompt_tokens must
 // clamp InputTokens to 0, not go negative.
 func TestSplitTokens_OpenAIJSON_ClampNegativeInput(t *testing.T) {