@@ -7,15 +7,18 @@ package translator
77
88import (
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+ }
0 commit comments