diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go index 221c056ad..83bc17ef4 100644 --- a/authbridge/authlib/pipeline/extensions.go +++ b/authbridge/authlib/pipeline/extensions.go @@ -157,12 +157,41 @@ 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 + // (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 +201,22 @@ 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. + // + // 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"` } // 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..2d078d732 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,89 @@ 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 +} + +// 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, // 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 +296,30 @@ 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 + state.totalAnthropicUsage() + } + 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,14 +337,17 @@ 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 - // 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() } } } @@ -260,12 +368,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..680040597 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,281 @@ 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_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 +// 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_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 +// 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..1e0a7ac5b 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,15 +260,18 @@ 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 { + // + // 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} } @@ -259,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 } } @@ -392,10 +446,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 +481,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 +501,28 @@ 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. +// +// 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 + } + 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..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{ @@ -531,6 +561,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"}