Skip to content

Commit 65e760a

Browse files
committed
feat(usage): include cache tokens in total token calculation and add tests
- Updated `TotalTokens` calculation to account for `CacheReadTokens` and `CacheCreationTokens`. - Added tests to validate accurate token aggregation and fallback behavior for `CachedTokens`.
1 parent 2bcc762 commit 65e760a

2 files changed

Lines changed: 35 additions & 1 deletion

File tree

internal/runtime/executor/helps/usage_helpers.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -527,7 +527,7 @@ func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
527527
if detail.CachedTokens == 0 {
528528
detail.CachedTokens = detail.CacheCreationTokens
529529
}
530-
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
530+
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.CacheReadTokens + detail.CacheCreationTokens
531531
return detail
532532
}
533533

internal/runtime/executor/helps/usage_helpers_test.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,40 @@ func TestParseOpenAIStreamUsageResponsesFields(t *testing.T) {
8989
}
9090
}
9191

92+
func TestParseClaudeUsageIncludesCacheTokensInTotal(t *testing.T) {
93+
data := []byte(`{"usage":{"input_tokens":3085,"output_tokens":253,"cache_read_input_tokens":7,"cache_creation_input_tokens":19514}}`)
94+
detail := ParseClaudeUsage(data)
95+
if detail.InputTokens != 3085 {
96+
t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 3085)
97+
}
98+
if detail.OutputTokens != 253 {
99+
t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 253)
100+
}
101+
if detail.CacheReadTokens != 7 {
102+
t.Fatalf("cache read tokens = %d, want %d", detail.CacheReadTokens, 7)
103+
}
104+
if detail.CacheCreationTokens != 19514 {
105+
t.Fatalf("cache creation tokens = %d, want %d", detail.CacheCreationTokens, 19514)
106+
}
107+
if detail.CachedTokens != 7 {
108+
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 7)
109+
}
110+
if detail.TotalTokens != 22859 {
111+
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 22859)
112+
}
113+
}
114+
115+
func TestParseClaudeUsageFallsBackCachedTokensToCacheCreation(t *testing.T) {
116+
data := []byte(`{"usage":{"input_tokens":3085,"output_tokens":253,"cache_creation_input_tokens":19514}}`)
117+
detail := ParseClaudeUsage(data)
118+
if detail.CachedTokens != 19514 {
119+
t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 19514)
120+
}
121+
if detail.TotalTokens != 22852 {
122+
t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 22852)
123+
}
124+
}
125+
92126
func TestParseGeminiCLIUsage_TopLevelUsageMetadata(t *testing.T) {
93127
data := []byte(`{"usageMetadata":{"promptTokenCount":11,"candidatesTokenCount":7,"thoughtsTokenCount":3,"totalTokenCount":21,"cachedContentTokenCount":5}}`)
94128
detail := ParseGeminiCLIUsage(data)

0 commit comments

Comments
 (0)