History

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.

Luis Pater committed 6 months agov7.0.3

Commit bd8c05a1 parent 041ccf03 files · +37 −31Browse files

Unified · Split

3 files
  1. internal/redisqueue/plugin.go +14 −10
  2. internal/runtime/executor/helps/usage_helpers.go +16 −16
  3. sdk/cliproxy/usage/manager.go +7 −5
internal/redisqueue/plugin.go+14−10
@@ -49,11 +49,13 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec · 48 unmodified lines
49 requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))49 requestID := strings.TrimSpace(internallogging.GetRequestID(ctx))
50 50
51 tokens := tokenStats{51 tokens := tokenStats{
52 InputTokens: record.Detail.InputTokens,52 InputTokens: record.Detail.InputTokens,
53 OutputTokens: record.Detail.OutputTokens,53 OutputTokens: record.Detail.OutputTokens,
54 ReasoningTokens: record.Detail.ReasoningTokens,54 ReasoningTokens: record.Detail.ReasoningTokens,
55 CachedTokens: record.Detail.CachedTokens,55 CachedTokens: record.Detail.CachedTokens,
56 TotalTokens: record.Detail.TotalTokens,56 CacheReadTokens: record.Detail.CacheReadTokens,
57 CacheCreationTokens: record.Detail.CacheCreationTokens,
58 TotalTokens: record.Detail.TotalTokens,
57 }59 }
58 if tokens.TotalTokens == 0 {60 if tokens.TotalTokens == 0 {
59 tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens61 tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
@@ -116,11 +118,13 @@ type requestDetail struct { · 56 unmodified lines
116}118}
117 119
118type tokenStats struct {120type tokenStats struct {
119 InputTokens int64 `json:"input_tokens"`121 InputTokens int64 `json:"input_tokens"`
120 OutputTokens int64 `json:"output_tokens"`122 OutputTokens int64 `json:"output_tokens"`
121 ReasoningTokens int64 `json:"reasoning_tokens"`123 ReasoningTokens int64 `json:"reasoning_tokens"`
122 CachedTokens int64 `json:"cached_tokens"`124 CachedTokens int64 `json:"cached_tokens"`
123 TotalTokens int64 `json:"total_tokens"`125 CacheReadTokens int64 `json:"cache_read_tokens"`
126 CacheCreationTokens int64 `json:"cache_creation_tokens"`
127 TotalTokens int64 `json:"total_tokens"`
124}128}
125 129
126type failDetail struct {130type failDetail struct {
internal/runtime/executor/helps/usage_helpers.go+16−16
@@ -116,6 +116,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool { · 115 unmodified lines
116 detail.OutputTokens != 0 ||116 detail.OutputTokens != 0 ||
117 detail.ReasoningTokens != 0 ||117 detail.ReasoningTokens != 0 ||
118 detail.CachedTokens != 0 ||118 detail.CachedTokens != 0 ||
119 detail.CacheReadTokens != 0 ||
120 detail.CacheCreationTokens != 0 ||
119 detail.TotalTokens != 0121 detail.TotalTokens != 0
120}122}
121 123
@@ -356,17 +358,7 @@ func ParseClaudeUsage(data []byte) usage.Detail { · 234 unmodified lines
356 if !usageNode.Exists() {358 if !usageNode.Exists() {
357 return usage.Detail{}359 return usage.Detail{}
358 }360 }
359 detail := usage.Detail{361 return parseClaudeUsageNode(usageNode)
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
370}362}
371 363
372func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {364func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
@@ -378,16 +370,24 @@ func ParseClaudeStreamUsage(line []byte) (usage.Detail, bool) { · 5 unmodified lines
378 if !usageNode.Exists() {370 if !usageNode.Exists() {
379 return usage.Detail{}, false371 return usage.Detail{}, false
380 }372 }
373 return parseClaudeUsageNode(usageNode), true
374}
375
376func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
377 cacheReadTokens := usageNode.Get("cache_read_input_tokens").Int()
378 cacheCreationTokens := usageNode.Get("cache_creation_input_tokens").Int()
381 detail := usage.Detail{379 detail := usage.Detail{
382 InputTokens: usageNode.Get("input_tokens").Int(),380 InputTokens: usageNode.Get("input_tokens").Int(),
383 OutputTokens: usageNode.Get("output_tokens").Int(),381 OutputTokens: usageNode.Get("output_tokens").Int(),
384 CachedTokens: usageNode.Get("cache_read_input_tokens").Int(),382 CachedTokens: cacheReadTokens,
383 CacheReadTokens: cacheReadTokens,
384 CacheCreationTokens: cacheCreationTokens,
385 }385 }
386 if detail.CachedTokens == 0 {386 if detail.CachedTokens == 0 {
387 detail.CachedTokens = usageNode.Get("cache_creation_input_tokens").Int()387 detail.CachedTokens = detail.CacheCreationTokens
388 }388 }
389 detail.TotalTokens = detail.InputTokens + detail.OutputTokens389 detail.TotalTokens = detail.InputTokens + detail.OutputTokens
390 return detail, true390 return detail
391}391}
392 392
393func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {393func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
sdk/cliproxy/usage/manager.go+7−5
@@ -34,11 +34,13 @@ type Failure struct { · 33 unmodified lines
34 34
35// Detail holds the token usage breakdown.35// Detail holds the token usage breakdown.
36type Detail struct {36type Detail struct {
37 InputTokens int6437 InputTokens int64
38 OutputTokens int6438 OutputTokens int64
39 ReasoningTokens int6439 ReasoningTokens int64
40 CachedTokens int6440 CachedTokens int64
41 TotalTokens int6441 CacheReadTokens int64
42 CacheCreationTokens int64
43 TotalTokens int64
42}44}
43 45
44type requestedModelAliasContextKey struct{}46type requestedModelAliasContextKey struct{}