cpa-plugin/internal/data/usage.go

// Package data 定义 cpa-ext 各模块共享且与 Provider 无关的数据事实。
package data

import (
	"errors"
	"fmt"
	"math"
	"strings"
	"time"
)

// UsageQuality 表示 Provider Token 数据经过规范化后的可信程度。
// 使用方必须检查该值，不能把所有零值都当成真实测量结果。
type UsageQuality string

const (
	UsageQualityComplete     UsageQuality = "complete"
	UsageQualityNormalized   UsageQuality = "normalized"
	UsageQualityPartial      UsageQuality = "partial"
	UsageQualityInconsistent UsageQuality = "inconsistent"
	UsageQualityUnclassified UsageQuality = "unclassified"
	UsageQualityMissing      UsageQuality = "missing"
)

// TokenUsage 保存互不重叠的 Token 分项。每个 Token 只能进入一个分项，
// 防止计价和统计重复计算。
type TokenUsage struct {
	UncachedInputTokens int64
	CacheReadTokens     int64
	CacheCreationTokens int64
	NonReasoningTokens  int64
	ReasoningTokens     int64
	UnclassifiedTokens  int64
	TotalTokens         int64
}

// Usage 是与 Provider 无关的最小用量事实。该类型不保存凭证、认证请求头、
// 请求体或响应内容。
type Usage struct {
	Provider    string
	Model       string
	RequestedAt time.Time
	Generated   bool
	Failed      bool
	Quality     UsageQuality
	Tokens      TokenUsage
}

// Validate 检查 Usage 的结构约束，但不会修改观测到的数据。
// inconsistent 类型允许保留互相矛盾的总数和分项，供后续诊断使用。
func (u Usage) Validate() error {
	if strings.TrimSpace(u.Provider) == "" {
		return errors.New("usage provider is required")
	}
	if strings.TrimSpace(u.Model) == "" {
		return errors.New("usage model is required")
	}
	if u.RequestedAt.IsZero() {
		return errors.New("usage requested time is required")
	}
	if !u.Quality.valid() {
		return fmt.Errorf("unknown usage quality %q", u.Quality)
	}
	return u.Tokens.validate(u.Quality)
}

func (q UsageQuality) valid() bool {
	switch q {
	case UsageQualityComplete,
		UsageQualityNormalized,
		UsageQualityPartial,
		UsageQualityInconsistent,
		UsageQualityUnclassified,
		UsageQualityMissing:
		return true
	default:
		return false
	}
}

func (t TokenUsage) validate(quality UsageQuality) error {
	counts := []struct {
		name  string
		value int64
	}{
		{"uncached input tokens", t.UncachedInputTokens},
		{"cache read tokens", t.CacheReadTokens},
		{"cache creation tokens", t.CacheCreationTokens},
		{"non-reasoning output tokens", t.NonReasoningTokens},
		{"reasoning tokens", t.ReasoningTokens},
		{"unclassified tokens", t.UnclassifiedTokens},
		{"total tokens", t.TotalTokens},
	}
	for _, count := range counts {
		if count.value < 0 {
			return fmt.Errorf("%s cannot be negative", count.name)
		}
	}

	sum, ok := t.bucketSum()
	if !ok {
		return errors.New("token bucket sum overflows int64")
	}

	switch quality {
	case UsageQualityMissing:
		if sum != 0 || t.TotalTokens != 0 {
			return errors.New("missing usage cannot contain token counts")
		}
	case UsageQualityComplete, UsageQualityNormalized:
		if t.UnclassifiedTokens != 0 {
			return errors.New("classified usage cannot contain unclassified tokens")
		}
		if sum != t.TotalTokens {
			return errors.New("classified token buckets do not equal total tokens")
		}
	case UsageQualityPartial, UsageQualityUnclassified:
		if sum != t.TotalTokens {
			return errors.New("token buckets do not equal total tokens")
		}
	case UsageQualityInconsistent:
		// 矛盾的计数需要作为明确的诊断事实保留下来。
	}
	return nil
}

func (t TokenUsage) bucketSum() (int64, bool) {
	values := [...]int64{
		t.UncachedInputTokens,
		t.CacheReadTokens,
		t.CacheCreationTokens,
		t.NonReasoningTokens,
		t.ReasoningTokens,
		t.UnclassifiedTokens,
	}
	var sum int64
	for _, value := range values {
		if value > math.MaxInt64-sum {
			return 0, false
		}
		sum += value
	}
	return sum, true
}