From 817becab58e28189447ff413f16eb01fc5eb5846 Mon Sep 17 00:00:00 2001 From: MK Date: Mon, 7 Sep 2026 19:51:42 -0400 Subject: [PATCH] fix(anthropic): capture input + cache tokens on the streaming path (#433) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit readAnthropicStream handled only message_delta (output_tokens), so a streamed Anthropic call dropped ALL input tokens — and, under prompt caching, the cache read/creation counts too. This is the streaming sibling of the non-streaming fix in #431/#432. Anthropic reports input_tokens (+ cache_read/creation) on message_start and accumulates output_tokens onto message_delta. Now: - parse message_start usage into locals, - emit ONE complete UsageInfo on the terminal message_delta (input + cache read + creation + output + TotalTokens), so the streamed usage matches the non-streaming path. A single authoritative Usage is correct whether a consumer overwrites (result.Usage = *delta.Usage, the existing pattern) or sums per-delta — avoiding both input-loss on overwrite and double-count on sum. Tests: cache-heavy stream recovers input/cache/output/total; a non-cached stream populates input+output with zero cache fields. golangci-lint clean; full llm suite passes. --- forge-core/llm/providers/anthropic.go | 38 +++++- .../llm/providers/anthropic_stream_test.go | 115 ++++++++++++++++++ 2 files changed, 152 insertions(+), 1 deletion(-) create mode 100644 forge-core/llm/providers/anthropic_stream_test.go diff --git a/forge-core/llm/providers/anthropic.go b/forge-core/llm/providers/anthropic.go index b5fc693..4db400e 100644 --- a/forge-core/llm/providers/anthropic.go +++ b/forge-core/llm/providers/anthropic.go @@ -417,6 +417,20 @@ type anthropicContentBlockDelta struct { } `json:"delta"` } +// anthropicMessageStart carries the initial usage on a streaming response: +// Anthropic reports input_tokens (+ prompt-cache read/creation) on the +// message_start event, while output_tokens accumulates onto message_delta. +// Without parsing this, the streaming path drops ALL input tokens (#433). +type anthropicMessageStart struct { + Message struct { + Usage struct { + InputTokens int `json:"input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` + } `json:"usage"` + } `json:"message"` +} + type anthropicMessageDelta struct { Delta struct { StopReason string `json:"stop_reason"` @@ -430,6 +444,13 @@ func (c *AnthropicClient) readAnthropicStream(r io.Reader, ch chan<- llm.StreamD scanner := bufio.NewScanner(r) var currentToolCall *llm.ToolCall var eventType string + // Input + prompt-cache tokens arrive on message_start; output accumulates + // onto message_delta. Capture the input side here and emit the COMPLETE + // usage once, on the terminal message_delta — a single authoritative + // UsageInfo is correct whether a consumer overwrites (result.Usage = + // *delta.Usage) or sums per-delta, and avoids losing input on the common + // overwrite path (#433). + var inputTokens, cacheReadTokens, cacheCreationTokens int for scanner.Scan() { line := scanner.Text() @@ -445,6 +466,15 @@ func (c *AnthropicClient) readAnthropicStream(r io.Reader, ch chan<- llm.StreamD } switch eventType { + case "message_start": + var ev anthropicMessageStart + if json.Unmarshal([]byte(after), &ev) != nil { + continue + } + inputTokens = ev.Message.Usage.InputTokens + cacheReadTokens = ev.Message.Usage.CacheReadInputTokens + cacheCreationTokens = ev.Message.Usage.CacheCreationInputTokens + case "content_block_start": var ev anthropicContentBlockStart if json.Unmarshal([]byte(after), &ev) != nil { @@ -494,7 +524,13 @@ func (c *AnthropicClient) readAnthropicStream(r io.Reader, ch chan<- llm.StreamD ch <- llm.StreamDelta{ FinishReason: finishReason, Usage: &llm.UsageInfo{ - OutputTokens: ev.Usage.OutputTokens, + InputTokens: inputTokens, + OutputTokens: ev.Usage.OutputTokens, + CacheReadInputTokens: cacheReadTokens, + CacheCreationInputTokens: cacheCreationTokens, + // True total incl. cached prefix, matching the + // non-streaming path (#431/#432). + TotalTokens: inputTokens + cacheReadTokens + cacheCreationTokens + ev.Usage.OutputTokens, }, } diff --git a/forge-core/llm/providers/anthropic_stream_test.go b/forge-core/llm/providers/anthropic_stream_test.go new file mode 100644 index 0000000..d0eeeae --- /dev/null +++ b/forge-core/llm/providers/anthropic_stream_test.go @@ -0,0 +1,115 @@ +package providers + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/initializ/forge/forge-core/llm" +) + +// #433: the Anthropic streaming path must capture input + prompt-cache tokens +// from message_start (they are NOT on message_delta, which carries only +// output), so a streamed call's final usage matches the non-streaming path +// instead of undercounting to zero input. +func TestAnthropicChatStream_CapturesInputAndCacheTokens(t *testing.T) { + // event:-framed SSE, exactly as Anthropic streams it. input + cache ride + // on message_start; output accumulates onto message_delta. + sse := "event: message_start\n" + + `data: {"type":"message_start","message":{"usage":{"input_tokens":12,"cache_read_input_tokens":4000,"cache_creation_input_tokens":200,"output_tokens":1}}}` + "\n\n" + + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":25}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = r.Body.Close() + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(sse)) + })) + defer srv.Close() + + c := NewAnthropicClient(llm.ClientConfig{APIKey: "x", BaseURL: srv.URL, Model: "claude-sonnet-4-6"}) + ch, err := c.ChatStream(context.Background(), &llm.ChatRequest{ + Model: "claude-sonnet-4-6", + Messages: []llm.ChatMessage{{Role: llm.RoleUser, Content: "hi"}}, + }) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + + var content string + var usage llm.UsageInfo + for delta := range ch { + content += delta.Content + if delta.Usage != nil { + usage = *delta.Usage // overwrite — the provider emits one complete Usage + } + } + + if content != "hi" { + t.Errorf("content = %q, want %q", content, "hi") + } + if usage.InputTokens != 12 { + t.Errorf("InputTokens = %d, want 12 (from message_start, previously dropped)", usage.InputTokens) + } + if usage.OutputTokens != 25 { + t.Errorf("OutputTokens = %d, want 25 (from message_delta)", usage.OutputTokens) + } + if usage.CacheReadInputTokens != 4000 { + t.Errorf("CacheReadInputTokens = %d, want 4000", usage.CacheReadInputTokens) + } + if usage.CacheCreationInputTokens != 200 { + t.Errorf("CacheCreationInputTokens = %d, want 200", usage.CacheCreationInputTokens) + } + if got := usage.TotalInputTokens(); got != 4212 { + t.Errorf("TotalInputTokens() = %d, want 4212 (12+4000+200)", got) + } + if usage.TotalTokens != 4237 { + t.Errorf("TotalTokens = %d, want 4237 (input+cache_read+cache_creation+output)", usage.TotalTokens) + } +} + +// A non-cached stream: input + output populate, cache fields stay zero. +func TestAnthropicChatStream_NoCaching(t *testing.T) { + sse := "event: message_start\n" + + `data: {"type":"message_start","message":{"usage":{"input_tokens":30,"output_tokens":1}}}` + "\n\n" + + "event: message_delta\n" + + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":8}}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte(sse)) + })) + defer srv.Close() + + c := NewAnthropicClient(llm.ClientConfig{APIKey: "x", BaseURL: srv.URL, Model: "claude-sonnet-4-6"}) + ch, err := c.ChatStream(context.Background(), &llm.ChatRequest{ + Model: "claude-sonnet-4-6", + Messages: []llm.ChatMessage{{Role: llm.RoleUser, Content: "hi"}}, + }) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + var usage llm.UsageInfo + for delta := range ch { + if delta.Usage != nil { + usage = *delta.Usage + } + } + if usage.InputTokens != 30 || usage.OutputTokens != 8 { + t.Errorf("usage = %+v, want input=30 output=8", usage) + } + if usage.CacheReadInputTokens != 0 || usage.CacheCreationInputTokens != 0 { + t.Errorf("cache tokens should be zero for a non-cached stream, got read=%d creation=%d", + usage.CacheReadInputTokens, usage.CacheCreationInputTokens) + } + if usage.TotalTokens != 38 { + t.Errorf("TotalTokens = %d, want 38", usage.TotalTokens) + } +}