Skip to content

Commit 885ffac

Browse files
luispaterluispater
authored
Merge pull request #4528 from router-for-me/fix/issue-81-normalized-token-accounting
fix(usage): add canonical token accounting v2
2 parents e572987 + fe8a616 commit 885ffac

7 files changed

Lines changed: 639 additions & 47 deletions

File tree

‎internal/redisqueue/plugin.go‎

Lines changed: 23 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -65,21 +65,16 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
6565
}
6666
responseServiceTier := strings.TrimSpace(record.ResponseServiceTier)
6767

68+
usageDetail := coreusage.EnsureTokenBreakdownForProvider(record.Detail, record.Provider, record.ExecutorType)
6869
tokens := tokenStats{
69-
InputTokens: record.Detail.InputTokens,
70-
OutputTokens: record.Detail.OutputTokens,
71-
ReasoningTokens: record.Detail.ReasoningTokens,
72-
CachedTokens: record.Detail.CachedTokens,
73-
CacheReadTokens: record.Detail.CacheReadTokens,
70+
InputTokens: usageDetail.InputTokens,
71+
OutputTokens: usageDetail.OutputTokens,
72+
ReasoningTokens: usageDetail.ReasoningTokens,
73+
CachedTokens: usageDetail.CachedTokens,
74+
CacheReadTokens: usageDetail.CacheReadTokens,
7475
CacheReadTokensPresent: true,
75-
CacheCreationTokens: record.Detail.CacheCreationTokens,
76-
TotalTokens: record.Detail.TotalTokens,
77-
}
78-
if tokens.TotalTokens == 0 {
79-
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
80-
}
81-
if tokens.TotalTokens == 0 {
82-
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
76+
CacheCreationTokens: usageDetail.CacheCreationTokens,
77+
TotalTokens: usageDetail.TotalTokens,
8378
}
8479

8580
failed := record.Failed
@@ -103,6 +98,8 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
10398

10499
payload, err := json.Marshal(queuedUsageDetail{
105100
requestDetail: detail,
101+
AccountingVersion: coreusage.TokenAccountingSchemaVersion,
102+
TokenBreakdown: usageDetail.TokenBreakdown,
106103
Provider: provider,
107104
ExecutorType: executorType,
108105
Model: modelName,
@@ -123,17 +120,19 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
123120

124121
type queuedUsageDetail struct {
125122
requestDetail
126-
Provider string `json:"provider"`
127-
ExecutorType string `json:"executor_type"`
128-
Model string `json:"model"`
129-
Alias string `json:"alias"`
130-
Endpoint string `json:"endpoint"`
131-
AuthType string `json:"auth_type"`
132-
APIKey string `json:"api_key"`
133-
RequestID string `json:"request_id"`
134-
ReasoningEffort string `json:"reasoning_effort"`
135-
ServiceTier string `json:"service_tier"`
136-
ResponseServiceTier string `json:"response_service_tier,omitempty"`
123+
AccountingVersion int `json:"accounting_version"`
124+
TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"`
125+
Provider string `json:"provider"`
126+
ExecutorType string `json:"executor_type"`
127+
Model string `json:"model"`
128+
Alias string `json:"alias"`
129+
Endpoint string `json:"endpoint"`
130+
AuthType string `json:"auth_type"`
131+
APIKey string `json:"api_key"`
132+
RequestID string `json:"request_id"`
133+
ReasoningEffort string `json:"reasoning_effort"`
134+
ServiceTier string `json:"service_tier"`
135+
ResponseServiceTier string `json:"response_service_tier,omitempty"`
137136
}
138137

139138
type requestDetail struct {

‎internal/redisqueue/plugin_test.go‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
6161
requireStringField(t, payload, "service_tier", "auto")
6262
requireMissingField(t, payload, "request_service_tier")
6363
requireStringField(t, payload, "response_service_tier", "default")
64+
requireIntField(t, payload, "accounting_version", coreusage.TokenAccountingSchemaVersion)
65+
requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, 30)
6466
requireTokensBoolField(t, payload, "cache_read_tokens_present", true)
6567
requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"})
6668
requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"})
@@ -70,6 +72,38 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
7072
})
7173
}
7274

75+
func TestUsageQueuePluginNormalizesDirectSDKUsageByProvider(t *testing.T) {
76+
tests := []struct {
77+
provider string
78+
wantTotal int
79+
}{
80+
{provider: "openai", wantTotal: 130},
81+
{provider: "gemini", wantTotal: 142},
82+
}
83+
for _, tt := range tests {
84+
t.Run(tt.provider, func(t *testing.T) {
85+
withEnabledQueue(t, func() {
86+
ctx := internallogging.WithResponseStatusHolder(context.Background())
87+
internallogging.SetResponseStatus(ctx, http.StatusOK)
88+
89+
(&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{
90+
Provider: tt.provider,
91+
Model: "direct-sdk-model",
92+
Detail: coreusage.Detail{
93+
InputTokens: 100,
94+
OutputTokens: 30,
95+
ReasoningTokens: 12,
96+
},
97+
})
98+
99+
payload := popSinglePayload(t)
100+
requireIntField(t, requireTokensPayload(t, payload), "total_tokens", tt.wantTotal)
101+
requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, int64(tt.wantTotal))
102+
})
103+
})
104+
}
105+
}
106+
73107
func TestUsageQueuePluginPayloadIncludesGenerateFalse(t *testing.T) {
74108
withEnabledQueue(t, func() {
75109
ctx := internallogging.WithResponseStatusHolder(context.Background())
@@ -397,6 +431,38 @@ func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, w
397431
}
398432
}
399433

434+
func requireIntField(t *testing.T, payload map[string]json.RawMessage, key string, want int) {
435+
t.Helper()
436+
437+
raw, ok := payload[key]
438+
if !ok {
439+
t.Fatalf("payload missing %q", key)
440+
}
441+
var got int
442+
if err := json.Unmarshal(raw, &got); err != nil {
443+
t.Fatalf("unmarshal %q: %v", key, err)
444+
}
445+
if got != want {
446+
t.Fatalf("%s = %d, want %d", key, got, want)
447+
}
448+
}
449+
450+
func requireTokenBreakdown(t *testing.T, payload map[string]json.RawMessage, quality coreusage.TokenAccountingQuality, total int64) {
451+
t.Helper()
452+
453+
raw, ok := payload["token_breakdown"]
454+
if !ok {
455+
t.Fatal("payload missing token_breakdown")
456+
}
457+
var breakdown coreusage.TokenBreakdown
458+
if err := json.Unmarshal(raw, &breakdown); err != nil {
459+
t.Fatalf("unmarshal token_breakdown: %v", err)
460+
}
461+
if !breakdown.Valid() || breakdown.Quality != quality || breakdown.TotalTokens != total {
462+
t.Fatalf("token_breakdown = %+v, want quality=%s total=%d", breakdown, quality, total)
463+
}
464+
}
465+
400466
func requireMissingField(t *testing.T, payload map[string]json.RawMessage, key string) {
401467
t.Helper()
402468

‎internal/runtime/executor/helps/usage_helpers.go‎

Lines changed: 50 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,7 @@ func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.De
177177
if model == "" {
178178
return usage.Record{}, false
179179
}
180-
detail = normalizeUsageDetailTotal(detail)
180+
detail = normalizeUsageDetailTotal(detail, r.provider, r.executorType)
181181
if !hasNonZeroTokenUsage(detail) {
182182
return usage.Record{}, false
183183
}
@@ -201,20 +201,14 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det
201201
if r == nil {
202202
return
203203
}
204-
detail = normalizeUsageDetailTotal(detail)
204+
detail = normalizeUsageDetailTotal(detail, r.provider, r.executorType)
205205
r.once.Do(func() {
206206
r.publishRecord(ctx, r.buildRecord(detail, failed, fail))
207207
})
208208
}
209209

210-
func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail {
211-
if detail.TotalTokens == 0 {
212-
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
213-
if total > 0 {
214-
detail.TotalTokens = total
215-
}
216-
}
217-
return detail
210+
func normalizeUsageDetailTotal(detail usage.Detail, provider, executorType string) usage.Detail {
211+
return usage.EnsureTokenBreakdownForProvider(detail, provider, executorType)
218212
}
219213

220214
func hasNonZeroTokenUsage(detail usage.Detail) bool {
@@ -224,7 +218,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool {
224218
detail.CachedTokens != 0 ||
225219
detail.CacheReadTokens != 0 ||
226220
detail.CacheCreationTokens != 0 ||
227-
detail.TotalTokens != 0
221+
detail.TotalTokens != 0 ||
222+
detail.TokenBreakdown.TotalTokens != 0
228223
}
229224

230225
// ensurePublished guarantees that a usage record is emitted exactly once.
@@ -556,11 +551,14 @@ func hasOpenAIStyleUsageTokenFields(usageNode gjson.Result) bool {
556551
if !usageNode.Exists() || !usageNode.IsObject() {
557552
return false
558553
}
554+
return usageNode.Get("total_tokens").Exists() || hasOpenAIStyleUsageBucketFields(usageNode)
555+
}
556+
557+
func hasOpenAIStyleUsageBucketFields(usageNode gjson.Result) bool {
559558
return usageNode.Get("prompt_tokens").Exists() ||
560559
usageNode.Get("input_tokens").Exists() ||
561560
usageNode.Get("completion_tokens").Exists() ||
562561
usageNode.Get("output_tokens").Exists() ||
563-
usageNode.Get("total_tokens").Exists() ||
564562
usageNode.Get("prompt_tokens_details.cached_tokens").Exists() ||
565563
usageNode.Get("input_tokens_details.cached_tokens").Exists() ||
566564
usageNode.Get("prompt_tokens_details.cache_write_tokens").Exists() ||
@@ -610,6 +608,21 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail {
610608
if reasoning.Exists() {
611609
detail.ReasoningTokens = reasoning.Int()
612610
}
611+
if hasOpenAIStyleUsageBucketFields(usageNode) {
612+
detail.TokenBreakdown = usage.NewSubsetTokenBreakdown(
613+
detail.InputTokens,
614+
detail.CacheReadTokens,
615+
detail.CacheCreationTokens,
616+
detail.OutputTokens,
617+
detail.ReasoningTokens,
618+
detail.TotalTokens,
619+
)
620+
} else {
621+
detail.TokenBreakdown = usage.NewUnclassifiedTokenBreakdown(detail.TotalTokens)
622+
}
623+
if detail.TotalTokens == 0 {
624+
detail.TotalTokens = detail.TokenBreakdown.TotalTokens
625+
}
613626
return detail
614627
}
615628

@@ -665,6 +678,14 @@ func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail {
665678
detail.CachedTokens = detail.CacheCreationTokens
666679
}
667680
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.CacheReadTokens + detail.CacheCreationTokens
681+
detail.TokenBreakdown = usage.NewIndependentTokenBreakdown(
682+
detail.InputTokens,
683+
detail.CacheReadTokens,
684+
detail.CacheCreationTokens,
685+
detail.OutputTokens,
686+
detail.ReasoningTokens,
687+
detail.TotalTokens,
688+
)
668689
return detail
669690
}
670691

@@ -681,6 +702,14 @@ func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
681702
if detail.TotalTokens == 0 {
682703
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
683704
}
705+
detail.TokenBreakdown = usage.NewSeparateReasoningTokenBreakdown(
706+
detail.InputTokens,
707+
detail.CacheReadTokens,
708+
detail.CacheCreationTokens,
709+
detail.OutputTokens,
710+
detail.ReasoningTokens,
711+
detail.TotalTokens,
712+
)
684713
return detail
685714
}
686715

@@ -699,11 +728,16 @@ func parseInteractionsUsageDetail(node gjson.Result) usage.Detail {
699728
detail.CacheReadTokens = detail.CachedTokens
700729
}
701730
if detail.TotalTokens == 0 {
702-
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + detail.CacheCreationTokens
703-
if cacheRead.Exists() {
704-
detail.TotalTokens += detail.CacheReadTokens
705-
}
731+
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
706732
}
733+
detail.TokenBreakdown = usage.NewSeparateReasoningTokenBreakdown(
734+
detail.InputTokens,
735+
detail.CacheReadTokens,
736+
detail.CacheCreationTokens,
737+
detail.OutputTokens,
738+
detail.ReasoningTokens,
739+
detail.TotalTokens,
740+
)
707741
return detail
708742
}
709743

0 commit comments

Comments
 (0)