Skip to content

Commit 4ae1ec4

Browse files
hustxiayangaabchookanurag94aishwaryaraimule21gemini-code-assist[bot]
authored
fix: double counting for the usage in stream mode for anthropic models (#2239)
**Description** `ExtractTokenUsageFromExplicitCaching`  is called in both start message and the delta message, as a result. The usage report is not correct. Fix: remove the accumulation in the message delta; also the previous tests can pass because the test cases did not mock the real cases, update all related tests cases as well. --------- Signed-off-by: yxia216 <yxia216@bloomberg.net> Signed-off-by: Aaron Choo <achoo30@bloomberg.net> Signed-off-by: Anurag Aggarwal <kanurag94@gmail.com> Signed-off-by: Aishwarya <aishwarya.raimule@nutanix.com> Signed-off-by: aishwaryaraimule21 <aishwarya.raimule@nutanix.com> Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: Hritik003 <hritik.raj@nutanix.com> Signed-off-by: Linus Schlumberger <linus.schlumberger@siemens.com> Signed-off-by: Ignasi Barrera <nacx@apache.org> Signed-off-by: Chris Burns <29541485+ChrisJBurns@users.noreply.github.com> Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com> Signed-off-by: hustxiayang <yxia216@bloomberg.net> Signed-off-by: hustxiayang <yangxiast@gmail.com> Co-authored-by: Aaron Choo <achoo30@bloomberg.net> Co-authored-by: Anurag Aggarwal <kanurag94@gmail.com> Co-authored-by: aishwaryaraimule21 <aishwarya.raimule@nutanix.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Ignasi Barrera <ignasi@tetrate.io> Co-authored-by: Hritik Raj <hritik.raj@nutanix.com> Co-authored-by: Linus Schlumberger <linus.schlumberger@siemens.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: Chris Burns <29541485+ChrisJBurns@users.noreply.github.com> Co-authored-by: Takeshi Yoneda <t.y.mathetake@gmail.com> Co-authored-by: Ignasi Barrera <nacx@apache.org>
1 parent 2cc2417 commit 4ae1ec4

3 files changed

Lines changed: 161 additions & 24 deletions

File tree

‎internal/translator/anthropic_helper.go‎

Lines changed: 2 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -966,8 +966,8 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
966966
&u.CacheReadInputTokens,
967967
&u.CacheCreationInputTokens,
968968
)
969-
// For message_start, we store the initial usage but don't add to the accumulated
970-
// The message_delta event will contain the final totals
969+
// Set all input token counts (input, cache read, cache creation) from message_start.
970+
// message_delta may also contain these fields but only output_tokens is used from it.
971971
if input, ok := usage.InputTokens(); ok {
972972
p.tokenUsage.SetInputTokens(input)
973973
}
@@ -1061,17 +1061,6 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
10611061
if output, ok := usage.OutputTokens(); ok {
10621062
p.tokenUsage.AddOutputTokens(output)
10631063
}
1064-
// Update input tokens to include any cache tokens from delta
1065-
if cached, ok := usage.CachedInputTokens(); ok {
1066-
p.tokenUsage.AddInputTokens(cached)
1067-
// Accumulate any additional cache tokens from delta
1068-
p.tokenUsage.AddCachedInputTokens(cached)
1069-
}
1070-
if cacheCreation, ok := usage.CacheCreationInputTokens(); ok {
1071-
p.tokenUsage.AddInputTokens(cacheCreation)
1072-
// Accumulate cache creation tokens
1073-
p.tokenUsage.AddCacheCreationInputTokens(cacheCreation)
1074-
}
10751064
p.tokenUsage.SetReasoningTokens(uint32(u.OutputTokensDetails.ThinkingTokens)) //nolint:gosec
10761065
if event.Delta.StopReason != "" {
10771066
p.stopReason = event.Delta.StopReason

‎internal/translator/anthropic_helper_test.go‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,18 @@ package translator
77

88
import (
99
"fmt"
10+
"strings"
1011
"testing"
1112

1213
"github.com/anthropics/anthropic-sdk-go"
1314
"github.com/anthropics/anthropic-sdk-go/shared/constant"
15+
"github.com/stretchr/testify/assert"
1416
"github.com/stretchr/testify/require"
1517
"k8s.io/utils/ptr"
1618

1719
"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
1820
"github.com/envoyproxy/ai-gateway/internal/internalapi"
21+
"github.com/envoyproxy/ai-gateway/internal/metrics"
1922
)
2023

2124
// mockErrorReader is a helper for testing io.Reader failures.
@@ -1157,3 +1160,148 @@ func TestBuildAnthropicParamsWithReasoningEffort(t *testing.T) {
11571160
require.Equal(t, anthropic.OutputConfigEffort(""), params.OutputConfig.Effort)
11581161
})
11591162
}
1163+
1164+
func TestAnthropicStreamParser_StreamingTokenUsage(t *testing.T) {
1165+
tests := []struct {
1166+
name string
1167+
events string
1168+
expectedInputTokens uint32
1169+
expectedOutputTokens uint32
1170+
expectedTotalTokens uint32
1171+
expectedCachedTokens uint32
1172+
expectedCacheCreationTokens uint32
1173+
}{
1174+
{
1175+
name: "with cache tokens",
1176+
events: `event: message_start
1177+
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}}}
1178+
1179+
event: content_block_start
1180+
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
1181+
1182+
event: content_block_delta
1183+
data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}}
1184+
1185+
event: content_block_stop
1186+
data: {"type": "content_block_stop", "index": 0}
1187+
1188+
event: message_delta
1189+
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}}
1190+
1191+
event: message_stop
1192+
data: {"type": "message_stop"}
1193+
1194+
`,
1195+
expectedInputTokens: 14041, // 678 + 13363 + 0
1196+
expectedOutputTokens: 5,
1197+
expectedTotalTokens: 14046, // 14041 + 5
1198+
expectedCachedTokens: 13363,
1199+
expectedCacheCreationTokens: 0,
1200+
},
1201+
{
1202+
name: "without cache tokens",
1203+
events: `event: message_start
1204+
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}}}
1205+
1206+
event: content_block_start
1207+
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
1208+
1209+
event: content_block_delta
1210+
data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}
1211+
1212+
event: content_block_stop
1213+
data: {"type": "content_block_stop", "index": 0}
1214+
1215+
event: message_delta
1216+
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}}
1217+
1218+
event: message_stop
1219+
data: {"type": "message_stop"}
1220+
1221+
`,
1222+
expectedInputTokens: 100,
1223+
expectedOutputTokens: 10,
1224+
expectedTotalTokens: 110,
1225+
expectedCachedTokens: 0,
1226+
expectedCacheCreationTokens: 0,
1227+
},
1228+
{
1229+
name: "with cache creation tokens",
1230+
events: `event: message_start
1231+
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}}}
1232+
1233+
event: content_block_start
1234+
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
1235+
1236+
event: content_block_delta
1237+
data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Response"}}
1238+
1239+
event: content_block_stop
1240+
data: {"type": "content_block_stop", "index": 0}
1241+
1242+
event: message_delta
1243+
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}}
1244+
1245+
event: message_stop
1246+
data: {"type": "message_stop"}
1247+
1248+
`,
1249+
expectedInputTokens: 5200, // 200 + 5000 + 0
1250+
expectedOutputTokens: 8,
1251+
expectedTotalTokens: 5208, // 5200 + 8
1252+
expectedCachedTokens: 0,
1253+
expectedCacheCreationTokens: 5000,
1254+
},
1255+
}
1256+
1257+
for _, tt := range tests {
1258+
t.Run(tt.name, func(t *testing.T) {
1259+
parser := newAnthropicStreamParser("claude-sonnet-4-6")
1260+
1261+
// Feed each event block separately (simulating chunked SSE delivery),
1262+
// with the last chunk marked as endOfStream.
1263+
chunks := splitSSEEvents(tt.events)
1264+
var tokenUsage metrics.TokenUsage
1265+
for i, chunk := range chunks {
1266+
endOfStream := i == len(chunks)-1
1267+
_, _, usage, _, err := parser.Process(strings.NewReader(chunk), endOfStream, nil)
1268+
require.NoError(t, err)
1269+
if endOfStream {
1270+
tokenUsage = usage
1271+
}
1272+
}
1273+
1274+
inputTokens, inputSet := tokenUsage.InputTokens()
1275+
assert.True(t, inputSet, "InputTokens should be set")
1276+
assert.Equal(t, tt.expectedInputTokens, inputTokens, "InputTokens mismatch")
1277+
1278+
outputTokens, outputSet := tokenUsage.OutputTokens()
1279+
assert.True(t, outputSet, "OutputTokens should be set")
1280+
assert.Equal(t, tt.expectedOutputTokens, outputTokens, "OutputTokens mismatch")
1281+
1282+
totalTokens, totalSet := tokenUsage.TotalTokens()
1283+
assert.True(t, totalSet, "TotalTokens should be set")
1284+
assert.Equal(t, tt.expectedTotalTokens, totalTokens, "TotalTokens mismatch")
1285+
1286+
cachedTokens, cachedSet := tokenUsage.CachedInputTokens()
1287+
assert.True(t, cachedSet, "CachedInputTokens should be set")
1288+
assert.Equal(t, tt.expectedCachedTokens, cachedTokens, "CachedInputTokens mismatch")
1289+
1290+
cacheCreation, cacheCreationSet := tokenUsage.CacheCreationInputTokens()
1291+
assert.True(t, cacheCreationSet, "CacheCreationInputTokens should be set")
1292+
assert.Equal(t, tt.expectedCacheCreationTokens, cacheCreation, "CacheCreationInputTokens mismatch")
1293+
})
1294+
}
1295+
}
1296+
1297+
func splitSSEEvents(data string) []string {
1298+
parts := strings.Split(data, "\n\n")
1299+
var events []string
1300+
for _, p := range parts {
1301+
trimmed := strings.TrimSpace(p)
1302+
if trimmed != "" {
1303+
events = append(events, p+"\n\n")
1304+
}
1305+
}
1306+
return events
1307+
}

‎tests/data-plane/testupstream_test.go‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -551,7 +551,7 @@ data: [DONE]
551551
expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken},
552552
responseStatus: strconv.Itoa(http.StatusOK),
553553
responseBody: `event: message_start
554-
data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 15}}}
554+
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}}}
555555
556556
event: content_block_start
557557
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
@@ -566,7 +566,7 @@ event: content_block_stop
566566
data: {"type": "content_block_stop", "index": 0}
567567
568568
event: message_delta
569-
data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 12, "cache_read_input_tokens":10}}
569+
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}}
570570
571571
event: message_stop
572572
data: {"type": "message_stop"}
@@ -614,7 +614,7 @@ data: [DONE]
614614
expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken},
615615
responseStatus: strconv.Itoa(http.StatusOK),
616616
responseBody: `event: message_start
617-
data: {"type": "message_start", "message": {"id": "msg_123", "usage": {"input_tokens": 50}}}
617+
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}}}
618618
619619
event: content_block_start
620620
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
629629
data: {"type": "content_block_stop", "index": 0}
630630
631631
event: message_delta
632-
data: {"type": "message_delta", "delta": {"stop_reason": "tool_use"}, "usage": {"output_tokens": 20}}
632+
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}}
633633
634634
event: message_stop
635635
data: {"type": "message_stop"}`,
@@ -910,7 +910,7 @@ data: [DONE]
910910
expRequestHeaders: map[string]string{"Authorization": "Bearer " + fakeGCPAuthToken},
911911
responseStatus: strconv.Itoa(http.StatusOK),
912912
responseBody: `event: message_start
913-
data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8}}}
913+
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}}}
914914
915915
event: content_block_start
916916
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
@@ -922,15 +922,15 @@ event: content_block_stop
922922
data: {"type": "content_block_stop", "index": 0}
923923
924924
event: message_delta
925-
data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 15}}
925+
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}}
926926
927927
event: message_stop
928928
data: {"type": "message_stop"}
929929
930930
`,
931931
expStatus: http.StatusOK,
932932
expResponseBody: `event: message_start
933-
data: {"type": "message_start", "message": {"id": "msg_789", "usage": {"input_tokens": 8}}}
933+
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}}}
934934
935935
event: content_block_start
936936
data: {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
@@ -942,7 +942,7 @@ event: content_block_stop
942942
data: {"type": "content_block_stop", "index": 0}
943943
944944
event: message_delta
945-
data: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 15}}
945+
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}}
946946
947947
event: message_stop
948948
data: {"type": "message_stop"}
@@ -1052,7 +1052,7 @@ event: content_block_stop
10521052
data: {"type":"content_block_stop","index":0}
10531053
10541054
event: message_delta
1055-
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":9,"output_tokens":10}}
1055+
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}}
10561056
10571057
event: message_stop
10581058
data: {"type":"message_stop"}
@@ -1090,7 +1090,7 @@ data: {"type":"message_stop"}
10901090
{"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19kZWx0YSIsImluZGV4IjowLCJkZWx0YSI6eyJ0eXBlIjoidGV4dF9kZWx0YSIsInRleHQiOiLwn5GLIEhvdyJ9fQ==","p":"abcdefghijklmnopqrstuvwxyzABCDEFG"}
10911091
{"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19kZWx0YSIsImluZGV4IjowLCJkZWx0YSI6eyJ0eXBlIjoidGV4dF9kZWx0YSIsInRleHQiOiIgYXJlIHlvdSBkb2luZyB0b2RheT8ifX0=","p":"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ01234"}
10921092
{"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19zdG9wIiwiaW5kZXgiOjB9","p":"abcdefghijklmnopqrstuvwxyz"}
1093-
{"bytes":"eyJ0eXBlIjoibWVzc2FnZV9kZWx0YSIsImRlbHRhIjp7InN0b3BfcmVhc29uIjoiZW5kX3R1cm4iLCJzdG9wX3NlcXVlbmNlIjpudWxsfSwidXNhZ2UiOnsib3V0cHV0X3Rva2VucyI6MTV9fQ==","p":"abcdefghijklmnopqrstu"}
1093+
{"bytes":"eyJ0eXBlIjoibWVzc2FnZV9kZWx0YSIsImRlbHRhIjp7InN0b3BfcmVhc29uIjoiZW5kX3R1cm4iLCJzdG9wX3NlcXVlbmNlIjpudWxsfSwidXNhZ2UiOnsiaW5wdXRfdG9rZW5zIjoxMCwiY2FjaGVfY3JlYXRpb25faW5wdXRfdG9rZW5zIjowLCJjYWNoZV9yZWFkX2lucHV0X3Rva2VucyI6MCwib3V0cHV0X3Rva2VucyI6MTV9fQ==","p":"abcdefghijklmnopqrstu"}
10941094
{"bytes":"eyJ0eXBlIjoibWVzc2FnZV9zdG9wIiwiYW1hem9uLWJlZHJvY2staW52b2NhdGlvbk1ldHJpY3MiOnsiaW5wdXRUb2tlbkNvdW50IjoxMCwib3V0cHV0VG9rZW5Db3VudCI6MTUsImludm9jYXRpb25MYXRlbmN5IjoxNzk4LCJmaXJzdEJ5dGVMYXRlbmN5IjoxNTA3fX0=","p":"ab"}
10951095
`,
10961096
expStatus: http.StatusOK,
@@ -1119,7 +1119,7 @@ event: content_block_stop
11191119
data: {"type":"content_block_stop","index":0}
11201120
11211121
event: message_delta
1122-
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":15}}
1122+
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}}
11231123
11241124
event: message_stop
11251125
data: {"type":"message_stop","amazon-bedrock-invocationMetrics":{"inputTokenCount":10,"outputTokenCount":15,"invocationLatency":1798,"firstByteLatency":1507}}

0 commit comments

Comments
 (0)