396 lines
14 KiB
Go
396 lines
14 KiB
Go
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,
|
|
},
|
|
}
|
|
}
|
|
|
|
// NewPartialSubsetTokenBreakdown preserves known subset buckets while assigning
|
|
// an authoritative remainder to the unclassified bucket.
|
|
func NewPartialSubsetTokenBreakdown(inputTotal, cacheRead, cacheWrite, outputTotal, reasoning, total int64) TokenBreakdown {
|
|
cacheTotal, okCache := nonNegativeSum(cacheRead, cacheWrite)
|
|
expectedTotal, okExpected := nonNegativeSum(inputTotal, outputTotal)
|
|
if !okCache || !okExpected || inputTotal < 0 || outputTotal < 0 || reasoning < 0 ||
|
|
cacheTotal > inputTotal || reasoning > outputTotal || total < 0 {
|
|
return inconsistentTokenBreakdown(total, expectedTotal)
|
|
}
|
|
resolvedTotal := total
|
|
if resolvedTotal == 0 {
|
|
resolvedTotal = expectedTotal
|
|
}
|
|
if resolvedTotal < expectedTotal {
|
|
return inconsistentTokenBreakdown(total, expectedTotal)
|
|
}
|
|
unclassified := resolvedTotal - expectedTotal
|
|
quality := TokenAccountingQualityComplete
|
|
if unclassified > 0 {
|
|
quality = TokenAccountingQualityUnclassified
|
|
}
|
|
return TokenBreakdown{
|
|
SchemaVersion: TokenAccountingSchemaVersion,
|
|
Quality: quality,
|
|
TotalTokens: resolvedTotal,
|
|
Input: TokenInputBreakdown{
|
|
TotalTokens: inputTotal,
|
|
UncachedTokens: inputTotal - cacheTotal,
|
|
CacheReadTokens: cacheRead,
|
|
CacheWriteTokens: cacheWrite,
|
|
},
|
|
Output: TokenOutputBreakdown{
|
|
TotalTokens: outputTotal,
|
|
NonReasoningTokens: outputTotal - reasoning,
|
|
ReasoningTokens: reasoning,
|
|
},
|
|
UnclassifiedTokens: unclassified,
|
|
}
|
|
}
|
|
|
|
// 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() {
|
|
semantics := tokenAccountingSemanticsFor(provider, executorType)
|
|
if detail.CacheReadTokens == 0 && detail.CachedTokens > 0 && detail.InputTokens == 0 &&
|
|
detail.OutputTokens == 0 && detail.ReasoningTokens == 0 && detail.CacheCreationTokens == 0 && detail.TotalTokens == 0 &&
|
|
(semantics == tokenAccountingSemanticsSubset || semantics == tokenAccountingSemanticsSeparateReasoning) {
|
|
detail.CacheReadTokens = detail.CachedTokens
|
|
}
|
|
detail.TokenBreakdown = tokenBreakdownForSemantics(detail, semantics)
|
|
}
|
|
if detail.TotalTokens == 0 {
|
|
detail.TotalTokens = detail.TokenBreakdown.TotalTokens
|
|
}
|
|
return detail
|
|
}
|
|
|
|
func tokenBreakdownForSemantics(detail Detail, semantics tokenAccountingSemantics) TokenBreakdown {
|
|
if detail.TotalTokens == 0 && detail.InputTokens == 0 && detail.OutputTokens == 0 {
|
|
if total, okTotal := unclassifiedTokenLowerBound(detail); !okTotal {
|
|
return inconsistentTokenBreakdown(detail.TotalTokens, 0)
|
|
} else if total > 0 && (semantics == tokenAccountingSemanticsUnknown ||
|
|
semantics == tokenAccountingSemanticsSubset ||
|
|
(semantics == tokenAccountingSemanticsSeparateReasoning &&
|
|
(detail.CacheReadTokens > 0 || detail.CacheCreationTokens > 0 || detail.CachedTokens > 0))) {
|
|
return NewUnclassifiedTokenBreakdown(total)
|
|
}
|
|
}
|
|
switch semantics {
|
|
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 = unclassifiedTokenLowerBound(detail)
|
|
if !okTotal {
|
|
return inconsistentTokenBreakdown(detail.TotalTokens, 0)
|
|
}
|
|
}
|
|
return NewUnclassifiedTokenBreakdown(total)
|
|
}
|
|
}
|
|
|
|
func unclassifiedTokenLowerBound(detail Detail) (int64, bool) {
|
|
cacheTokens, okCache := nonNegativeSum(detail.CacheReadTokens, detail.CacheCreationTokens)
|
|
if !okCache || detail.InputTokens < 0 || detail.OutputTokens < 0 || detail.ReasoningTokens < 0 || detail.CachedTokens < 0 {
|
|
return 0, false
|
|
}
|
|
inputTotal := detail.InputTokens
|
|
if cacheTokens > inputTotal {
|
|
inputTotal = cacheTokens
|
|
}
|
|
if detail.CachedTokens > inputTotal {
|
|
inputTotal = detail.CachedTokens
|
|
}
|
|
outputTotal := detail.OutputTokens
|
|
if detail.ReasoningTokens > outputTotal {
|
|
outputTotal = detail.ReasoningTokens
|
|
}
|
|
return nonNegativeSum(inputTotal, outputTotal)
|
|
}
|
|
|
|
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
|
|
}
|