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] 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 0968606f..784532eb 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 3db9189f..682e02cd 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 c4b8d827..d5f2309e 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 c54b38e4..0f5adba6 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 00000000..307a607e --- /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 00000000..f5e7834e --- /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 0c9529f1..7fa60416 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 }