lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-09-09 01:20:13 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4528 from router-for-me/fix/issue-81-normalized-token-accounting

Browse Source 
fix(usage): add canonical token accounting v2
This commit is contained in:
Luis Pater
2026-07-24 01:04:02 +08:00
committed by GitHub
parent e5729872fa fe8a616aa3
commit 885ffac42c
7 changed files with 639 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
internal/redisqueue/plugin.go
View File

@@ -65,21 +65,16 @@ func (p *usageQueuePlugin) HandleUsage(ctx context.Context, record coreusage.Rec
}
responseServiceTier := strings.TrimSpace(record.ResponseServiceTier)
usageDetail := coreusage.EnsureTokenBreakdownForProvider(record.Detail, record.Provider, record.ExecutorType)
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 {

66
internal/redisqueue/plugin_test.go
View File

@@ -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.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"})
@@ -70,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())
@@ -397,6 +431,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()

66
internal/runtime/executor/helps/usage_helpers.go
View File

@@ -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,20 +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 {
if detail.TotalTokens == 0 {
total := detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
if total > 0 {
detail.TotalTokens = total
}
}
return detail
func normalizeUsageDetailTotal(detail usage.Detail, provider, executorType string) usage.Detail {
return usage.EnsureTokenBreakdownForProvider(detail, provider, executorType)
}
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.
@@ -556,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() ||
@@ -610,6 +608,21 @@ func parseOpenAIStyleUsageNode(usageNode gjson.Result) usage.Detail {
if reasoning.Exists() {
detail.ReasoningTokens = reasoning.Int()
}
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
}
return detail
}
@@ -665,6 +678,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 +702,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 +728,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
}

64
internal/runtime/executor/helps/usage_helpers_test.go
View File

@@ -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,24 @@ 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 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) {
@@ -87,6 +111,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 +310,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 +334,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 +359,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,
}, "openai", "")
if detail.TotalTokens != 130 {
t.Fatalf("total tokens = %d, want 130", detail.TotalTokens)
}
if detail.TokenBreakdown.Quality != usage.TokenAccountingQualityComplete || detail.TokenBreakdown.Output.ReasoningTokens != 12 {
t.Fatalf("token breakdown = %+v", detail.TokenBreakdown)
}
}
func TestParseInteractionsUsageNormalizesCacheWriteAlias(t *testing.T) {

321
sdk/cliproxy/usage/accounting.go Normal file
View File

@@ -0,0 +1,321 @@
package usage
import "strings"
// 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"
)
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"`
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 {
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() {
detail.TokenBreakdown = tokenBreakdownForProvider(detail, provider, executorType)
}
if detail.TotalTokens == 0 {
detail.TotalTokens = detail.TokenBreakdown.TotalTokens
}
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 {
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
}

121
sdk/cliproxy/usage/accounting_test.go Normal file
View File

@@ -0,0 +1,121 @@
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)
}
}
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)
}
}

1
sdk/cliproxy/usage/manager.go
View File

@@ -68,6 +68,7 @@ type Detail struct {
CacheReadTokens int64
CacheCreationTokens int64
TotalTokens int64
TokenBreakdown TokenBreakdown
ResponseServiceTier string
}