From 416a08017447d060c3b91bd928e2f2b631b601ef Mon Sep 17 00:00:00 2001 From: Dylan <1990016@gmail.com> Date: Thu, 23 Jul 2026 21:27:03 +0800 Subject: [PATCH 1/2] fix(usage): add normalized token accounting v2 --- internal/redisqueue/plugin.go | 47 ++-- internal/redisqueue/plugin_test.go | 34 +++ .../runtime/executor/helps/usage_helpers.go | 51 +++- .../executor/helps/usage_helpers_test.go | 49 +++- sdk/cliproxy/usage/accounting.go | 239 ++++++++++++++++++ sdk/cliproxy/usage/accounting_test.go | 56 ++++ sdk/cliproxy/usage/manager.go | 1 + 7 files changed, 434 insertions(+), 43 deletions(-) create mode 100644 sdk/cliproxy/usage/accounting.go create mode 100644 sdk/cliproxy/usage/accounting_test.go diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 0968606f5..784532eb4 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -65,21 +65,16 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } responseServiceTier := strings.TrimSpace(record.ResponseServiceTier) + usageDetail := coreusage.EnsureTokenBreakdown(record.Detail) tokens := tokenStats{ - InputTokens: record.Detail.InputTokens, - OutputTokens: record.Detail.OutputTokens, - ReasoningTokens: record.Detail.ReasoningTokens, - CachedTokens: record.Detail.CachedTokens, - CacheReadTokens: record.Detail.CacheReadTokens, + InputTokens: usageDetail.InputTokens, + OutputTokens: usageDetail.OutputTokens, + ReasoningTokens: usageDetail.ReasoningTokens, + CachedTokens: usageDetail.CachedTokens, + CacheReadTokens: usageDetail.CacheReadTokens, CacheReadTokensPresent: true, - CacheCreationTokens: record.Detail.CacheCreationTokens, - TotalTokens: record.Detail.TotalTokens, - } - if tokens.TotalTokens == 0 { - tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens - } - if tokens.TotalTokens == 0 { - tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens + CacheCreationTokens: usageDetail.CacheCreationTokens, + TotalTokens: usageDetail.TotalTokens, } failed := record.Failed @@ -103,6 +98,8 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec payload, err := json.Marshal(queuedUsageDetail{ requestDetail: detail, + AccountingVersion: coreusage.TokenAccountingSchemaVersion, + TokenBreakdown: usageDetail.TokenBreakdown, Provider: provider, ExecutorType: executorType, Model: modelName, @@ -123,17 +120,19 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec type queuedUsageDetail struct { requestDetail - Provider string `json:"provider"` - ExecutorType string `json:"executor_type"` - Model string `json:"model"` - Alias string `json:"alias"` - Endpoint string `json:"endpoint"` - AuthType string `json:"auth_type"` - APIKey string `json:"api_key"` - RequestID string `json:"request_id"` - ReasoningEffort string `json:"reasoning_effort"` - ServiceTier string `json:"service_tier"` - ResponseServiceTier string `json:"response_service_tier,omitempty"` + AccountingVersion int `json:"accounting_version"` + TokenBreakdown coreusage.TokenBreakdown `json:"token_breakdown"` + Provider string `json:"provider"` + ExecutorType string `json:"executor_type"` + Model string `json:"model"` + Alias string `json:"alias"` + Endpoint string `json:"endpoint"` + AuthType string `json:"auth_type"` + APIKey string `json:"api_key"` + RequestID string `json:"request_id"` + ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` + ResponseServiceTier string `json:"response_service_tier,omitempty"` } type requestDetail struct { diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 3db9189fe..682e02cd3 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -61,6 +61,8 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireStringField(t, payload, "service_tier", "auto") requireMissingField(t, payload, "request_service_tier") requireStringField(t, payload, "response_service_tier", "default") + requireIntField(t, payload, "accounting_version", coreusage.TokenAccountingSchemaVersion) + requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityUnclassified, 30) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"}) @@ -397,6 +399,38 @@ func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, w } } +func requireIntField(t *testing.T, payload map[string]json.RawMessage, key string, want int) { + t.Helper() + + raw, ok := payload[key] + if !ok { + t.Fatalf("payload missing %q", key) + } + var got int + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshal %q: %v", key, err) + } + if got != want { + t.Fatalf("%s = %d, want %d", key, got, want) + } +} + +func requireTokenBreakdown(t *testing.T, payload map[string]json.RawMessage, quality coreusage.TokenAccountingQuality, total int64) { + t.Helper() + + raw, ok := payload["token_breakdown"] + if !ok { + t.Fatal("payload missing token_breakdown") + } + var breakdown coreusage.TokenBreakdown + if err := json.Unmarshal(raw, &breakdown); err != nil { + t.Fatalf("unmarshal token_breakdown: %v", err) + } + if !breakdown.Valid() || breakdown.Quality != quality || breakdown.TotalTokens != total { + t.Fatalf("token_breakdown = %+v, want quality=%s total=%d", breakdown, quality, total) + } +} + func requireMissingField(t *testing.T, payload map[string]json.RawMessage, key string) { t.Helper() diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index c4b8d8274..d5f2309e2 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -208,13 +208,7 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det } func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail { - if detail.TotalTokens == 0 { - total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens - if total > 0 { - detail.TotalTokens = total - } - } - return detail + return usage.EnsureTokenBreakdown(detail) } func hasNonZeroTokenUsage(detail usage.Detail) bool { @@ -224,7 +218,8 @@ func hasNonZeroTokenUsage(detail usage.Detail) bool { detail.CachedTokens != 0 || detail.CacheReadTokens != 0 || detail.CacheCreationTokens != 0 || - detail.TotalTokens != 0 + detail.TotalTokens != 0 || + detail.TokenBreakdown.TotalTokens != 0 } // ensurePublished guarantees that a usage record is emitted exactly once. @@ -610,6 +605,17 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { if reasoning.Exists() { detail.ReasoningTokens = reasoning.Int() } + detail.TokenBreakdown = usage.NewSubsetTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) + if detail.TotalTokens == 0 { + detail.TotalTokens = detail.TokenBreakdown.TotalTokens + } return detail } @@ -665,6 +671,14 @@ func parseClaudeUsageNode(usageNode gjson.Result) usage.Detail { detail.CachedTokens = detail.CacheCreationTokens } detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.CacheReadTokens + detail.CacheCreationTokens + detail.TokenBreakdown = usage.NewIndependentTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) return detail } @@ -681,6 +695,14 @@ func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail { if detail.TotalTokens == 0 { detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens } + detail.TokenBreakdown = usage.NewSeparateReasoningTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) return detail } @@ -699,11 +721,16 @@ func parseInteractionsUsageDetail(node gjson.Result) usage.Detail { detail.CacheReadTokens = detail.CachedTokens } if detail.TotalTokens == 0 { - detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens + detail.CacheCreationTokens - if cacheRead.Exists() { - detail.TotalTokens += detail.CacheReadTokens - } + detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens } + detail.TokenBreakdown = usage.NewSeparateReasoningTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) return detail } diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index c54b38e4e..0f5adba6a 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -12,16 +12,16 @@ import ( ) func TestParseOpenAIUsageChatCompletions(t *testing.T) { - data := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":4},"completion_tokens_details":{"reasoning_tokens":5}}}`) + data := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":6,"total_tokens":16,"prompt_tokens_details":{"cached_tokens":4},"completion_tokens_details":{"reasoning_tokens":5}}}`) detail := ParseOpenAIUsage(data) - if detail.InputTokens != 1 { - t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 1) + if detail.InputTokens != 10 { + t.Fatalf("input tokens = %d, want %d", detail.InputTokens, 10) } - if detail.OutputTokens != 2 { - t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 2) + if detail.OutputTokens != 6 { + t.Fatalf("output tokens = %d, want %d", detail.OutputTokens, 6) } - if detail.TotalTokens != 3 { - t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 3) + if detail.TotalTokens != 16 { + t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 16) } if detail.CachedTokens != 4 { t.Fatalf("cached tokens = %d, want %d", detail.CachedTokens, 4) @@ -32,6 +32,12 @@ func TestParseOpenAIUsageChatCompletions(t *testing.T) { if detail.ReasoningTokens != 5 { t.Fatalf("reasoning tokens = %d, want %d", detail.ReasoningTokens, 5) } + if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } + if detail.TokenBreakdown.Input.UncachedTokens != 6 || detail.TokenBreakdown.Output.NonReasoningTokens != 1 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseOpenAIUsageResponses(t *testing.T) { @@ -58,6 +64,9 @@ func TestParseOpenAIUsageResponses(t *testing.T) { if detail.ResponseServiceTier != "default" { t.Fatalf("response service tier = %q, want default", detail.ResponseServiceTier) } + if detail.TokenBreakdown.Input.UncachedTokens != 3 || detail.TokenBreakdown.Output.NonReasoningTokens != 11 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) { @@ -87,6 +96,9 @@ func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) { if detail.ResponseServiceTier != "priority" { t.Fatalf("response service tier = %q, want priority", detail.ResponseServiceTier) } + if detail.TokenBreakdown.Input.UncachedTokens != 30 || detail.TokenBreakdown.Input.CacheWriteTokens != 40 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseOpenAIUsageNormalizesCacheCreationAlias(t *testing.T) { @@ -283,6 +295,9 @@ func TestParseClaudeUsageIncludesCacheTokensInTotal(t *testing.T) { if detail.TotalTokens != 22859 { t.Fatalf("total tokens = %d, want %d", detail.TotalTokens, 22859) } + if detail.TokenBreakdown.Input.TotalTokens != 22606 || detail.TokenBreakdown.Input.UncachedTokens != 3085 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseClaudeUsageFallsBackCachedTokensToCacheCreation(t *testing.T) { @@ -304,6 +319,9 @@ func TestParseGeminiUsageNormalizesCachedContent(t *testing.T) { if detail.CacheReadTokens != 4 { t.Fatalf("cache read tokens = %d, want 4", detail.CacheReadTokens) } + if detail.TokenBreakdown.Input.UncachedTokens != 6 || detail.TokenBreakdown.TotalTokens != 12 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseInteractionsUsage(t *testing.T) { @@ -326,6 +344,23 @@ func TestParseInteractionsUsage(t *testing.T) { if detail.CacheReadTokens != 2 { t.Fatalf("cache read tokens = %d, want 2", detail.CacheReadTokens) } + if detail.TokenBreakdown.Input.UncachedTokens != 1 || detail.TokenBreakdown.Output.TotalTokens != 9 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } +} + +func TestNormalizeUsageDetailTotalDoesNotDoubleCountReasoning(t *testing.T) { + detail := normalizeUsageDetailTotal(usage.Detail{ + InputTokens: 100, + OutputTokens: 30, + ReasoningTokens: 12, + }) + if detail.TotalTokens != 130 { + t.Fatalf("total tokens = %d, want 130", detail.TotalTokens) + } + if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityUnclassified || detail.TokenBreakdown.UnclassifiedTokens != 130 { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } } func TestParseInteractionsUsageNormalizesCacheWriteAlias(t *testing.T) { diff --git a/sdk/cliproxy/usage/accounting.go b/sdk/cliproxy/usage/accounting.go new file mode 100644 index 000000000..307a607e5 --- /dev/null +++ b/sdk/cliproxy/usage/accounting.go @@ -0,0 +1,239 @@ +package usage + +// TokenAccountingSchemaVersion identifies the canonical token accounting contract. +const TokenAccountingSchemaVersion = 2 + +// TokenAccountingQuality describes how confidently a token total can be classified. +type TokenAccountingQuality string + +const ( + TokenAccountingQualityComplete TokenAccountingQuality = "complete" + TokenAccountingQualityInconsistent TokenAccountingQuality = "inconsistent" + TokenAccountingQualityUnclassified TokenAccountingQuality = "unclassified" +) + +// TokenInputBreakdown contains mutually exclusive input token buckets. +type TokenInputBreakdown struct { + TotalTokens int64 `json:"total_tokens"` + UncachedTokens int64 `json:"uncached_tokens"` + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheWriteTokens int64 `json:"cache_write_tokens"` +} + +// TokenOutputBreakdown contains mutually exclusive output token buckets. +type TokenOutputBreakdown struct { + TotalTokens int64 `json:"total_tokens"` + NonReasoningTokens int64 `json:"non_reasoning_tokens"` + ReasoningTokens int64 `json:"reasoning_tokens"` +} + +// TokenBreakdown is the canonical, non-overlapping token accounting contract. +type TokenBreakdown struct { + SchemaVersion int `json:"schema_version"` + Quality TokenAccountingQuality `json:"quality"` + TotalTokens int64 `json:"total_tokens"` + Input TokenInputBreakdown `json:"input"` + Output TokenOutputBreakdown `json:"output"` + UnclassifiedTokens int64 `json:"unclassified_tokens"` +} + +// Valid reports whether the breakdown satisfies the v2 accounting invariants. +func (b TokenBreakdown) Valid() bool { + if b.SchemaVersion != TokenAccountingSchemaVersion || !validTokenAccountingQuality(b.Quality) { + return false + } + if b.TotalTokens < 0 || b.UnclassifiedTokens < 0 || + b.Input.TotalTokens < 0 || b.Input.UncachedTokens < 0 || + b.Input.CacheReadTokens < 0 || b.Input.CacheWriteTokens < 0 || + b.Output.TotalTokens < 0 || b.Output.NonReasoningTokens < 0 || + b.Output.ReasoningTokens < 0 { + return false + } + if b.Input.TotalTokens != b.Input.UncachedTokens+b.Input.CacheReadTokens+b.Input.CacheWriteTokens { + return false + } + if b.Output.TotalTokens != b.Output.NonReasoningTokens+b.Output.ReasoningTokens { + return false + } + if b.TotalTokens != b.Input.TotalTokens+b.Output.TotalTokens+b.UnclassifiedTokens { + return false + } + if b.Quality == TokenAccountingQualityComplete && b.UnclassifiedTokens != 0 { + return false + } + return true +} + +func validTokenAccountingQuality(quality TokenAccountingQuality) bool { + switch quality { + case TokenAccountingQualityComplete, TokenAccountingQualityInconsistent, TokenAccountingQualityUnclassified: + return true + default: + return false + } +} + +// NewSubsetTokenBreakdown normalizes protocols where cache tokens are included +// in input totals and reasoning tokens are included in output totals. +func NewSubsetTokenBreakdown(inputTotal, cacheRead, cacheWrite, outputTotal, reasoning, total int64) TokenBreakdown { + expectedTotal, okExpected := nonNegativeSum(inputTotal, outputTotal) + if !okExpected || cacheRead < 0 || cacheWrite < 0 || reasoning < 0 || + cacheRead+cacheWrite > inputTotal || reasoning > outputTotal { + return inconsistentTokenBreakdown(total, expectedTotal) + } + resolvedTotal, okTotal := resolveAccountingTotal(total, expectedTotal) + if !okTotal { + return inconsistentTokenBreakdown(total, expectedTotal) + } + return TokenBreakdown{ + SchemaVersion: TokenAccountingSchemaVersion, + Quality: TokenAccountingQualityComplete, + TotalTokens: resolvedTotal, + Input: TokenInputBreakdown{ + TotalTokens: inputTotal, + UncachedTokens: inputTotal - cacheRead - cacheWrite, + CacheReadTokens: cacheRead, + CacheWriteTokens: cacheWrite, + }, + Output: TokenOutputBreakdown{ + TotalTokens: outputTotal, + NonReasoningTokens: outputTotal - reasoning, + ReasoningTokens: reasoning, + }, + } +} + +// NewIndependentTokenBreakdown normalizes protocols where uncached input, +// cache reads, cache writes, non-reasoning output, and reasoning are separate. +func NewIndependentTokenBreakdown(uncachedInput, cacheRead, cacheWrite, nonReasoningOutput, reasoning, total int64) TokenBreakdown { + inputTotal, okInput := nonNegativeSum(uncachedInput, cacheRead, cacheWrite) + outputTotal, okOutput := nonNegativeSum(nonReasoningOutput, reasoning) + expectedTotal, okExpected := nonNegativeSum(inputTotal, outputTotal) + if !okInput || !okOutput || !okExpected { + return inconsistentTokenBreakdown(total, expectedTotal) + } + resolvedTotal, okTotal := resolveAccountingTotal(total, expectedTotal) + if !okTotal { + return inconsistentTokenBreakdown(total, expectedTotal) + } + return TokenBreakdown{ + SchemaVersion: TokenAccountingSchemaVersion, + Quality: TokenAccountingQualityComplete, + TotalTokens: resolvedTotal, + Input: TokenInputBreakdown{ + TotalTokens: inputTotal, + UncachedTokens: uncachedInput, + CacheReadTokens: cacheRead, + CacheWriteTokens: cacheWrite, + }, + Output: TokenOutputBreakdown{ + TotalTokens: outputTotal, + NonReasoningTokens: nonReasoningOutput, + ReasoningTokens: reasoning, + }, + } +} + +// NewSeparateReasoningTokenBreakdown normalizes protocols where cache tokens +// are included in input totals while reasoning is separate from ordinary output. +func NewSeparateReasoningTokenBreakdown(inputTotal, cacheRead, cacheWrite, nonReasoningOutput, reasoning, total int64) TokenBreakdown { + if inputTotal < 0 || cacheRead < 0 || cacheWrite < 0 || cacheRead+cacheWrite > inputTotal { + return inconsistentTokenBreakdown(total, 0) + } + outputTotal, okOutput := nonNegativeSum(nonReasoningOutput, reasoning) + expectedTotal, okExpected := nonNegativeSum(inputTotal, outputTotal) + if !okOutput || !okExpected { + return inconsistentTokenBreakdown(total, expectedTotal) + } + resolvedTotal, okTotal := resolveAccountingTotal(total, expectedTotal) + if !okTotal { + return inconsistentTokenBreakdown(total, expectedTotal) + } + return TokenBreakdown{ + SchemaVersion: TokenAccountingSchemaVersion, + Quality: TokenAccountingQualityComplete, + TotalTokens: resolvedTotal, + Input: TokenInputBreakdown{ + TotalTokens: inputTotal, + UncachedTokens: inputTotal - cacheRead - cacheWrite, + CacheReadTokens: cacheRead, + CacheWriteTokens: cacheWrite, + }, + Output: TokenOutputBreakdown{ + TotalTokens: outputTotal, + NonReasoningTokens: nonReasoningOutput, + ReasoningTokens: reasoning, + }, + } +} + +// NewUnclassifiedTokenBreakdown preserves an authoritative total without +// guessing how an unknown protocol partitions it. +func NewUnclassifiedTokenBreakdown(total int64) TokenBreakdown { + if total <= 0 { + quality := TokenAccountingQualityComplete + if total < 0 { + quality = TokenAccountingQualityInconsistent + } + return TokenBreakdown{SchemaVersion: TokenAccountingSchemaVersion, Quality: quality} + } + return TokenBreakdown{ + SchemaVersion: TokenAccountingSchemaVersion, + Quality: TokenAccountingQualityUnclassified, + TotalTokens: total, + UnclassifiedTokens: total, + } +} + +// EnsureTokenBreakdown attaches a valid v2 breakdown to legacy or direct SDK +// usage details without guessing whether reasoning is already inside output. +func EnsureTokenBreakdown(detail Detail) Detail { + if !detail.TokenBreakdown.Valid() { + total := detail.TotalTokens + if total == 0 { + total = detail.InputTokens + detail.OutputTokens + } + detail.TokenBreakdown = NewUnclassifiedTokenBreakdown(total) + } + if detail.TotalTokens == 0 { + detail.TotalTokens = detail.TokenBreakdown.TotalTokens + } + return detail +} + +func inconsistentTokenBreakdown(total, fallback int64) TokenBreakdown { + resolved := total + if resolved <= 0 { + resolved = fallback + } + if resolved < 0 { + resolved = 0 + } + return TokenBreakdown{ + SchemaVersion: TokenAccountingSchemaVersion, + Quality: TokenAccountingQualityInconsistent, + TotalTokens: resolved, + UnclassifiedTokens: resolved, + } +} + +func resolveAccountingTotal(total, expected int64) (int64, bool) { + if total < 0 || expected < 0 { + return 0, false + } + if total == 0 { + return expected, true + } + return total, total == expected +} + +func nonNegativeSum(values ...int64) (int64, bool) { + var total int64 + for _, value := range values { + if value < 0 || total > int64(^uint64(0)>>1)-value { + return 0, false + } + total += value + } + return total, true +} diff --git a/sdk/cliproxy/usage/accounting_test.go b/sdk/cliproxy/usage/accounting_test.go new file mode 100644 index 000000000..f5e7834ea --- /dev/null +++ b/sdk/cliproxy/usage/accounting_test.go @@ -0,0 +1,56 @@ +package usage + +import "testing" + +func TestNewSubsetTokenBreakdownAvoidsCacheAndReasoningDoubleCount(t *testing.T) { + breakdown := NewSubsetTokenBreakdown(100, 40, 10, 30, 12, 130) + if !breakdown.Valid() { + t.Fatalf("breakdown is invalid: %+v", breakdown) + } + if breakdown.Input.UncachedTokens != 50 || breakdown.Output.NonReasoningTokens != 18 { + t.Fatalf("breakdown = %+v", breakdown) + } + if breakdown.TotalTokens != 130 { + t.Fatalf("total = %d, want 130", breakdown.TotalTokens) + } +} + +func TestNewIndependentTokenBreakdownKeepsClaudeCacheBucketsIndependent(t *testing.T) { + breakdown := NewIndependentTokenBreakdown(30, 7, 13, 5, 0, 55) + if !breakdown.Valid() { + t.Fatalf("breakdown is invalid: %+v", breakdown) + } + if breakdown.Input.TotalTokens != 50 || breakdown.TotalTokens != 55 { + t.Fatalf("breakdown = %+v", breakdown) + } +} + +func TestNewSeparateReasoningTokenBreakdownAddsReasoningToOutput(t *testing.T) { + breakdown := NewSeparateReasoningTokenBreakdown(20, 5, 0, 7, 3, 30) + if !breakdown.Valid() { + t.Fatalf("breakdown is invalid: %+v", breakdown) + } + if breakdown.Output.TotalTokens != 10 || breakdown.TotalTokens != 30 { + t.Fatalf("breakdown = %+v", breakdown) + } +} + +func TestTokenBreakdownMarksContradictoryParentsInconsistent(t *testing.T) { + breakdown := NewSubsetTokenBreakdown(10, 4, 0, 3, 1, 20) + if !breakdown.Valid() { + t.Fatalf("breakdown is invalid: %+v", breakdown) + } + if breakdown.Quality != TokenAccountingQualityInconsistent || breakdown.UnclassifiedTokens != 20 { + t.Fatalf("breakdown = %+v", breakdown) + } +} + +func TestNewUnclassifiedTokenBreakdownDoesNotGuessBuckets(t *testing.T) { + breakdown := NewUnclassifiedTokenBreakdown(42) + if !breakdown.Valid() { + t.Fatalf("breakdown is invalid: %+v", breakdown) + } + if breakdown.Quality != TokenAccountingQualityUnclassified || breakdown.UnclassifiedTokens != 42 { + t.Fatalf("breakdown = %+v", breakdown) + } +} diff --git a/sdk/cliproxy/usage/manager.go b/sdk/cliproxy/usage/manager.go index 0c9529f19..7fa604168 100644 --- a/sdk/cliproxy/usage/manager.go +++ b/sdk/cliproxy/usage/manager.go @@ -68,6 +68,7 @@ type Detail struct { CacheReadTokens int64 CacheCreationTokens int64 TotalTokens int64 + TokenBreakdown TokenBreakdown ResponseServiceTier string } From fe8a616aa3037043241556cdea956fd26d0620a6 Mon Sep 17 00:00:00 2001 From: Dylan <1990016@gmail.com> Date: Fri, 24 Jul 2026 00:31:32 +0800 Subject: [PATCH 2/2] fix(usage): classify partial token accounting correctly --- internal/redisqueue/plugin.go | 2 +- internal/redisqueue/plugin_test.go | 34 ++++++- .../runtime/executor/helps/usage_helpers.go | 33 ++++--- .../executor/helps/usage_helpers_test.go | 19 +++- sdk/cliproxy/usage/accounting.go | 92 ++++++++++++++++++- sdk/cliproxy/usage/accounting_test.go | 65 +++++++++++++ 6 files changed, 223 insertions(+), 22 deletions(-) diff --git a/internal/redisqueue/plugin.go b/internal/redisqueue/plugin.go index 784532eb4..f43eadd26 100644 --- a/internal/redisqueue/plugin.go +++ b/internal/redisqueue/plugin.go @@ -65,7 +65,7 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec } responseServiceTier := strings.TrimSpace(record.ResponseServiceTier) - usageDetail := coreusage.EnsureTokenBreakdown(record.Detail) + usageDetail := coreusage.EnsureTokenBreakdownForProvider(record.Detail, record.Provider, record.ExecutorType) tokens := tokenStats{ InputTokens: usageDetail.InputTokens, OutputTokens: usageDetail.OutputTokens, diff --git a/internal/redisqueue/plugin_test.go b/internal/redisqueue/plugin_test.go index 682e02cd3..5f41507f2 100644 --- a/internal/redisqueue/plugin_test.go +++ b/internal/redisqueue/plugin_test.go @@ -62,7 +62,7 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { requireMissingField(t, payload, "request_service_tier") requireStringField(t, payload, "response_service_tier", "default") requireIntField(t, payload, "accounting_version", coreusage.TokenAccountingSchemaVersion) - requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityUnclassified, 30) + requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, 30) requireTokensBoolField(t, payload, "cache_read_tokens_present", true) requireHeaderField(t, payload, "response_headers", "X-Upstream-Request-Id", []string{"upstream-req-1"}) requireHeaderField(t, payload, "response_headers", "Retry-After", []string{"30"}) @@ -72,6 +72,38 @@ func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) { }) } +func TestUsageQueuePluginNormalizesDirectSDKUsageByProvider(t *testing.T) { + tests := []struct { + provider string + wantTotal int + }{ + {provider: "openai", wantTotal: 130}, + {provider: "gemini", wantTotal: 142}, + } + for _, tt := range tests { + t.Run(tt.provider, func(t *testing.T) { + withEnabledQueue(t, func() { + ctx := internallogging.WithResponseStatusHolder(context.Background()) + internallogging.SetResponseStatus(ctx, http.StatusOK) + + (&usageQueuePlugin{}).HandleUsage(ctx, coreusage.Record{ + Provider: tt.provider, + Model: "direct-sdk-model", + Detail: coreusage.Detail{ + InputTokens: 100, + OutputTokens: 30, + ReasoningTokens: 12, + }, + }) + + payload := popSinglePayload(t) + requireIntField(t, requireTokensPayload(t, payload), "total_tokens", tt.wantTotal) + requireTokenBreakdown(t, payload, coreusage.TokenAccountingQualityComplete, int64(tt.wantTotal)) + }) + }) + } +} + func TestUsageQueuePluginPayloadIncludesGenerateFalse(t *testing.T) { withEnabledQueue(t, func() { ctx := internallogging.WithResponseStatusHolder(context.Background()) diff --git a/internal/runtime/executor/helps/usage_helpers.go b/internal/runtime/executor/helps/usage_helpers.go index d5f2309e2..fc8f71afa 100644 --- a/internal/runtime/executor/helps/usage_helpers.go +++ b/internal/runtime/executor/helps/usage_helpers.go @@ -177,7 +177,7 @@ func (r *UsageReporter) buildAdditionalModelRecord(model string, detail usage.De if model == "" { return usage.Record{}, false } - detail = normalizeUsageDetailTotal(detail) + detail = normalizeUsageDetailTotal(detail, r.provider, r.executorType) if !hasNonZeroTokenUsage(detail) { return usage.Record{}, false } @@ -201,14 +201,14 @@ func (r *UsageReporter) publishWithOutcome(ctx context.Context, detail usage.Det if r == nil { return } - detail = normalizeUsageDetailTotal(detail) + detail = normalizeUsageDetailTotal(detail, r.provider, r.executorType) r.once.Do(func() { r.publishRecord(ctx, r.buildRecord(detail, failed, fail)) }) } -func normalizeUsageDetailTotal(detail usage.Detail) usage.Detail { - return usage.EnsureTokenBreakdown(detail) +func normalizeUsageDetailTotal(detail usage.Detail, provider, executorType string) usage.Detail { + return usage.EnsureTokenBreakdownForProvider(detail, provider, executorType) } func hasNonZeroTokenUsage(detail usage.Detail) bool { @@ -551,11 +551,14 @@ func hasOpenAIStyleUsageTokenFields(usageNode gjson.Result) bool { if !usageNode.Exists() || !usageNode.IsObject() { return false } + return usageNode.Get("total_tokens").Exists() || hasOpenAIStyleUsageBucketFields(usageNode) +} + +func hasOpenAIStyleUsageBucketFields(usageNode gjson.Result) bool { return usageNode.Get("prompt_tokens").Exists() || usageNode.Get("input_tokens").Exists() || usageNode.Get("completion_tokens").Exists() || usageNode.Get("output_tokens").Exists() || - usageNode.Get("total_tokens").Exists() || usageNode.Get("prompt_tokens_details.cached_tokens").Exists() || usageNode.Get("input_tokens_details.cached_tokens").Exists() || usageNode.Get("prompt_tokens_details.cache_write_tokens").Exists() || @@ -605,14 +608,18 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail { if reasoning.Exists() { detail.ReasoningTokens = reasoning.Int() } - detail.TokenBreakdown = usage.NewSubsetTokenBreakdown( - detail.InputTokens, - detail.CacheReadTokens, - detail.CacheCreationTokens, - detail.OutputTokens, - detail.ReasoningTokens, - detail.TotalTokens, - ) + if hasOpenAIStyleUsageBucketFields(usageNode) { + detail.TokenBreakdown = usage.NewSubsetTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) + } else { + detail.TokenBreakdown = usage.NewUnclassifiedTokenBreakdown(detail.TotalTokens) + } if detail.TotalTokens == 0 { detail.TotalTokens = detail.TokenBreakdown.TotalTokens } diff --git a/internal/runtime/executor/helps/usage_helpers_test.go b/internal/runtime/executor/helps/usage_helpers_test.go index 0f5adba6a..4511e033b 100644 --- a/internal/runtime/executor/helps/usage_helpers_test.go +++ b/internal/runtime/executor/helps/usage_helpers_test.go @@ -69,6 +69,21 @@ func TestParseOpenAIUsageResponses(t *testing.T) { } } +func TestParseOpenAIUsageTotalOnlyIsUnclassified(t *testing.T) { + detail := ParseOpenAIUsage([]byte(`{"usage":{"total_tokens":42}}`)) + if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityUnclassified || + detail.TotalTokens != 42 || detail.TokenBreakdown.UnclassifiedTokens != 42 { + t.Fatalf("detail = %+v", detail) + } +} + +func TestParseOpenAIUsageExplicitZeroBucketsRemainInconsistent(t *testing.T) { + detail := ParseOpenAIUsage([]byte(`{"usage":{"input_tokens":0,"output_tokens":0,"total_tokens":42}}`)) + if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != usage.TokenAccountingQualityInconsistent { + t.Fatalf("detail = %+v", detail) + } +} + func TestParseCodexUsageIncludesCacheWriteTokens(t *testing.T) { data := []byte(`{"response":{"service_tier":"priority","usage":{"input_tokens":100,"output_tokens":20,"total_tokens":120,"input_tokens_details":{"cached_tokens":30,"cache_write_tokens":40}}}}`) detail, ok := ParseCodexUsage(data) @@ -354,11 +369,11 @@ func TestNormalizeUsageDetailTotalDoesNotDoubleCountReasoning(t *testing.T) { InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, - }) + }, "openai", "") if detail.TotalTokens != 130 { t.Fatalf("total tokens = %d, want 130", detail.TotalTokens) } - if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityUnclassified || detail.TokenBreakdown.UnclassifiedTokens != 130 { + if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete || detail.TokenBreakdown.Output.ReasoningTokens != 12 { t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) } } diff --git a/sdk/cliproxy/usage/accounting.go b/sdk/cliproxy/usage/accounting.go index 307a607e5..6429cb7fc 100644 --- a/sdk/cliproxy/usage/accounting.go +++ b/sdk/cliproxy/usage/accounting.go @@ -1,5 +1,7 @@ package usage +import "strings" + // TokenAccountingSchemaVersion identifies the canonical token accounting contract. const TokenAccountingSchemaVersion = 2 @@ -12,6 +14,15 @@ const ( TokenAccountingQualityUnclassified TokenAccountingQuality = "unclassified" ) +type tokenAccountingSemantics uint8 + +const ( + tokenAccountingSemanticsUnknown tokenAccountingSemantics = iota + tokenAccountingSemanticsSubset + tokenAccountingSemanticsIndependent + tokenAccountingSemanticsSeparateReasoning +) + // TokenInputBreakdown contains mutually exclusive input token buckets. type TokenInputBreakdown struct { TotalTokens int64 `json:"total_tokens"` @@ -188,12 +199,15 @@ func NewUnclassifiedTokenBreakdown(total int64) TokenBreakdown { // EnsureTokenBreakdown attaches a valid v2 breakdown to legacy or direct SDK // usage details without guessing whether reasoning is already inside output. func EnsureTokenBreakdown(detail Detail) Detail { + return EnsureTokenBreakdownForProvider(detail, "", "") +} + +// EnsureTokenBreakdownForProvider attaches a valid v2 breakdown to legacy or +// direct SDK usage details using the known provider's token semantics. Unknown +// providers remain unclassified instead of guessing how their buckets overlap. +func EnsureTokenBreakdownForProvider(detail Detail, provider, executorType string) Detail { if !detail.TokenBreakdown.Valid() { - total := detail.TotalTokens - if total == 0 { - total = detail.InputTokens + detail.OutputTokens - } - detail.TokenBreakdown = NewUnclassifiedTokenBreakdown(total) + detail.TokenBreakdown = tokenBreakdownForProvider(detail, provider, executorType) } if detail.TotalTokens == 0 { detail.TotalTokens = detail.TokenBreakdown.TotalTokens @@ -201,6 +215,74 @@ func EnsureTokenBreakdown(detail Detail) Detail { return detail } +func tokenBreakdownForProvider(detail Detail, provider, executorType string) TokenBreakdown { + switch tokenAccountingSemanticsFor(provider, executorType) { + case tokenAccountingSemanticsSubset: + return NewSubsetTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) + case tokenAccountingSemanticsIndependent: + return NewIndependentTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) + case tokenAccountingSemanticsSeparateReasoning: + return NewSeparateReasoningTokenBreakdown( + detail.InputTokens, + detail.CacheReadTokens, + detail.CacheCreationTokens, + detail.OutputTokens, + detail.ReasoningTokens, + detail.TotalTokens, + ) + default: + total := detail.TotalTokens + if total == 0 { + var okTotal bool + total, okTotal = nonNegativeSum(detail.InputTokens, detail.OutputTokens) + if !okTotal { + return inconsistentTokenBreakdown(detail.TotalTokens, 0) + } + } + return NewUnclassifiedTokenBreakdown(total) + } +} + +func tokenAccountingSemanticsFor(provider, executorType string) tokenAccountingSemantics { + normalizedProvider := strings.ToLower(strings.TrimSpace(provider)) + normalizedExecutor := strings.ToLower(strings.TrimSpace(executorType)) + value := strings.TrimSpace(normalizedProvider + " " + normalizedExecutor) + if value == "" || value == "unknown" || value == "unknown unknown" { + return tokenAccountingSemanticsUnknown + } + if normalizedExecutor == "openaicompatexecutor" || normalizedProvider == "openai-compatibility" || strings.HasPrefix(normalizedProvider, "openai-compatible-") { + return tokenAccountingSemanticsSubset + } + if strings.Contains(value, "claude") || strings.Contains(value, "anthropic") { + return tokenAccountingSemanticsIndependent + } + for _, marker := range []string{"gemini", "aistudio", "antigravity", "vertex", "interaction"} { + if strings.Contains(value, marker) { + return tokenAccountingSemanticsSeparateReasoning + } + } + for _, marker := range []string{"openai", "codex", "xai", "grok", "kimi", "qwen", "deepseek", "openrouter"} { + if strings.Contains(value, marker) { + return tokenAccountingSemanticsSubset + } + } + return tokenAccountingSemanticsUnknown +} + func inconsistentTokenBreakdown(total, fallback int64) TokenBreakdown { resolved := total if resolved <= 0 { diff --git a/sdk/cliproxy/usage/accounting_test.go b/sdk/cliproxy/usage/accounting_test.go index f5e7834ea..3cca043bd 100644 --- a/sdk/cliproxy/usage/accounting_test.go +++ b/sdk/cliproxy/usage/accounting_test.go @@ -54,3 +54,68 @@ func TestNewUnclassifiedTokenBreakdownDoesNotGuessBuckets(t *testing.T) { t.Fatalf("breakdown = %+v", breakdown) } } + +func TestEnsureTokenBreakdownForProviderUsesKnownSemantics(t *testing.T) { + tests := []struct { + name string + provider string + executorType string + detail Detail + wantTotal int64 + wantInput int64 + wantOutput int64 + }{ + { + name: "OpenAI subsets cache and reasoning", + provider: "openai", + detail: Detail{InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, CacheReadTokens: 40, CacheCreationTokens: 10}, + wantTotal: 130, + wantInput: 100, + wantOutput: 30, + }, + { + name: "OpenAI compatible executor takes precedence", + provider: "anthropic", + executorType: "OpenAICompatExecutor", + detail: Detail{InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, CacheReadTokens: 40, CacheCreationTokens: 10}, + wantTotal: 130, + wantInput: 100, + wantOutput: 30, + }, + { + name: "Gemini keeps reasoning separate", + provider: "gemini", + detail: Detail{InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, CacheReadTokens: 40, CacheCreationTokens: 10}, + wantTotal: 142, + wantInput: 100, + wantOutput: 42, + }, + { + name: "Claude keeps cache and reasoning independent", + provider: "anthropic", + detail: Detail{InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12, CacheReadTokens: 40, CacheCreationTokens: 10}, + wantTotal: 192, + wantInput: 150, + wantOutput: 42, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + detail := EnsureTokenBreakdownForProvider(tt.detail, tt.provider, tt.executorType) + if !detail.TokenBreakdown.Valid() || detail.TokenBreakdown.Quality != TokenAccountingQualityComplete { + t.Fatalf("token breakdown = %+v", detail.TokenBreakdown) + } + if detail.TotalTokens != tt.wantTotal || detail.TokenBreakdown.TotalTokens != tt.wantTotal || + detail.TokenBreakdown.Input.TotalTokens != tt.wantInput || detail.TokenBreakdown.Output.TotalTokens != tt.wantOutput { + t.Fatalf("detail = %+v, want total=%d input=%d output=%d", detail, tt.wantTotal, tt.wantInput, tt.wantOutput) + } + }) + } +} + +func TestEnsureTokenBreakdownForUnknownProviderDoesNotGuessReasoning(t *testing.T) { + detail := EnsureTokenBreakdownForProvider(Detail{InputTokens: 100, OutputTokens: 30, ReasoningTokens: 12}, "plugin-provider", "") + if detail.TotalTokens != 130 || detail.TokenBreakdown.Quality != TokenAccountingQualityUnclassified || detail.TokenBreakdown.UnclassifiedTokens != 130 { + t.Fatalf("detail = %+v", detail) + } +}