Skip to content

Commit bd8c05a

Browse files
committed
feat(usage): add support for detailed token breakdown in usage tracking
- Introduced `CacheReadTokens` and `CacheCreationTokens` to enhance token breakdown. - Refactored `parseClaudeUsageNode` for cleaner and reusable logic. - Adjusted helpers and updated token calculations to align with the new fields.
1 parent 041ccf0 commit bd8c05a

3 files changed

Lines changed: 37 additions & 31 deletions

File tree

internal/redisqueue/plugin.go

Lines changed: 14 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,13 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
4949
requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
5050

5151
tokens := tokenStats{
52-
InputTokens: record.Detail.InputTokens,
53-
OutputTokens: record.Detail.OutputTokens,
54-
ReasoningTokens: record.Detail.ReasoningTokens,
55-
CachedTokens: record.Detail.CachedTokens,
56-
TotalTokens: record.Detail.TotalTokens,
52+
InputTokens: record.Detail.InputTokens,
53+
OutputTokens: record.Detail.OutputTokens,
54+
ReasoningTokens: record.Detail.ReasoningTokens,
55+
CachedTokens: record.Detail.CachedTokens,
56+
CacheReadTokens: record.Detail.CacheReadTokens,
57+
CacheCreationTokens: record.Detail.CacheCreationTokens,
58+
TotalTokens: record.Detail.TotalTokens,
5759
}
5860
if tokens.TotalTokens == 0 {
5961
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
@@ -116,11 +118,13 @@ type requestDetail struct {
116118
}
117119

118120
type tokenStats struct {
119-
InputTokens int64 `json:"input_tokens"`
120-
OutputTokens int64 `json:"output_tokens"`
121-
ReasoningTokens int64 `json:"reasoning_tokens"`
122-
CachedTokens int64 `json:"cached_tokens"`
123-
TotalTokens int64 `json:"total_tokens"`
121+
InputTokens int64 `json:"input_tokens"`
122+
OutputTokens int64 `json:"output_tokens"`
123+
ReasoningTokens int64 `json:"reasoning_tokens"`
124+
CachedTokens int64 `json:"cached_tokens"`
125+
CacheReadTokens int64 `json:"cache_read_tokens"`
126+
CacheCreationTokens int64 `json:"cache_creation_tokens"`
127+
TotalTokens int64 `json:"total_tokens"`
124128
}
125129

126130
type failDetail struct {

internal/runtime/executor/helps/usage_helpers.go

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool {
116116
detail.OutputTokens != 0 ||
117117
detail.ReasoningTokens != 0 ||
118118
detail.CachedTokens != 0 ||
119+
detail.CacheReadTokens != 0 ||
120+
detail.CacheCreationTokens != 0 ||
119121
detail.TotalTokens != 0
120122
}
121123

@@ -356,17 +358,7 @@ func ParseClaudeUsage(data []byte) usage.Detail {
356358
if !usageNode.Exists() {
357359
return usage.Detail{}
358360
}
359-
detail := usage.Detail{
360-
InputTokens: usageNode.Get("input_tokens").Int(),
361-
OutputTokens: usageNode.Get("output_tokens").Int(),
362-
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
363-
}
364-
if detail.CachedTokens == 0 {
365-
// fall back to creation tokens when read tokens are absent
366-
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
367-
}
368-
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
369-
return detail
361+
return parseClaudeUsageNode(usageNode)
370362
}
371363

372364
func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
@@ -378,16 +370,24 @@ func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
378370
if !usageNode.Exists() {
379371
return usage.Detail{}, false
380372
}
373+
return parseClaudeUsageNode(usageNode), true
374+
}
375+
376+
func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
377+
cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int()
378+
cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int()
381379
detail := usage.Detail{
382-
InputTokens: usageNode.Get("input_tokens").Int(),
383-
OutputTokens: usageNode.Get("output_tokens").Int(),
384-
CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),
380+
InputTokens: usageNode.Get("input_tokens").Int(),
381+
OutputTokens: usageNode.Get("output_tokens").Int(),
382+
CachedTokens: cacheReadTokens,
383+
CacheReadTokens: cacheReadTokens,
384+
CacheCreationTokens: cacheCreationTokens,
385385
}
386386
if detail.CachedTokens == 0 {
387-
detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()
387+
detail.CachedTokens = detail.CacheCreationTokens
388388
}
389389
detail.TotalTokens = detail.InputTokens + detail.OutputTokens
390-
return detail, true
390+
return detail
391391
}
392392

393393
func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {

sdk/cliproxy/usage/manager.go

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,11 +34,13 @@ type Failure struct {
3434

3535
// Detail holds the token usage breakdown.
3636
type Detail struct {
37-
InputTokens int64
38-
OutputTokens int64
39-
ReasoningTokens int64
40-
CachedTokens int64
41-
TotalTokens int64
37+
InputTokens int64
38+
OutputTokens int64
39+
ReasoningTokens int64
40+
CachedTokens int64
41+
CacheReadTokens int64
42+
CacheCreationTokens int64
43+
TotalTokens int64
4244
}
4345

4446
type requestedModelAliasContextKey struct{}

0 commit comments

Comments
 (0)