diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go index 18937b274..d11940531 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic.go +++ b/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 + } } } } diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go index f6019977f..04cc36d6c 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go +++ b/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) + } +} diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go index 792280fa9..4d7696b31 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin.go +++ b/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)