Fix: Read prompt tokens from message_delta usage on the beta Messages path

authbridge/authlib/plugins/inferenceparser/anthropic.go

@@ -177,9 +177,10 @@ func parseAnthropicJSON(body []byte, ext *pipeline.InferenceExtension) {
 
 // anthropicStreamEvent is one SSE event's data payload. The Messages stream is
 // a sequence of typed events (vs OpenAI's uniform chat.completion.chunk):
-// message_start (carries usage.input_tokens), content_block_delta (text_delta /
+// message_start (carries usage — but see below), content_block_delta (text_delta /
 // input_json_delta / thinking_delta), message_delta (delta.stop_reason +
-// cumulative usage.output_tokens), message_stop, plus ping/content_block_*.
+// cumulative usage.output_tokens, and on the ?beta=true path the prompt-cache
+// counts too), message_stop, plus ping/content_block_*.
 type anthropicStreamEvent struct {
 	Type    string `json:"type"`
 	Message *struct {
@@ -194,7 +195,9 @@ type anthropicStreamEvent struct {
 }
 
 // foldAnthropicFrame folds one Messages SSE event into the running stream state.
-// input_tokens come from message_start; the completion accumulates from
+// 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,
+// message_delta on the ?beta=true path. The completion accumulates from
 // text_delta blocks; stop_reason and the cumulative output_tokens arrive in
 // message_delta. Unknown events (ping, content_block_start/stop, message_stop)
 // are ignored.
@@ -216,12 +219,27 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline
 		if ev.Delta != nil && ev.Delta.StopReason != "" {
 			ext.FinishReason = ev.Delta.StopReason
 		}
-		if ev.Usage != nil && ev.Usage.OutputTokens > 0 {
-			// usage.output_tokens in message_delta is cumulative — take the
-			// latest. TotalTokens must be non-zero for the shared finalize
-			// block to copy the counts onto the extension.
-			state.usage.CompletionTokens = ev.Usage.OutputTokens
-			state.usage.TotalTokens = state.usage.PromptTokens + ev.Usage.OutputTokens
+		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
+			}
+			if ev.Usage.OutputTokens > 0 {
+				// usage.output_tokens in message_delta is cumulative — take the
+				// latest. TotalTokens must be non-zero for the shared finalize
+				// block to copy the counts onto the extension.
+				state.usage.CompletionTokens = ev.Usage.OutputTokens
+				state.usage.TotalTokens = state.usage.PromptTokens + ev.Usage.OutputTokens
+			}
 		}
 	}
 }

authbridge/authlib/plugins/inferenceparser/anthropic_test.go

@@ -154,3 +154,101 @@ func TestInferenceParser_AnthropicMessages_StreamFoldsEvents(t *testing.T) {
 			ext.PromptTokens, ext.CompletionTokens, ext.TotalTokens)
 	}
 }
+
+// TestInferenceParser_AnthropicMessages_StreamBetaPathUsage covers the ?beta=true
+// Messages path, where the prompt-cache counts arrive in message_delta instead of
+// message_start. The frames below are the usage blocks captured verbatim from a
+// real Claude Code turn (anthropic-beta: claude-code-20250219) against an
+// Anthropic-compatible gateway: message_start carried only input_tokens, and the
+// 33,763 cached tokens appeared two events later. Reading the prompt size from
+// message_start alone recorded that turn as 9 tokens instead of 33,772.
+func TestInferenceParser_AnthropicMessages_StreamBetaPathUsage(t *testing.T) {
+	p := NewInferenceParser()
+	// Query-free Path: the HTTP listeners (forwardproxy, reverseproxy) populate
+	// Context.Path from r.URL.Path, so /v1/messages?beta=true arrives here as
+	// /v1/messages. extproc does NOT — it uses the :path pseudo-header, query
+	// included; TestInferenceParser_AnthropicMessages_QueryStringPath covers that.
+	pctx := &pipeline.Context{Path: "/v1/messages"}
+	pctx.Extensions.Inference = &pipeline.InferenceExtension{Model: "claude-haiku-4-5", Stream: true, IsAction: true}
+
+	frames := [][]byte{
+		[]byte(`{"type":"message_start","message":{"id":"msg_bdrk_1","type":"message","role":"assistant","usage":{"input_tokens":9,"output_tokens":0}}}`),
+		[]byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Done"}}`),
+		[]byte(`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":9,"output_tokens":399,"cache_creation_input_tokens":3755,"cache_read_input_tokens":30008}}`),
+		[]byte(`{"type":"message_stop"}`),
+	}
+	for _, f := range frames {
+		p.OnResponseFrame(context.Background(), pctx, f, false)
+	}
+	p.OnResponseFrame(context.Background(), pctx, nil, true)
+
+	ext := pctx.Extensions.Inference
+	// 9 + 3755 + 30008 — the cached context is still billed input.
+	if ext.PromptTokens != 33772 {
+		t.Errorf("PromptTokens = %d, want 33772 (message_delta usage ignored?)", ext.PromptTokens)
+	}
+	if ext.CompletionTokens != 399 || ext.TotalTokens != 34171 {
+		t.Errorf("tokens = completion %d / total %d, want 399/34171",
+			ext.CompletionTokens, ext.TotalTokens)
+	}
+	if ext.FinishReason != "end_turn" {
+		t.Errorf("FinishReason = %q, want end_turn", ext.FinishReason)
+	}
+}
+
+// TestInferenceParser_AnthropicMessages_QueryStringPath pins dialect dispatch
+// when Context.Path carries a query string. extproc populates Path from the
+// HTTP/2 :path pseudo-header, which includes the query, so Claude Code's
+// POST /v1/messages?beta=true arrives here as "/v1/messages?beta=true" — while
+// the HTTP listeners strip it via r.URL.Path.
+//
+// Two distinct failure modes are covered, both previously silent:
+//
+//   - OnRequest's exact-match switch fell to default, leaving
+//     Extensions.Inference nil so the whole exchange went unrecorded;
+//   - had dispatch matched but the four dialect-selection sites not been
+//     normalised, an Anthropic stream would have been folded by the OpenAI
+//     handler, which does not understand message_delta and would report zero
+//     tokens rather than fail.
+//
+// Asserting the token counts therefore checks the routing, not just the match.
+func TestInferenceParser_AnthropicMessages_QueryStringPath(t *testing.T) {
+	p := NewInferenceParser()
+	pctx := &pipeline.Context{
+		Path: "/v1/messages?beta=true",
+		Body: []byte(`{"model":"claude-haiku-4-5","messages":[{"role":"user","content":"hi"}],"stream":true}`),
+	}
+
+	if action := p.OnRequest(context.Background(), pctx); action.Type != pipeline.Continue {
+		t.Fatalf("expected Continue, got %v", action.Type)
+	}
+	ext := pctx.Extensions.Inference
+	if ext == nil {
+		t.Fatal("Extensions.Inference is nil — query string defeated the dispatch switch")
+	}
+	if ext.Model != "claude-haiku-4-5" {
+		t.Errorf("Model = %q, want claude-haiku-4-5", ext.Model)
+	}
+
+	frames := [][]byte{
+		[]byte(`{"type":"message_start","message":{"id":"msg_q1","type":"message","role":"assistant","usage":{"input_tokens":11,"output_tokens":0}}}`),
+		[]byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}`),
+		[]byte(`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":11,"output_tokens":7,"cache_read_input_tokens":500}}`),
+	}
+	for _, f := range frames {
+		p.OnResponseFrame(context.Background(), pctx, f, false)
+	}
+	p.OnResponseFrame(context.Background(), pctx, nil, true)
+
+	// 11 + 500 cached input; the OpenAI folder would leave these at 0.
+	if ext.PromptTokens != 511 || ext.CompletionTokens != 7 || ext.TotalTokens != 518 {
+		t.Errorf("tokens = prompt %d / completion %d / total %d, want 511/7/518 (wrong dialect?)",
+			ext.PromptTokens, ext.CompletionTokens, ext.TotalTokens)
+	}
+	if ext.Completion != "ok" {
+		t.Errorf("Completion = %q, want \"ok\"", ext.Completion)
+	}
+	if ext.FinishReason != "end_turn" {
+		t.Errorf("FinishReason = %q, want end_turn", ext.FinishReason)
+	}
+}

authbridge/authlib/plugins/inferenceparser/plugin.go

@@ -31,13 +31,31 @@ func (p *InferenceParser) Capabilities() pipeline.PluginCapabilities {
 	}
 }
 
+// endpointPath returns pctx.Path with any query string removed.
+//
+// The listeners disagree on what Path holds, and dialect dispatch below is
+// exact-match, so this has to be normalised in one place. The HTTP listeners
+// set Path from r.URL.Path, which already excludes the query; extproc sets it
+// from the HTTP/2 :path pseudo-header, which per RFC 9113 §8.3.1 includes it.
+//
+// Claude Code posts to /v1/messages?beta=true, so without this the request
+// falls to the default arm on the envoy-sidecar path and the parser records no
+// inference telemetry at all — and once OnRequest did match, the four
+// dialect-selection sites below would send an Anthropic stream to the OpenAI
+// parser. Both failure modes are silent, which is why every site normalises
+// rather than only the dispatch switch.
+func endpointPath(pctx *pipeline.Context) string {
+	path, _, _ := strings.Cut(pctx.Path, "?")
+	return path
+}
+
 func (p *InferenceParser) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action {
 	// Dispatch by endpoint dialect: OpenAI chat/completions vs Anthropic
 	// Messages. No Invocation is recorded when the parser doesn't apply
 	// (unrecognized path, empty body, or non-JSON body) — operators infer
 	// "inference-parser is in this pipeline" from config, not per-event rows.
 	var ext *pipeline.InferenceExtension
-	switch pctx.Path {
+	switch endpointPath(pctx) {
 	case "/v1/chat/completions", "/v1/completions":
 		ext = parseOpenAIRequest(pctx.Body)
 	case anthropicMessagesPath:
@@ -124,13 +142,13 @@ func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context)
 	}
 
 	if ext.Stream {
-		if pctx.Path == anthropicMessagesPath {
+		if endpointPath(pctx) == anthropicMessagesPath {
 			parseAnthropicSSE(pctx.ResponseBody, ext)
 		} else {
 			parseInferenceSSE(pctx.ResponseBody, ext)
 		}
 	} else {
-		if pctx.Path == anthropicMessagesPath {
+		if endpointPath(pctx) == anthropicMessagesPath {
 			parseAnthropicJSON(pctx.ResponseBody, ext)
 		} else {
 			parseInferenceJSON(pctx.ResponseBody, ext)
@@ -179,7 +197,7 @@ func (p *InferenceParser) OnResponseFrame(_ context.Context, pctx *pipeline.Cont
 			pctx.Skip("no_response_body")
 			return pipeline.Action{Type: pipeline.Continue}
 		}
-		if pctx.Path == anthropicMessagesPath {
+		if endpointPath(pctx) == anthropicMessagesPath {
 			parseAnthropicJSON(frame, ext)
 		} else {
 			parseInferenceJSON(frame, ext)
@@ -194,7 +212,7 @@ func (p *InferenceParser) OnResponseFrame(_ context.Context, pctx *pipeline.Cont
 	state := getOrCreateStreamState(pctx)
 
 	if len(frame) > 0 {
-		if pctx.Path == anthropicMessagesPath {
+		if endpointPath(pctx) == anthropicMessagesPath {
 			foldAnthropicFrame(frame, state, ext)
 		} else {
 			foldOpenAIFrame(frame, state, ext)