Feat: Record per-tool-call cost inputs on the Anthropic path

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.

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