From 045a780be9b01c4e7c59c690cc86c24953cc93b5 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 26 Aug 2026 15:44:01 -0400 Subject: [PATCH 1/4] feat: Record per-tool-call cost inputs on the Anthropic path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The inference-parser records enough to see that a turn happened, but not enough to say what it cost. Three gaps, all on data already on the wire: Streamed tool calls were dropped. A tool call arrives across three event types — id and name on content_block_start, arguments as input_json_delta fragments that are only valid JSON once concatenated — and the stream folder modeled none of them. A streaming turn recorded finishReason "tool_use" with an empty toolCalls list while the equivalent non-streaming response recorded the call in full. Fragments are routed by content block index, so interleaved calls don't merge into each other. Non-text request messages read as free. Content keeps only text blocks, so a tool_result — a whole file, in an agent loop — flattens to "" while the model was billed for every byte. InferenceMessage.ContentBytes is the wire size of the content value before that reduction. It is a byte count, not a token count: a size signal, not an exact one. Prompt-cache writes and reads were collapsed into PromptTokens. A provider that prices caching charges a premium to write an entry and a steep discount to read one — 12.5x apart for Anthropic — so two turns with identical PromptTokens can differ by an order of magnitude in cost. CacheWriteTokens and CacheReadTokens split it, from the same usage block the totals already come from. Scope: OpenAI *streaming* tool calls remain uncaptured — a separate pre-existing gap. Recording streamed tool-call Arguments is parity with the non-streaming path, which already records them. Depends on the message_delta usage fix (#811) — same hunk. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/extensions.go | 27 ++++ .../plugins/inferenceparser/anthropic.go | 120 ++++++++++++-- .../plugins/inferenceparser/anthropic_test.go | 151 ++++++++++++++++++ .../authlib/plugins/inferenceparser/plugin.go | 84 ++++++++-- .../plugins/inferenceparser/plugin_test.go | 38 +++++ authbridge/cmd/abctl/tui/detail_pane.go | 1 + 6 files changed, 397 insertions(+), 24 deletions(-) diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go index 221c056ad..55104ebf3 100644 --- a/authbridge/authlib/pipeline/extensions.go +++ b/authbridge/authlib/pipeline/extensions.go @@ -164,6 +164,23 @@ type InferenceExtension struct { TotalTokens int `json:"totalTokens,omitempty"` 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"` + // Classification — see MCPExtension.IsAction. IsAction bool `json:"isAction,omitempty"` } @@ -172,6 +189,16 @@ type InferenceExtension struct { type InferenceMessage struct { Role string `json:"role"` Content string `json:"content,omitempty"` + + // ContentBytes is the wire size of this message's content value as the + // client sent it, before the parser reduced it to text. Content keeps + // only text blocks, so a message whose payload is a tool result, an + // image, or any other non-text block flattens to "" and looks free — + // while the model was billed for all of it. ContentBytes is what those + // messages contribute, without recording their contents: it is a byte + // count of the raw JSON (syntax and escapes included), not a token + // count, and is a size signal rather than an exact one. + ContentBytes int `json:"contentBytes,omitempty"` } // InferenceTool is a function/tool the client declared the model may call. diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go index d11940531..921426031 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic.go +++ b/authbridge/authlib/plugins/inferenceparser/anthropic.go @@ -43,9 +43,16 @@ type anthropicTool struct { // content is a string or an array of content blocks (text / image / tool_use / // tool_result); reuse flattenContent, which keeps text blocks and drops the // rest — the same {"type":"text","text":...} shape OpenAI uses. +// +// ContentBytes records the size of what was there before that reduction, so +// the blocks flattenContent drops are still accounted for. In an agent loop +// the dropped blocks are the bulk of the conversation: every tool result comes +// back as a tool_result block, so a turn that reads a large file shows up as an +// empty Content and would otherwise look free. type anthropicReqMessage struct { - Role string - Content string + Role string + Content string + ContentBytes int } func (m *anthropicReqMessage) UnmarshalJSON(data []byte) error { @@ -58,6 +65,7 @@ func (m *anthropicReqMessage) UnmarshalJSON(data []byte) error { } m.Role = raw.Role m.Content = flattenContent(raw.Content) + m.ContentBytes = contentBytes(raw.Content) return nil } @@ -89,10 +97,14 @@ func parseAnthropicRequest(body []byte) *pipeline.InferenceExtension { // Surface it as a leading system message so downstream policy plugins // (IBAC, etc.) see it the same way they see OpenAI's system message. if sys := flattenContent(req.System); sys != "" { - ext.Messages = append(ext.Messages, pipeline.InferenceMessage{Role: "system", Content: sys}) + ext.Messages = append(ext.Messages, pipeline.InferenceMessage{ + Role: "system", Content: sys, ContentBytes: contentBytes(req.System), + }) } for _, msg := range req.Messages { - ext.Messages = append(ext.Messages, pipeline.InferenceMessage{Role: msg.Role, Content: msg.Content}) + ext.Messages = append(ext.Messages, pipeline.InferenceMessage{ + Role: msg.Role, Content: msg.Content, ContentBytes: msg.ContentBytes, + }) } for _, tool := range req.Tools { if tool.Name == "" { @@ -171,6 +183,8 @@ func parseAnthropicJSON(body []byte, ext *pipeline.InferenceExtension) { 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 } // --- streaming --- @@ -190,17 +204,78 @@ type anthropicStreamEvent struct { Type string `json:"type"` Text string `json:"text"` StopReason string `json:"stop_reason"` + // PartialJSON carries a fragment of a tool call's arguments on an + // input_json_delta. The model streams tool arguments as text that + // is only valid JSON once every fragment is concatenated. + PartialJSON string `json:"partial_json"` } `json:"delta"` Usage *anthropicUsage `json:"usage"` + + // Index identifies which content block an event belongs to. A response + // may contain several blocks (text plus one or more tool calls), and + // their deltas are only distinguishable by this index. + Index *int `json:"index"` + // ContentBlock is the opening descriptor on a content_block_start. For a + // tool call it carries the id and name; the arguments follow as deltas. + ContentBlock *struct { + Type string `json:"type"` + ID string `json:"id"` + Name string `json:"name"` + } `json:"content_block"` +} + +// anthropicToolCallState accumulates one streamed tool call. The id and name +// arrive on content_block_start; the arguments follow as a series of +// input_json_delta fragments that are only valid JSON once concatenated. +type anthropicToolCallState struct { + id string + name string + args strings.Builder +} + +// openAnthropicTool starts accumulating a tool call for content block index. +// The entry is pointer-held: a strings.Builder must not be copied once used, +// which a value slice would do the moment append reallocates. +func (s *inferenceStreamState) openAnthropicTool(index *int, id, name string) { + tc := &anthropicToolCallState{id: id, name: name} + s.toolCalls = append(s.toolCalls, tc) + if index != nil { + if s.toolsByIndex == nil { + s.toolsByIndex = map[int]*anthropicToolCallState{} + } + s.toolsByIndex[*index] = tc + } + s.openTool = tc +} + +// anthropicTool resolves the tool call a delta belongs to. Nil index falls +// back to the most recently opened call — blocks are emitted sequentially, +// so that is the same call the index would have named. +func (s *inferenceStreamState) anthropicTool(index *int) *anthropicToolCallState { + if index != nil { + if tc, ok := s.toolsByIndex[*index]; ok { + return tc + } + // An indexed delta for a block we never saw open is a text block's + // delta or a shape we don't model — not the open tool's arguments. + return nil + } + return s.openTool +} + +// closeAnthropicTool drops the fallback target at a content_block_stop, so a +// later unindexed delta can't append to a call that already ended. +func (s *inferenceStreamState) closeAnthropicTool() { + s.openTool = nil } // 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, // 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. +// text_delta blocks; tool calls accumulate from content_block_start plus +// input_json_delta; stop_reason and the cumulative output_tokens arrive in +// message_delta. Unknown events (ping, message_stop) are ignored. func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline.InferenceExtension) { var ev anthropicStreamEvent if err := json.Unmarshal(frame, &ev); err != nil { @@ -210,11 +285,29 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline 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 + } + case "content_block_start": + // A tool call opens here and is populated by later deltas. Text + // blocks need no setup — their deltas append to the completion. + if ev.ContentBlock != nil && ev.ContentBlock.Type == "tool_use" { + state.openAnthropicTool(ev.Index, ev.ContentBlock.ID, ev.ContentBlock.Name) } case "content_block_delta": - if ev.Delta != nil && ev.Delta.Type == "text_delta" { + if ev.Delta == nil { + return + } + switch ev.Delta.Type { + case "text_delta": state.completion.WriteString(ev.Delta.Text) + case "input_json_delta": + if tc := state.anthropicTool(ev.Index); tc != nil { + tc.args.WriteString(ev.Delta.PartialJSON) + } } + case "content_block_stop": + state.closeAnthropicTool() case "message_delta": if ev.Delta != nil && ev.Delta.StopReason != "" { ext.FinishReason = ev.Delta.StopReason @@ -232,6 +325,10 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline // 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 @@ -260,12 +357,7 @@ func parseAnthropicSSE(body []byte, ext *pipeline.InferenceExtension) { } foldAnthropicFrame(data, state, ext) } - ext.Completion = state.completion.String() - if state.usage.TotalTokens > 0 { - ext.PromptTokens = state.usage.PromptTokens - ext.CompletionTokens = state.usage.CompletionTokens - ext.TotalTokens = state.usage.TotalTokens - } + state.finalize(ext) } // rawMessageToMap decodes a JSON object into a map, returning nil for an absent diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go index 04cc36d6c..f1475094e 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go +++ b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go @@ -112,6 +112,11 @@ func TestInferenceParser_AnthropicMessages_NonStreamingResponse(t *testing.T) { t.Errorf("tokens = prompt %d / completion %d / total %d, want 27/8/35", ext.PromptTokens, ext.CompletionTokens, ext.TotalTokens) } + // The cached share of the prompt is recorded separately: 2 read, no writes. + if ext.CacheReadTokens != 2 || ext.CacheWriteTokens != 0 { + t.Errorf("cache = write %d / read %d, want 0/2", + ext.CacheWriteTokens, ext.CacheReadTokens) + } } func TestInferenceParser_AnthropicMessages_StreamFoldsEvents(t *testing.T) { @@ -194,6 +199,152 @@ func TestInferenceParser_AnthropicMessages_StreamBetaPathUsage(t *testing.T) { if ext.FinishReason != "end_turn" { t.Errorf("FinishReason = %q, want end_turn", ext.FinishReason) } + // The same event carries how the cached 33,763 split between writes and + // reads. A write bills 1.25x base and a read 0.1x, so collapsing both into + // PromptTokens leaves a 12.5x spread invisible: this turn cost roughly + // eleven times what the same prompt costs once the entry is warm. + if ext.CacheWriteTokens != 3755 || ext.CacheReadTokens != 30008 { + t.Errorf("cache = write %d / read %d, want 3755/30008", + ext.CacheWriteTokens, ext.CacheReadTokens) + } +} + +// TestInferenceParser_AnthropicMessages_StreamToolUse covers a streamed tool +// call. The pieces arrive across three event types — id and name on +// content_block_start, arguments as input_json_delta fragments that are only +// valid JSON once concatenated — so no single frame carries the call. Before +// this was folded in, a streaming turn recorded finishReason "tool_use" with an +// empty toolCalls list, while the equivalent non-streaming response recorded +// the call in full. +func TestInferenceParser_AnthropicMessages_StreamToolUse(t *testing.T) { + p := NewInferenceParser() + 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":{"usage":{"input_tokens":12,"output_tokens":1}}}`), + []byte(`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`), + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Checking."}}`), + []byte(`{"type":"content_block_stop","index":0}`), + []byte(`{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_1","name":"Read","input":{}}}`), + []byte(`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"file_pa"}}`), + []byte(`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"th\":\"/etc/hosts\"}"}}`), + []byte(`{"type":"content_block_stop","index":1}`), + []byte(`{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":40}}`), + []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 + if len(ext.ToolCalls) != 1 { + t.Fatalf("ToolCalls = %+v, want one call", ext.ToolCalls) + } + tc := ext.ToolCalls[0] + if tc.ID != "toolu_1" || tc.Name != "Read" { + t.Errorf("tool call id/name = %q/%q, want toolu_1/Read", tc.ID, tc.Name) + } + // The two partial_json fragments concatenate into the complete arguments. + if tc.Arguments != `{"file_path":"/etc/hosts"}` { + t.Errorf("Arguments = %q, want {\"file_path\":\"/etc/hosts\"}", tc.Arguments) + } + // The text block is unaffected — only text_delta feeds the completion. + if ext.Completion != "Checking." { + t.Errorf("Completion = %q, want \"Checking.\"", ext.Completion) + } + if ext.FinishReason != "tool_use" { + t.Errorf("FinishReason = %q, want tool_use", ext.FinishReason) + } +} + +// TestInferenceParser_AnthropicMessages_StreamToolUseInterleaved proves the +// fragments are routed by block index rather than by arrival order. The two +// calls' deltas alternate here, which is the shape that would silently +// concatenate one call's arguments into the other if index were ignored. +func TestInferenceParser_AnthropicMessages_StreamToolUseInterleaved(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{Path: "/v1/messages"} + pctx.Extensions.Inference = &pipeline.InferenceExtension{Model: "claude-haiku-4-5", Stream: true, IsAction: true} + + frames := [][]byte{ + []byte(`{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_a","name":"Bash","input":{}}}`), + []byte(`{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_b","name":"Grep","input":{}}}`), + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"cmd\":"}}`), + []byte(`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"pattern\":"}}`), + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"ls\"}"}}`), + []byte(`{"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"\"TODO\"}"}}`), + []byte(`{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":60}}`), + } + for _, f := range frames { + p.OnResponseFrame(context.Background(), pctx, f, false) + } + p.OnResponseFrame(context.Background(), pctx, nil, true) + + ext := pctx.Extensions.Inference + if len(ext.ToolCalls) != 2 { + t.Fatalf("ToolCalls = %+v, want two calls", ext.ToolCalls) + } + // Order follows the opening frames, not the delta ordering. + if ext.ToolCalls[0].Name != "Bash" || ext.ToolCalls[1].Name != "Grep" { + t.Errorf("names = %q/%q, want Bash/Grep", ext.ToolCalls[0].Name, ext.ToolCalls[1].Name) + } + if ext.ToolCalls[0].Arguments != `{"cmd":"ls"}` { + t.Errorf("Bash Arguments = %q, want {\"cmd\":\"ls\"}", ext.ToolCalls[0].Arguments) + } + if ext.ToolCalls[1].Arguments != `{"pattern":"TODO"}` { + t.Errorf("Grep Arguments = %q, want {\"pattern\":\"TODO\"}", ext.ToolCalls[1].Arguments) + } +} + +// TestInferenceParser_AnthropicMessages_RequestContentBytes covers the sizes of +// messages the text flattening discards. In an agent loop those are most of the +// conversation: a tool_result block flattens to "" and reads as free, while the +// model was billed for every byte of it. +func TestInferenceParser_AnthropicMessages_RequestContentBytes(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{ + Path: "/v1/messages", + Body: []byte(`{ + "model": "claude-haiku-4-5", + "max_tokens": 64, + "system": [{"type": "text", "text": "You are Claude Code."}], + "messages": [ + {"role": "user", "content": "read /etc/hosts"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "/etc/hosts"}} + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "toolu_1", "content": "127.0.0.1 localhost"} + ]} + ] + }`), + } + p.OnRequest(context.Background(), pctx) + + ext := pctx.Extensions.Inference + if ext == nil || len(ext.Messages) != 4 { + t.Fatalf("Messages = %+v, want [system, user, assistant, user]", ext) + } + for i, m := range ext.Messages { + if m.ContentBytes <= 0 { + t.Errorf("Messages[%d] (%s) ContentBytes = %d, want > 0", i, m.Role, m.ContentBytes) + } + } + // The two block-array messages carry no text, so Content is empty while + // ContentBytes still reports what the request spent on them. + for _, i := range []int{2, 3} { + if ext.Messages[i].Content != "" { + t.Errorf("Messages[%d] Content = %q, want empty (no text blocks)", i, ext.Messages[i].Content) + } + } + // The tool_result payload is the larger of the two — a real one is a whole + // file, which is exactly the cost this field exists to make visible. + if ext.Messages[3].ContentBytes <= ext.Messages[1].ContentBytes { + t.Errorf("tool_result ContentBytes (%d) should exceed the plain user turn (%d)", + ext.Messages[3].ContentBytes, ext.Messages[1].ContentBytes) + } } // TestInferenceParser_AnthropicMessages_QueryStringPath pins dialect dispatch diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go index 4d7696b31..f6c31eda8 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin.go @@ -104,8 +104,9 @@ func parseOpenAIRequest(body []byte) *pipeline.InferenceExtension { } for _, msg := range req.Messages { ext.Messages = append(ext.Messages, pipeline.InferenceMessage{ - Role: msg.Role, - Content: msg.Content, + Role: msg.Role, + Content: msg.Content, + ContentBytes: msg.ContentBytes, }) } for _, tool := range req.Tools { @@ -165,9 +166,48 @@ func (p *InferenceParser) OnResponse(_ context.Context, pctx *pipeline.Context) // 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. +// +// 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 +// 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 +// distinguishable by the block index the provider stamps on each frame. +// openTool is the fallback for a provider that omits the index. type inferenceStreamState struct { completion strings.Builder usage inferenceUsage + + toolCalls []*anthropicToolCallState + toolsByIndex map[int]*anthropicToolCallState + openTool *anthropicToolCallState +} + +// finalize copies the accumulated stream state onto the public extension +// fields. Every write is an assignment rather than an accumulation, so a +// second finalize on the same state (the buffered OnResponse path running +// 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 len(s.toolCalls) == 0 { + return + } + calls := make([]pipeline.InferenceToolCall, 0, len(s.toolCalls)) + for _, tc := range s.toolCalls { + calls = append(calls, pipeline.InferenceToolCall{ + ID: tc.id, + Name: tc.name, + Arguments: tc.args.String(), + }) + } + ext.ToolCalls = calls } // streamStateKey scopes the scratch state to this plugin in @@ -220,12 +260,7 @@ func (p *InferenceParser) OnResponseFrame(_ context.Context, pctx *pipeline.Cont } if last { - ext.Completion = state.completion.String() - if state.usage.TotalTokens > 0 { - ext.PromptTokens = state.usage.PromptTokens - ext.CompletionTokens = state.usage.CompletionTokens - ext.TotalTokens = state.usage.TotalTokens - } + state.finalize(ext) // Empty stream with no body and no chunks — record Skip to // pair the response row with the request row. if ext.Completion == "" && ext.FinishReason == "" && ext.TotalTokens == 0 { @@ -392,10 +427,21 @@ 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. type inferenceUsage struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` TotalTokens int `json:"total_tokens"` + + CacheWriteTokens int `json:"-"` + CacheReadTokens int `json:"-"` } type inferenceRequest struct { @@ -416,9 +462,14 @@ type inferenceRequest struct { // The array form is used for multi-modal input and tool-result messages. // Non-text parts (image_url, tool_use objects, etc.) are dropped since the // parser only exposes text for downstream policy plugins. +// +// ContentBytes records the size of the content value before that reduction, +// so a message the model was billed for doesn't read as empty just because +// none of it was text. type inferenceMessage struct { - Role string - Content string + Role string + Content string + ContentBytes int } func (m *inferenceMessage) UnmarshalJSON(data []byte) error { @@ -431,9 +482,22 @@ func (m *inferenceMessage) UnmarshalJSON(data []byte) error { } m.Role = raw.Role m.Content = flattenContent(raw.Content) + m.ContentBytes = contentBytes(raw.Content) return nil } +// contentBytes is the wire size of a message's content value, and the source +// of InferenceMessage.ContentBytes. Absent and null content report 0 rather +// than the 4 bytes the literal `null` occupies — the field is a size signal +// for content that exists, and an assistant turn that carries only tool_calls +// has none. +func contentBytes(raw json.RawMessage) int { + if len(raw) == 0 || bytes.Equal(raw, []byte("null")) { + return 0 + } + return len(raw) +} + // flattenContent returns the text representation of an OpenAI content value. // Returns "" when content is absent, null, or contains no text parts. func flattenContent(raw json.RawMessage) string { diff --git a/authbridge/authlib/plugins/inferenceparser/plugin_test.go b/authbridge/authlib/plugins/inferenceparser/plugin_test.go index ad3effe15..ec091b88d 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin_test.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin_test.go @@ -531,6 +531,44 @@ func TestInferenceParser_MultipartContent(t *testing.T) { } } +// TestInferenceParser_ContentBytes covers the OpenAI dialect's own +// ContentBytes accounting — a separate UnmarshalJSON from the Anthropic path. +// A tool-result message whose content is an array flattens to text only for +// the text parts; the dropped parts (images, tool payloads) still cost prompt +// tokens, and ContentBytes is what they contribute. +func TestInferenceParser_ContentBytes(t *testing.T) { + p := NewInferenceParser() + body := `{ + "model": "gpt-4", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "tool", "content": [{"type":"image_url","image_url":{"url":"http://x/very/long/path"}}]}, + {"role": "assistant", "content": null, "tool_calls": []} + ] + }` + pctx := &pipeline.Context{Path: "/v1/chat/completions", Body: []byte(body)} + p.OnRequest(context.Background(), pctx) + + msgs := pctx.Extensions.Inference.Messages + if len(msgs) != 3 { + t.Fatalf("expected 3 messages, got %d", len(msgs)) + } + // `"hi"` on the wire — the quotes are part of the JSON value. + if msgs[0].ContentBytes != 4 { + t.Errorf("msgs[0].ContentBytes = %d, want 4", msgs[0].ContentBytes) + } + // No text parts, so Content is empty — but the message is far from free. + if msgs[1].Content != "" || msgs[1].ContentBytes <= msgs[0].ContentBytes { + t.Errorf("msgs[1] = %q / %d bytes, want empty text and > %d bytes", + msgs[1].Content, msgs[1].ContentBytes, msgs[0].ContentBytes) + } + // content: null has no content to size — reporting the 4 bytes of the + // literal would make an empty assistant turn look like a small payload. + if msgs[2].ContentBytes != 0 { + t.Errorf("msgs[2].ContentBytes = %d, want 0 for null content", msgs[2].ContentBytes) + } +} + func TestInferenceParser_NullContent(t *testing.T) { // Assistant messages that only carry tool_calls have content: null. p := NewInferenceParser() diff --git a/authbridge/cmd/abctl/tui/detail_pane.go b/authbridge/cmd/abctl/tui/detail_pane.go index 546d3e514..e2a9f7ccd 100644 --- a/authbridge/cmd/abctl/tui/detail_pane.go +++ b/authbridge/cmd/abctl/tui/detail_pane.go @@ -157,6 +157,7 @@ var ( inferenceRespKeys = []string{ "model", "completion", "finishReason", "promptTokens", "completionTokens", "totalTokens", "toolCalls", + "cacheWriteTokens", "cacheReadTokens", } mcpReqKeys = []string{"method", "rpcId", "params"} mcpRespKeys = []string{"method", "rpcId", "result", "error"} From bc732020e43ec1e069cca0298783f07fb3f8a0c8 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 26 Aug 2026 16:35:06 -0400 Subject: [PATCH 2/4] fix(inference-parser): keep prompt counts when a stream reports no output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit finalize gates the whole usage copy on TotalTokens > 0, but the Anthropic fold only computed that total inside the output_tokens > 0 arm. Any stream that reported a prompt and no completion therefore refreshed PromptTokens and then threw it away. Two streams do that, and the provider billed the prompt in both: - a terminal message_delta carrying the prompt with output_tokens == 0 (a refusal, or a generation stopped immediately); - a turn the caller interrupted after message_start, which never reaches a message_delta at all — the shape an agent produces every time a user cancels a running turn. Derive the total from the parts after any usage update instead. The interrupted case now records a real response row with prompt tokens and zero completion rather than skip/no_response_body; recovery there is partial by construction, since the ?beta=true path defers the cache counts to message_delta. The OpenAI fold is untouched: it takes total_tokens off the wire, where it can legitimately differ from prompt + completion. Raised in review of #811 by CodeRabbit and @mrsabath. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- .../plugins/inferenceparser/anthropic.go | 19 ++++- .../plugins/inferenceparser/anthropic_test.go | 77 +++++++++++++++++++ 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic.go b/authbridge/authlib/plugins/inferenceparser/anthropic.go index 921426031..2d078d732 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic.go +++ b/authbridge/authlib/plugins/inferenceparser/anthropic.go @@ -269,6 +269,17 @@ 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, @@ -287,6 +298,7 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline state.usage.PromptTokens = ev.Message.Usage.promptTotal() state.usage.CacheWriteTokens = ev.Message.Usage.CacheCreationInputTokens state.usage.CacheReadTokens = ev.Message.Usage.CacheReadInputTokens + state.totalAnthropicUsage() } case "content_block_start": // A tool call opens here and is populated by later deltas. Text @@ -331,12 +343,11 @@ func foldAnthropicFrame(frame []byte, state *inferenceStreamState, ext *pipeline state.usage.CacheReadTokens = ev.Usage.CacheReadInputTokens } 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. + // usage.output_tokens in message_delta is cumulative — take + // the latest rather than accumulating. state.usage.CompletionTokens = ev.Usage.OutputTokens - state.usage.TotalTokens = state.usage.PromptTokens + ev.Usage.OutputTokens } + state.totalAnthropicUsage() } } } diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go index f1475094e..d91360000 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go +++ b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go @@ -209,6 +209,83 @@ func TestInferenceParser_AnthropicMessages_StreamBetaPathUsage(t *testing.T) { } } +// TestInferenceParser_AnthropicMessages_StreamZeroOutputUsage covers a stream +// whose terminal message_delta reports the prompt with output_tokens == 0 — a +// refusal or an immediately-stopped generation. The prompt was billed, cache +// reads included, so the counts must survive: recomputing the total only inside +// the output_tokens > 0 arm left TotalTokens at zero, and finalize gates the +// whole usage copy on that total being non-zero, so a 5,109-token prompt landed +// on the extension as nothing at all. +func TestInferenceParser_AnthropicMessages_StreamZeroOutputUsage(t *testing.T) { + p := NewInferenceParser() + 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_2","type":"message","role":"assistant","usage":{"input_tokens":9,"output_tokens":0}}}`), + []byte(`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"input_tokens":9,"output_tokens":0,"cache_creation_input_tokens":100,"cache_read_input_tokens":5000}}`), + []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 + 100 + 5000, all of it input the provider charged for. + if ext.PromptTokens != 5109 || ext.TotalTokens != 5109 { + t.Errorf("tokens = prompt %d / total %d, want 5109/5109 (usage discarded on the zero-output path?)", + ext.PromptTokens, ext.TotalTokens) + } + if ext.CompletionTokens != 0 { + t.Errorf("CompletionTokens = %d, want 0", ext.CompletionTokens) + } + if ext.CacheWriteTokens != 100 || ext.CacheReadTokens != 5000 { + t.Errorf("cache = write %d / read %d, want 100/5000", + ext.CacheWriteTokens, ext.CacheReadTokens) + } +} + +// TestInferenceParser_AnthropicMessages_StreamInterruptedAfterStart covers a +// stream the caller abandoned after message_start — the shape an agent produces +// every time a user cancels a running turn. No message_delta ever arrives, so +// message_start's input_tokens is the only usage the stream reported, and it is +// the one the provider billed. The turn previously recorded skip/no_response_body +// and no counts, making cancelled turns free in the accounting. +// +// Recovery is partial by construction: on the ?beta=true path the cache counts +// ride on message_delta, so an interrupted turn can only ever report the +// uncached input_tokens. +func TestInferenceParser_AnthropicMessages_StreamInterruptedAfterStart(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{Path: "/v1/messages"} + pctx.Extensions.Inference = &pipeline.InferenceExtension{Model: "claude-haiku-4-5", Stream: true, IsAction: true} + + p.OnResponseFrame(context.Background(), pctx, + []byte(`{"type":"message_start","message":{"id":"msg_bdrk_3","type":"message","role":"assistant","usage":{"input_tokens":24000,"output_tokens":0}}}`), false) + // The connection ends here: no message_delta, no message_stop. + p.OnResponseFrame(context.Background(), pctx, nil, true) + + ext := pctx.Extensions.Inference + if ext.PromptTokens != 24000 || ext.TotalTokens != 24000 { + t.Errorf("tokens = prompt %d / total %d, want 24000/24000 (interrupted stream discarded?)", + ext.PromptTokens, ext.TotalTokens) + } + if ext.CompletionTokens != 0 { + t.Errorf("CompletionTokens = %d, want 0", ext.CompletionTokens) + } + // Having counts to record, the response must be a real row rather than the + // skip that stands in for a stream that reported nothing. Direction's zero + // value is Inbound, so this context's invocations land on the inbound list. + if invs := pctx.Extensions.Invocations; invs != nil { + for _, inv := range invs.Inbound { + if inv.Reason == "no_response_body" { + t.Errorf("recorded %s/%s, want a real response row", inv.Action, inv.Reason) + } + } + } +} + // TestInferenceParser_AnthropicMessages_StreamToolUse covers a streamed tool // call. The pieces arrive across three event types — id and name on // content_block_start, arguments as input_json_delta fragments that are only From ffd62e18a50c899419b588eb393a0d81c3be2305 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 26 Aug 2026 16:49:34 -0400 Subject: [PATCH 3/4] docs(pipeline): record that ToolCalls is dialect-asymmetric on streams MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adding tool-call capture to the Anthropic streaming path made ToolCalls populated on three of the four response paths instead of two. Before that it was uniformly empty on both streaming dialects, so a consumer could read "empty" the same way everywhere; now an empty ToolCalls means different things depending on which dialect streamed the response. The asymmetry is structural, not an oversight: the OpenAI streaming chunk shape decodes only `choices[].delta.content`, so assembling `choices[].delta.tool_calls[]` would need a second index-keyed accumulator — its own change. Until then the only record of the gap is a PR description, which stops being reachable the moment it merges. Document it on the field, in the shape the neighbouring CacheWriteTokens and ContentBytes comments use: say what populates it, and say what its zero value does and does not license a consumer to conclude. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/extensions.go | 24 +++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go index 55104ebf3..802e78192 100644 --- a/authbridge/authlib/pipeline/extensions.go +++ b/authbridge/authlib/pipeline/extensions.go @@ -157,12 +157,24 @@ 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"` - ToolCalls []InferenceToolCall `json:"toolCalls,omitempty"` + 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"` + + // ToolCalls are the tool invocations the model requested. Populated on + // three of the four response paths — both non-streaming dialects and + // Anthropic streaming. An OpenAI *stream* leaves it empty: that dialect + // splits each call across `choices[].delta.tool_calls[]` fragments keyed + // by their own index, a shape the streaming chunk decoder does not read. + // + // So empty means "the model requested no tools" only for a non-streaming + // response or an Anthropic stream. A consumer that spans dialects — cost + // accounting, per-tool attribution — must not read absence as a negative + // on a streamed OpenAI turn, where it is indistinguishable from a turn + // 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 From 78bc732ca92ad7e52dad854ec61d7955fa10fbc8 Mon Sep 17 00:00:00 2001 From: Hai Huang Date: Wed, 26 Aug 2026 19:41:38 -0400 Subject: [PATCH 4/4] Fix: Address review findings on inference parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three findings from review of #814, all latent rather than live today. 1. A stream carrying only tool calls was recorded as a skip. The "no response body" guard tested completion text, finish reason and usage but not ToolCalls, so a turn cancelled while the model was still emitting tool arguments — none of those three present, but a captured call — was labelled as having carried nothing and dropped out of any timeline filtered on observe. Latent because a real Anthropic stream opens with message_start, which sets usage; reachable on a truncated one. 2. An OpenAI usage frame cleared the cache accumulator. foldOpenAIFrame assigned the decoded chunk's usage over state.usage wholesale. inferenceUsage doubles as the OpenAI wire shape and the dialect-neutral accumulator, and its two cache fields are json:"-", so they are always zero in a freshly decoded chunk — the assignment cleared whatever had accumulated. Now copied field by field. Latent because nothing on the OpenAI path fills those fields today, which is precisely what would make the clobber silent when something does. 3. ContentBytes counts whatever whitespace the client's serializer emitted. Documented rather than changed: the measure is of what was sent, so compacting first would buy comparability across clients at the cost of an allocation per message on every request-body parse. Comparable across messages from one client, not across clients. Both source fixes carry a regression test, each verified to fail with its fix reverted. Assisted-By: Claude (Anthropic AI) Signed-off-by: Hai Huang --- authbridge/authlib/pipeline/extensions.go | 6 +++ .../plugins/inferenceparser/anthropic_test.go | 52 +++++++++++++++++++ .../authlib/plugins/inferenceparser/plugin.go | 29 ++++++++++- .../plugins/inferenceparser/plugin_test.go | 30 +++++++++++ 4 files changed, 115 insertions(+), 2 deletions(-) diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go index 802e78192..83bc17ef4 100644 --- a/authbridge/authlib/pipeline/extensions.go +++ b/authbridge/authlib/pipeline/extensions.go @@ -210,6 +210,12 @@ type InferenceMessage struct { // messages contribute, without recording their contents: it is a byte // count of the raw JSON (syntax and escapes included), not a token // count, and is a size signal rather than an exact one. + // + // Whitespace counts too, because the measure is of what was sent, not of + // a normalized form of it. A client that pretty-prints its request bodies + // therefore reports a higher count than one sending compact JSON for the + // same content — tens of percent apart on a deeply nested tool result. + // Comparable across messages from one client; not across clients. ContentBytes int `json:"contentBytes,omitempty"` } diff --git a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go index d91360000..680040597 100644 --- a/authbridge/authlib/plugins/inferenceparser/anthropic_test.go +++ b/authbridge/authlib/plugins/inferenceparser/anthropic_test.go @@ -375,6 +375,58 @@ func TestInferenceParser_AnthropicMessages_StreamToolUseInterleaved(t *testing.T } } +// TestInferenceParser_AnthropicMessages_StreamToolUseOnlyIsNotASkip pins the +// finalize guard against the one stream shape where every other signal is +// absent: a turn cancelled while the model was still emitting tool arguments. +// There is no completion text, no stop_reason, and no usage block, so the +// guard's other three terms all hold — and recording a skip here would label +// a stream that carried a tool call as having had no response body, hiding it +// from any timeline filtered on observe. +// +// A real Anthropic stream opens with message_start, whose input_tokens keeps +// TotalTokens non-zero, so this is a latent case rather than a live one. It is +// also what makes the interleaved test above fragile: drop its trailing +// message_delta and the two calls it proves are captured would vanish into a +// skip. +func TestInferenceParser_AnthropicMessages_StreamToolUseOnlyIsNotASkip(t *testing.T) { + p := NewInferenceParser() + pctx := &pipeline.Context{Path: "/v1/messages"} + pctx.Extensions.Inference = &pipeline.InferenceExtension{Model: "claude-haiku-4-5", Stream: true, IsAction: true} + + frames := [][]byte{ + []byte(`{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_x","name":"Read","input":{}}}`), + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"file\":"}}`), + } + for _, f := range frames { + p.OnResponseFrame(context.Background(), pctx, f, false) + } + p.OnResponseFrame(context.Background(), pctx, nil, true) + + ext := pctx.Extensions.Inference + // Preconditions: this is the shape the guard's other three terms match. + if ext.Completion != "" || ext.FinishReason != "" || ext.TotalTokens != 0 { + t.Fatalf("fixture no longer exercises the guard: Completion=%q FinishReason=%q TotalTokens=%d", + ext.Completion, ext.FinishReason, ext.TotalTokens) + } + if len(ext.ToolCalls) != 1 || ext.ToolCalls[0].Name != "Read" { + t.Fatalf("ToolCalls = %+v, want one Read call", ext.ToolCalls) + } + // Arguments stay as the model left them — truncated, not discarded. + if ext.ToolCalls[0].Arguments != `{"file":` { + t.Errorf("Arguments = %q, want the partial fragment", ext.ToolCalls[0].Arguments) + } + + inv := pctx.Extensions.Invocations + if inv == nil || len(inv.Inbound) == 0 { + t.Fatalf("Invocations = %+v, want a recorded response row", inv) + } + last := inv.Inbound[len(inv.Inbound)-1] + if last.Action != pipeline.ActionObserve { + t.Errorf("action = %s/%s, want observe (a captured tool call is a response body)", + last.Action, last.Reason) + } +} + // TestInferenceParser_AnthropicMessages_RequestContentBytes covers the sizes of // messages the text flattening discards. In an agent loop those are most of the // conversation: a tool_result block flattens to "" and reads as free, while the diff --git a/authbridge/authlib/plugins/inferenceparser/plugin.go b/authbridge/authlib/plugins/inferenceparser/plugin.go index f6c31eda8..1e0a7ac5b 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin.go @@ -263,7 +263,15 @@ func (p *InferenceParser) OnResponseFrame(_ context.Context, pctx *pipeline.Cont state.finalize(ext) // Empty stream with no body and no chunks — record Skip to // pair the response row with the request row. - if ext.Completion == "" && ext.FinishReason == "" && ext.TotalTokens == 0 { + // + // Tool calls count as a body. A turn cancelled while the model was + // still emitting tool arguments has no completion text, no finish + // reason, and no usage block, but finalize has captured the call — + // so skipping here would label a stream that demonstrably carried + // content as having none, and drop it out of any timeline filtered + // on observe. + if ext.Completion == "" && ext.FinishReason == "" && ext.TotalTokens == 0 && + len(ext.ToolCalls) == 0 { pctx.Skip("no_response_body") return pipeline.Action{Type: pipeline.Continue} } @@ -294,8 +302,19 @@ 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. if chunk.Usage.TotalTokens > 0 { - state.usage = chunk.Usage + state.usage.PromptTokens = chunk.Usage.PromptTokens + state.usage.CompletionTokens = chunk.Usage.CompletionTokens + state.usage.TotalTokens = chunk.Usage.TotalTokens } } @@ -491,6 +510,12 @@ func (m *inferenceMessage) UnmarshalJSON(data []byte) error { // than the 4 bytes the literal `null` occupies — the field is a size signal // for content that exists, and an assistant turn that carries only tool_calls // has none. +// +// raw is the client's bytes verbatim, so the count includes any whitespace the +// client's serializer emitted. That is deliberate: this measures what was +// sent. Compacting first would buy comparability across clients at the cost of +// an allocation per message on every request-body parse, and would no longer +// answer "how big was this on the wire". func contentBytes(raw json.RawMessage) int { if len(raw) == 0 || bytes.Equal(raw, []byte("null")) { return 0 diff --git a/authbridge/authlib/plugins/inferenceparser/plugin_test.go b/authbridge/authlib/plugins/inferenceparser/plugin_test.go index ec091b88d..50a29ca6a 100644 --- a/authbridge/authlib/plugins/inferenceparser/plugin_test.go +++ b/authbridge/authlib/plugins/inferenceparser/plugin_test.go @@ -470,6 +470,36 @@ 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{