diff --git a/internal/translator/anthropic_helper.go b/internal/translator/anthropic_helper.go index 2e919789ee..a852d1a678 100644 --- a/internal/translator/anthropic_helper.go +++ b/internal/translator/anthropic_helper.go @@ -966,8 +966,8 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat &u.CacheReadInputTokens, &u.CacheCreationInputTokens, ) - // For message_start, we store the initial usage but don't add to the accumulated - // The message_delta event will contain the final totals + // Set all input token counts (input, cache read, cache creation) from message_start. + // message_delta may also contain these fields but only output_tokens is used from it. if input, ok := usage.InputTokens(); ok { p.tokenUsage.SetInputTokens(input) } @@ -1061,17 +1061,6 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat if output, ok := usage.OutputTokens(); ok { p.tokenUsage.AddOutputTokens(output) } - // Update input tokens to include any cache tokens from delta - if cached, ok := usage.CachedInputTokens(); ok { - p.tokenUsage.AddInputTokens(cached) - // Accumulate any additional cache tokens from delta - p.tokenUsage.AddCachedInputTokens(cached) - } - if cacheCreation, ok := usage.CacheCreationInputTokens(); ok { - p.tokenUsage.AddInputTokens(cacheCreation) - // Accumulate cache creation tokens - p.tokenUsage.AddCacheCreationInputTokens(cacheCreation) - } p.tokenUsage.SetReasoningTokens(uint32(u.OutputTokensDetails.ThinkingTokens)) //nolint:gosec if event.Delta.StopReason != "" { p.stopReason = event.Delta.StopReason diff --git a/internal/translator/anthropic_helper_test.go b/internal/translator/anthropic_helper_test.go index eb39929731..c395094854 100644 --- a/internal/translator/anthropic_helper_test.go +++ b/internal/translator/anthropic_helper_test.go @@ -7,15 +7,18 @@ package translator import ( "fmt" + "strings" "testing" "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/shared/constant" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "k8s.io/utils/ptr" "github.com/envoyproxy/ai-gateway/internal/apischema/openai" "github.com/envoyproxy/ai-gateway/internal/internalapi" + "github.com/envoyproxy/ai-gateway/internal/metrics" ) // mockErrorReader is a helper for testing io.Reader failures. @@ -1157,3 +1160,148 @@ func TestBuildAnthropicParamsWithReasoningEffort(t *testing.T) { require.Equal(t, anthropic.OutputConfigEffort(""), params.OutputConfig.Effort) }) } + +func TestAnthropicStreamParser_StreamingTokenUsage(t *testing.T) { + tests := []struct { + name string + events string + expectedInputTokens uint32 + expectedOutputTokens uint32 + expectedTotalTokens uint32 + expectedCachedTokens uint32 + expectedCacheCreationTokens uint32 + }{ + { + name: "with cache tokens", + events: `event: message_start +data: {"type": "message_start", "message": {"id": "msg_abc123", "type": "message", "role": "assistant", "content": [], "model": "claude-sonnet-4-6", "usage": {"input_tokens": 678, "cache_read_input_tokens": 13363, "cache_creation_input_tokens": 0, "output_tokens": 1}}} + +event: content_block_start +data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + +event: content_block_delta +data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}} + +event: content_block_stop +data: {"type": "content_block_stop", "index": 0} + +event: message_delta +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 678, "cache_read_input_tokens": 13363, "cache_creation_input_tokens": 0, "output_tokens": 5}} + +event: message_stop +data: {"type": "message_stop"} + +`, + expectedInputTokens: 14041, // 678 + 13363 + 0 + expectedOutputTokens: 5, + expectedTotalTokens: 14046, // 14041 + 5 + expectedCachedTokens: 13363, + expectedCacheCreationTokens: 0, + }, + { + name: "without cache tokens", + events: `event: message_start +data: {"type": "message_start", "message": {"id": "msg_abc456", "type": "message", "role": "assistant", "content": [], "model": "claude-sonnet-4-6", "usage": {"input_tokens": 100, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, "output_tokens": 1}}} + +event: content_block_start +data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + +event: content_block_delta +data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}} + +event: content_block_stop +data: {"type": "content_block_stop", "index": 0} + +event: message_delta +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 100, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 0, "output_tokens": 10}} + +event: message_stop +data: {"type": "message_stop"} + +`, + expectedInputTokens: 100, + expectedOutputTokens: 10, + expectedTotalTokens: 110, + expectedCachedTokens: 0, + expectedCacheCreationTokens: 0, + }, + { + name: "with cache creation tokens", + events: `event: message_start +data: {"type": "message_start", "message": {"id": "msg_abc789", "type": "message", "role": "assistant", "content": [], "model": "claude-sonnet-4-6", "usage": {"input_tokens": 200, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 5000, "output_tokens": 1}}} + +event: content_block_start +data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + +event: content_block_delta +data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Response"}} + +event: content_block_stop +data: {"type": "content_block_stop", "index": 0} + +event: message_delta +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 200, "cache_read_input_tokens": 0, "cache_creation_input_tokens": 5000, "output_tokens": 8}} + +event: message_stop +data: {"type": "message_stop"} + +`, + expectedInputTokens: 5200, // 200 + 5000 + 0 + expectedOutputTokens: 8, + expectedTotalTokens: 5208, // 5200 + 8 + expectedCachedTokens: 0, + expectedCacheCreationTokens: 5000, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + parser := newAnthropicStreamParser("claude-sonnet-4-6") + + // Feed each event block separately (simulating chunked SSE delivery), + // with the last chunk marked as endOfStream. + chunks := splitSSEEvents(tt.events) + var tokenUsage metrics.TokenUsage + for i, chunk := range chunks { + endOfStream := i == len(chunks)-1 + _, _, usage, _, err := parser.Process(strings.NewReader(chunk), endOfStream, nil) + require.NoError(t, err) + if endOfStream { + tokenUsage = usage + } + } + + inputTokens, inputSet := tokenUsage.InputTokens() + assert.True(t, inputSet, "InputTokens should be set") + assert.Equal(t, tt.expectedInputTokens, inputTokens, "InputTokens mismatch") + + outputTokens, outputSet := tokenUsage.OutputTokens() + assert.True(t, outputSet, "OutputTokens should be set") + assert.Equal(t, tt.expectedOutputTokens, outputTokens, "OutputTokens mismatch") + + totalTokens, totalSet := tokenUsage.TotalTokens() + assert.True(t, totalSet, "TotalTokens should be set") + assert.Equal(t, tt.expectedTotalTokens, totalTokens, "TotalTokens mismatch") + + cachedTokens, cachedSet := tokenUsage.CachedInputTokens() + assert.True(t, cachedSet, "CachedInputTokens should be set") + assert.Equal(t, tt.expectedCachedTokens, cachedTokens, "CachedInputTokens mismatch") + + cacheCreation, cacheCreationSet := tokenUsage.CacheCreationInputTokens() + assert.True(t, cacheCreationSet, "CacheCreationInputTokens should be set") + assert.Equal(t, tt.expectedCacheCreationTokens, cacheCreation, "CacheCreationInputTokens mismatch") + }) + } +} + +func splitSSEEvents(data string) []string { + parts := strings.Split(data, "\n\n") + var events []string + for _, p := range parts { + trimmed := strings.TrimSpace(p) + if trimmed != "" { + events = append(events, p+"\n\n") + } + } + return events +} diff --git a/tests/data-plane/testupstream_test.go b/tests/data-plane/testupstream_test.go index 19bf7d503f..c3801e27a4 100644 --- a/tests/data-plane/testupstream_test.go +++ b/tests/data-plane/testupstream_test.go @@ -551,7 +551,7 @@ data: [DONE] expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), responseBody: `event: message_start -data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 15}}} +data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 15, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 10, "output_tokens": 1}}} event: content_block_start data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} @@ -566,7 +566,7 @@ event: content_block_stop data: {"type": "content_block_stop", "index": 0} event: message_delta -data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 12, "cache_read_input_tokens":10}} +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 15, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 10, "output_tokens": 12}} event: message_stop data: {"type": "message_stop"} @@ -614,7 +614,7 @@ data: [DONE] expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), responseBody: `event: message_start -data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 50}}} +data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 50, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 1}}} event: content_block_start data: {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_abc123", "name": "get_weather", "input": {}}} @@ -629,7 +629,7 @@ event: content_block_stop data: {"type": "content_block_stop", "index": 0} event: message_delta -data: {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 20}} +data: {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"input_tokens": 50, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 20}} event: message_stop data: {"type": "message_stop"}`, @@ -910,7 +910,7 @@ data: [DONE] expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken}, responseStatus: strconv.Itoa(http.StatusOK), responseBody: `event: message_start -data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8}}} +data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 1}}} event: content_block_start data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} @@ -922,7 +922,7 @@ event: content_block_stop data: {"type": "content_block_stop", "index": 0} event: message_delta -data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 15}} +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 15}} event: message_stop data: {"type": "message_stop"} @@ -930,7 +930,7 @@ data: {"type": "message_stop"} `, expStatus: http.StatusOK, expResponseBody: `event: message_start -data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8}}} +data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 1}}} event: content_block_start data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} @@ -942,7 +942,7 @@ event: content_block_stop data: {"type": "content_block_stop", "index": 0} event: message_delta -data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 15}} +data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 15}} event: message_stop data: {"type": "message_stop"} @@ -1052,7 +1052,7 @@ event: content_block_stop data: {"type":"content_block_stop","index":0} event: message_delta -data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":9,"output_tokens":10}} +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":9,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":10}} event: message_stop data: {"type":"message_stop"} @@ -1090,7 +1090,7 @@ data: {"type":"message_stop"} {"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19kZWx0YSIsImluZGV4IjowLCJkZWx0YSI6eyJ0eXBlIjoidGV4dF9kZWx0YSIsInRleHQiOiLwn5GLIEhvdyJ9fQ==","p":"abcdefghijklmnopqrstuvwxyzABCDEFG"} {"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19kZWx0YSIsImluZGV4IjowLCJkZWx0YSI6eyJ0eXBlIjoidGV4dF9kZWx0YSIsInRleHQiOiIgYXJlIHlvdSBkb2luZyB0b2RheT8ifX0=","p":"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234"} {"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19zdG9wIiwiaW5kZXgiOjB9","p":"abcdefghijklmnopqrstuvwxyz"} -{"bytes":"eyJ0eXBlIjoibWVzc2FnZV9kZWx0YSIsImRlbHRhIjp7InN0b3BfcmVhc29uIjoiZW5kX3R1cm4iLCJzdG9wX3NlcXVlbmNlIjpudWxsfSwidXNhZ2UiOnsib3V0cHV0X3Rva2VucyI6MTV9fQ==","p":"abcdefghijklmnopqrstu"} +{"bytes":"eyJ0eXBlIjoibWVzc2FnZV9kZWx0YSIsImRlbHRhIjp7InN0b3BfcmVhc29uIjoiZW5kX3R1cm4iLCJzdG9wX3NlcXVlbmNlIjpudWxsfSwidXNhZ2UiOnsiaW5wdXRfdG9rZW5zIjoxMCwiY2FjaGVfY3JlYXRpb25faW5wdXRfdG9rZW5zIjowLCJjYWNoZV9yZWFkX2lucHV0X3Rva2VucyI6MCwib3V0cHV0X3Rva2VucyI6MTV9fQ==","p":"abcdefghijklmnopqrstu"} {"bytes":"eyJ0eXBlIjoibWVzc2FnZV9zdG9wIiwiYW1hem9uLWJlZHJvY2staW52b2NhdGlvbk1ldHJpY3MiOnsiaW5wdXRUb2tlbkNvdW50IjoxMCwib3V0cHV0VG9rZW5Db3VudCI6MTUsImludm9jYXRpb25MYXRlbmN5IjoxNzk4LCJmaXJzdEJ5dGVMYXRlbmN5IjoxNTA3fX0=","p":"ab"} `, expStatus: http.StatusOK, @@ -1119,7 +1119,7 @@ event: content_block_stop data: {"type":"content_block_stop","index":0} event: message_delta -data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":15}} +data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":10,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":15}} event: message_stop data: {"type":"message_stop","amazon-bedrock-invocationMetrics":{"inputTokenCount":10,"outputTokenCount":15,"invocationLatency":1798,"firstByteLatency":1507}}