From 6a461bf1f3e730db4670f9b69140fe8af8efb72e Mon Sep 17 00:00:00 2001 From: chuan Date: Fri, 14 Aug 2026 21:08:02 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E6=8C=81=E4=B9=85?= =?UTF-8?q?=E5=8C=96=E7=94=A8=E9=87=8F=E8=AE=A1=E8=B4=B9=E4=B8=8E=E7=AE=A1?= =?UTF-8?q?=E7=90=86=E9=9D=A2=E6=9D=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 使用 SQLite 保存关联后的请求生命周期和用量记录 - 支持长上下文阶梯价格和可配置的 Fast 计费倍率 - 添加管理接口和可配置列的用量面板 - 保留失败、取消、重试和 compact 请求，便于计费核对 --- .gitignore | 4 + AGENTS.md | 2 +- config.example.yaml | 2 +- docs/modules/billing.md | 4 +- docs/modules/pricing.md | 4 +- go.mod | 5 +- go.sum | 2 + internal/collection/record.go | 58 ++++ internal/collection/service.go | 36 ++ internal/data/usage.go | 144 ++++++++ internal/data/usage_test.go | 146 ++++++++ internal/plugin/app.go | 174 +++++++++- internal/plugin/app_test.go | 33 +- internal/plugin/config.go | 10 +- internal/plugin/management.go | 210 +++++++++++ internal/plugin/management_test.go | 385 ++++++++++++++++++++ internal/plugin/pricing_management.go | 222 ++++++++++++ internal/plugin/types.go | 78 ++++- internal/pricing/pricing.go | 163 +++++++++ internal/pricing/pricing_test.go | 91 +++++ internal/repository/sqlite_pricing.go | 178 ++++++++++ internal/repository/sqlite_pricing_test.go | 88 +++++ internal/repository/sqlite_usage.go | 309 +++++++++++++++++ internal/repository/sqlite_usage_test.go | 94 +++++ internal/web/embed.go | 12 + internal/web/ui.html | 386 +++++++++++++++++++++ 26 files changed, 2809 insertions(+), 31 deletions(-) create mode 100644 internal/collection/record.go create mode 100644 internal/collection/service.go create mode 100644 internal/data/usage.go create mode 100644 internal/data/usage_test.go create mode 100644 internal/plugin/management.go create mode 100644 internal/plugin/management_test.go create mode 100644 internal/plugin/pricing_management.go create mode 100644 internal/pricing/pricing.go create mode 100644 internal/pricing/pricing_test.go create mode 100644 internal/repository/sqlite_pricing.go create mode 100644 internal/repository/sqlite_pricing_test.go create mode 100644 internal/repository/sqlite_usage.go create mode 100644 internal/repository/sqlite_usage_test.go create mode 100644 internal/web/embed.go create mode 100644 internal/web/ui.html diff --git a/.gitignore b/.gitignore index babd493..21233d9 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,7 @@ coverage.out .cache/ /.runtime/ + +/CLIProxyAPI/ +/cpa-plugin-key-billing/ +/cpa-usage-keeper/ \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md index 76ac635..1fa1fe7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,7 +3,7 @@ 当前项目的本质是对 `CLIProxyAPI` 的核心扩展，功能边界会非常收敛 - 更好的统计、展示 -- 初期专注于codex +- 初期优先保证codex的正常使用 - 能最简单的管理多用户key ## 参考 diff --git a/config.example.yaml b/config.example.yaml index dd21674..bac75e6 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -6,4 +6,4 @@ plugins: enabled: true priority: 100 codex_only: true - + database_path: data/cpa-ext.db diff --git a/docs/modules/billing.md b/docs/modules/billing.md index a67bebc..9ef165b 100644 --- a/docs/modules/billing.md +++ b/docs/modules/billing.md @@ -271,7 +271,7 @@ Reasoning Token 用于后台明细，但若上游语义表明其已经包含在 ### 5.3 Fast / priority 2.5× -Fast/priority 是金额倍率，不是另一种余额单位： +Fast/priority 是金额倍率，不是另一种余额单位。OpenAI/Codex 通过 `service_tier=priority` 表达，Anthropic 通过 `speed=fast` 表达： ```text Fast 最终金额 = 基础金额 × 2.5 @@ -286,7 +286,7 @@ effective_service_tier = 否则 standard ``` -然后只应用一次倍率。倍率必须保存在账本价格快照中。 +两种表达先归一化为一个 `fast_requested` 事实，然后只应用一次倍率。价格配置可以关闭 Fast 加价；关闭后请求仍使用 Fast，但金额不乘倍率。倍率和是否实际应用必须保存在账本价格快照中。 Fast 规则必须是版本化价格政策的一部分，不能依赖管理员每次手工补规则。实现时需要用当前 OpenAI 官方资料再次核对支持模型和倍率。 diff --git a/docs/modules/pricing.md b/docs/modules/pricing.md index 11eee2a..e5e5956 100644 --- a/docs/modules/pricing.md +++ b/docs/modules/pricing.md @@ -221,7 +221,7 @@ Reasoning Token 若已经包含在输出中，不增加第五段价格。后台 ## 6. Fast / priority 政策 -Fast 是金额倍率，不是独立余额单位。GPT-5.6 Fast/priority 的当前产品政策保存为版本化 `5/2`，只应用一次。 +Fast 是金额倍率，不是独立余额单位。GPT-5.6 Fast/priority 的当前产品政策保存为版本化 `5/2`，只应用一次。请求可能使用 OpenAI/Codex 的 `service_tier=priority`，也可能使用 Anthropic 的 `speed=fast`。采集层必须把两种输入和最终上游请求中的 Fast 状态归一化为一个 `fast_requested` 事实。 首先合并请求和响应事实： @@ -232,7 +232,7 @@ effective_service_tier = 否则 standard ``` -然后只对 `effective_service_tier=priority` 应用一次倍率。禁止同时配置： +`effective_service_tier=priority` 或请求 `speed=fast` 都表示 Fast。价格政策中的 `fast_pricing_enabled` 决定是否应用倍率，因此 Fast 请求可以按标准价格结算。禁止同时配置： - `service_tier=priority × 2.5`；以及 - `response_service_tier=priority × 2.5` diff --git a/go.mod b/go.mod index b59271a..935e9fd 100644 --- a/go.mod +++ b/go.mod @@ -2,4 +2,7 @@ module cpa-ext go 1.24 -require gopkg.in/yaml.v3 v3.0.1 +require ( + github.com/mattn/go-sqlite3 v1.14.48 + gopkg.in/yaml.v3 v3.0.1 +) diff --git a/go.sum b/go.sum index a62c313..09d46ca 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/mattn/go-sqlite3 v1.14.48 h1:7XHIgl0a8HwOaiK4E47ozLkST78rR9+OtNGx27D/TFs= +github.com/mattn/go-sqlite3 v1.14.48/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/internal/collection/record.go b/internal/collection/record.go new file mode 100644 index 0000000..7399768 --- /dev/null +++ b/internal/collection/record.go @@ -0,0 +1,58 @@ +// Package collection 定义 CPA 用量观察及其采集服务。 +package collection + +import "time" + +// Record 是当前请求明细功能保存的用量观察。 +// 它不包含汇总结果，暂时拿不到的字段使用零值或 NULL。 +type Record struct { + RequestID string + ExecutionID string + TraceID string + RequestedAt time.Time + APIKey string + KeyAlias string + Model string + ReasoningEffort string + ServiceTier string + Speed string + Failed bool + ExecutorType string + RequestType string + Endpoint string + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CachedTokens int64 + CacheReadTokens int64 + CacheWriteTokens int64 + ReasoningTokens int64 + TTFT time.Duration + Latency time.Duration + CostMicros *int64 + PriceTier string + FastRequested bool + FastPricingApplied bool + PriceMultiplierNumerator int64 + PriceMultiplierDenominator int64 + ClientIP string + Outcome string + StatusCode int + Error string +} + +// RequestRecord 保存一次模型执行的终态，即使没有产生 Usage 也必须存在。 +type RequestRecord struct { + RequestID string + TraceID string + RequestedAt time.Time + CompletedAt time.Time + Model string + RequestedModel string + SourceFormat string + Stream bool + Outcome string + StatusCode int + Error string + Endpoint string +} diff --git a/internal/collection/service.go b/internal/collection/service.go new file mode 100644 index 0000000..04f3be3 --- /dev/null +++ b/internal/collection/service.go @@ -0,0 +1,36 @@ +package collection + +import "context" + +// Repository 是采集模块需要的最小持久化接口。 +type Repository interface { + Insert(context.Context, Record) error + UpsertRequest(context.Context, RequestRecord) error + ListRecent(context.Context, int) ([]Record, error) + Close() error +} + +// Service 负责保存观察和读取请求明细，不包含 SQL 实现。 +type Service struct { + repository Repository +} + +func NewService(repository Repository) *Service { + return &Service{repository: repository} +} + +func (s *Service) Observe(ctx context.Context, record Record) error { + return s.repository.Insert(ctx, record) +} + +func (s *Service) ObserveRequest(ctx context.Context, record RequestRecord) error { + return s.repository.UpsertRequest(ctx, record) +} + +func (s *Service) Recent(ctx context.Context, limit int) ([]Record, error) { + return s.repository.ListRecent(ctx, limit) +} + +func (s *Service) Close() error { + return s.repository.Close() +} diff --git a/internal/data/usage.go b/internal/data/usage.go new file mode 100644 index 0000000..f149772 --- /dev/null +++ b/internal/data/usage.go @@ -0,0 +1,144 @@ +// 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 +} diff --git a/internal/data/usage_test.go b/internal/data/usage_test.go new file mode 100644 index 0000000..fd7c4ea --- /dev/null +++ b/internal/data/usage_test.go @@ -0,0 +1,146 @@ +package data_test + +import ( + "math" + "testing" + "time" + + "cpa-ext/internal/data" +) + +func validUsage() data.Usage { + return data.Usage{ + Provider: "codex", + Model: "gpt-5.5", + RequestedAt: time.Date(2026, time.August, 14, 8, 0, 0, 0, time.UTC), + Generated: true, + Quality: data.UsageQualityComplete, + Tokens: data.TokenUsage{ + UncachedInputTokens: 10, + CacheReadTokens: 20, + CacheCreationTokens: 30, + NonReasoningTokens: 40, + ReasoningTokens: 50, + TotalTokens: 150, + }, + } +} + +func TestUsageValidateAcceptsCompleteBreakdown(t *testing.T) { + t.Parallel() + + if err := validUsage().Validate(); err != nil { + t.Fatalf("Validate() error = %v", err) + } +} + +func TestUsageDistinguishesMissingFromMeasuredZero(t *testing.T) { + t.Parallel() + + measured := validUsage() + measured.Tokens = data.TokenUsage{} + if err := measured.Validate(); err != nil { + t.Fatalf("measured zero Validate() error = %v", err) + } + + missing := measured + missing.Quality = data.UsageQualityMissing + if err := missing.Validate(); err != nil { + t.Fatalf("missing Validate() error = %v", err) + } + + missing.Tokens.TotalTokens = 1 + if err := missing.Validate(); err == nil { + t.Fatal("missing usage with tokens unexpectedly passed validation") + } +} + +func TestUsageValidateRejectsInvalidIdentityAndQuality(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mutate func(*data.Usage) + }{ + {"missing provider", func(u *data.Usage) { u.Provider = " " }}, + {"missing model", func(u *data.Usage) { u.Model = "" }}, + {"missing requested time", func(u *data.Usage) { u.RequestedAt = time.Time{} }}, + {"unknown quality", func(u *data.Usage) { u.Quality = "other" }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + usage := validUsage() + test.mutate(&usage) + if err := usage.Validate(); err == nil { + t.Fatal("Validate() unexpectedly succeeded") + } + }) + } +} + +func TestUsageValidateRejectsInvalidTokenBreakdown(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + mutate func(*data.Usage) + }{ + {"negative token", func(u *data.Usage) { u.Tokens.CacheReadTokens = -1 }}, + {"wrong total", func(u *data.Usage) { u.Tokens.TotalTokens++ }}, + {"unclassified complete", func(u *data.Usage) { + u.Tokens.UnclassifiedTokens = 1 + u.Tokens.TotalTokens++ + }}, + {"overflow", func(u *data.Usage) { + u.Tokens = data.TokenUsage{ + UncachedInputTokens: math.MaxInt64, + CacheReadTokens: 1, + TotalTokens: math.MaxInt64, + } + }}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + usage := validUsage() + test.mutate(&usage) + if err := usage.Validate(); err == nil { + t.Fatal("Validate() unexpectedly succeeded") + } + }) + } +} + +func TestUsageValidatePreservesExplicitQualityStates(t *testing.T) { + t.Parallel() + + normalized := validUsage() + normalized.Quality = data.UsageQualityNormalized + if err := normalized.Validate(); err != nil { + t.Fatalf("normalized Validate() error = %v", err) + } + + partial := validUsage() + partial.Quality = data.UsageQualityPartial + if err := partial.Validate(); err != nil { + t.Fatalf("partial Validate() error = %v", err) + } + + unclassified := validUsage() + unclassified.Quality = data.UsageQualityUnclassified + unclassified.Tokens.UnclassifiedTokens = 5 + unclassified.Tokens.TotalTokens += 5 + if err := unclassified.Validate(); err != nil { + t.Fatalf("unclassified Validate() error = %v", err) + } + + inconsistent := validUsage() + inconsistent.Quality = data.UsageQualityInconsistent + inconsistent.Tokens.TotalTokens++ + if err := inconsistent.Validate(); err != nil { + t.Fatalf("inconsistent Validate() error = %v", err) + } +} diff --git a/internal/plugin/app.go b/internal/plugin/app.go index 28a5dea..5afd99f 100644 --- a/internal/plugin/app.go +++ b/internal/plugin/app.go @@ -1,25 +1,31 @@ package plugin import ( + "context" "encoding/json" "fmt" "net/http" + "strings" "sync" "sync/atomic" + + "cpa-ext/internal/collection" + "cpa-ext/internal/pricing" + "cpa-ext/internal/repository" ) type App struct { - config atomic.Pointer[Config] - mu sync.Mutex + mu sync.RWMutex + config Config + usage *collection.Service + store *repository.SQLiteUsageRepository + prices map[string]pricing.Policy closed bool seen atomic.Uint64 } func NewApp() *App { - a := &App{} - cfg := defaultConfig() - a.config.Store(&cfg) - return a + return &App{config: defaultConfig(), prices: make(map[string]pricing.Policy)} } func (a *App) HandleMethod(method string, request []byte) (response []byte, err error) { @@ -35,6 +41,12 @@ func (a *App) HandleMethod(method string, request []byte) (response []byte, err return a.configure(request) case MethodUsageHandle: return a.handleUsage(request) + case MethodRequestComplete: + return a.handleRequestComplete(request) + case MethodManagementRegister: + return OKEnvelope(managementRegistration()) + case MethodManagementHandle: + return a.handleManagement(request) case MethodPluginShutdown: a.Shutdown() return OKEnvelope(struct{}{}) @@ -58,13 +70,35 @@ func (a *App) configure(raw []byte) ([]byte, error) { if err != nil { return nil, err } + usageRepository, err := repository.OpenSQLiteUsage(cfg.DatabasePath) + if err != nil { + return nil, err + } + nextUsage := collection.NewService(usageRepository) + policies, err := usageRepository.ListPrices(context.Background()) + if err != nil { + _ = nextUsage.Close() + return nil, err + } + nextPrices := make(map[string]pricing.Policy, len(policies)) + for _, policy := range policies { + nextPrices[policy.Model] = policy + } a.mu.Lock() defer a.mu.Unlock() if a.closed { + _ = nextUsage.Close() return nil, fmt.Errorf("插件已经关闭") } - a.config.Store(&cfg) + previousUsage := a.usage + a.config = cfg + a.usage = nextUsage + a.store = usageRepository + a.prices = nextPrices + if previousUsage != nil { + _ = previousUsage.Close() + } return OKEnvelope(registration(negotiated)) } @@ -79,9 +113,10 @@ func registration(schemaVersion uint32) Registration { ConfigFields: []ConfigField{ {Name: "enabled", Type: "boolean", Description: "启用 CPA 扩展。"}, {Name: "codex_only", Type: "boolean", Description: "只接收 Codex/OpenAI 模型的用量事件。"}, + {Name: "database_path", Type: "string", Description: "SQLite 数据库文件路径。"}, }, }, - Capabilities: Capabilities{UsagePlugin: true}, + Capabilities: Capabilities{RequestLifecyclePlugin: true, UsagePlugin: true, ManagementAPI: true}, } } @@ -90,21 +125,134 @@ func (a *App) handleUsage(raw []byte) ([]byte, error) { if err := json.Unmarshal(raw, &record); err != nil { return nil, fmt.Errorf("解析用量事件: %w", err) } - cfg := a.config.Load() - if cfg != nil && cfg.accepts(record) { - // The first milestone only proves ingestion. Persistence and aggregation - // belong in a separate package added behind this boundary. + a.mu.RLock() + defer a.mu.RUnlock() + if a.closed { + return nil, fmt.Errorf("插件已经关闭") + } + if a.config.accepts(record) { + if a.usage == nil { + return nil, fmt.Errorf("用量数据库尚未初始化") + } + // CPA wire 类型只存在于适配层，采集模块接收与协议无关的观察值。 + observed := collection.Record{ + RequestID: record.RequestID, + ExecutionID: record.ExecutionID, + TraceID: record.TraceID, + RequestedAt: record.RequestedAt, + APIKey: record.APIKey, + Model: record.Model, + ReasoningEffort: record.ReasoningEffort, + ServiceTier: record.ServiceTier, + Speed: record.Speed, + Failed: record.Failed, + ExecutorType: record.ExecutorType, + RequestType: usageRequestType(record.Endpoint), + Endpoint: record.Endpoint, + InputTokens: record.Detail.InputTokens, + OutputTokens: record.Detail.OutputTokens, + TotalTokens: record.Detail.TotalTokens, + CachedTokens: record.Detail.CachedTokens, + CacheReadTokens: record.Detail.CacheReadTokens, + CacheWriteTokens: record.Detail.CacheCreationTokens, + ReasoningTokens: record.Detail.ReasoningTokens, + TTFT: record.TTFT, + Latency: record.Latency, + ClientIP: record.ClientIP, + } + if policy, found := a.prices[strings.TrimSpace(record.Model)]; found { + cost, calculateErr := pricing.Calculate(policy, pricing.Usage{ + InputTokens: record.Detail.InputTokens, CacheReadTokens: record.Detail.CacheReadTokens, + CacheWriteTokens: record.Detail.CacheCreationTokens, OutputTokens: record.Detail.OutputTokens, + ServiceTier: record.ServiceTier, Speed: record.Speed, + }) + if calculateErr == nil { + observed.CostMicros = &cost.CostMicros + observed.PriceTier = cost.PriceTier + observed.FastRequested = cost.FastRequested + observed.FastPricingApplied = cost.FastApplied + observed.PriceMultiplierNumerator = cost.MultiplierNumerator + observed.PriceMultiplierDenominator = cost.MultiplierDenominator + } + } + if err := a.usage.Observe(context.Background(), observed); err != nil { + return nil, err + } a.seen.Add(1) } return OKEnvelope(struct{}{}) } +func (a *App) handleRequestComplete(raw []byte) ([]byte, error) { + var completion RequestCompletion + if err := json.Unmarshal(raw, &completion); err != nil { + return nil, fmt.Errorf("解析请求终态: %w", err) + } + a.mu.RLock() + defer a.mu.RUnlock() + if a.closed { + return nil, fmt.Errorf("插件已经关闭") + } + if a.usage == nil { + return nil, fmt.Errorf("用量数据库尚未初始化") + } + endpoint, _ := completion.Metadata["request_path"].(string) + if err := a.usage.ObserveRequest(context.Background(), collection.RequestRecord{ + RequestID: completion.RequestID, TraceID: completion.TraceID, + RequestedAt: completion.StartedAt, CompletedAt: completion.CompletedAt, + Model: completion.Model, RequestedModel: completion.RequestedModel, + SourceFormat: completion.SourceFormat, Stream: completion.Stream, + Outcome: string(completion.Outcome), StatusCode: completion.StatusCode, + Error: completion.Error, Endpoint: strings.TrimSpace(endpoint), + }); err != nil { + return nil, err + } + return OKEnvelope(struct{}{}) +} + +func usageRequestType(endpoint string) string { + if isCompactEndpoint(endpoint) { + return "JSON" + } + parts := strings.Fields(endpoint) + if len(parts) == 0 { + return "" + } + switch strings.ToUpper(parts[0]) { + case http.MethodPost: + return "SSE" + case http.MethodGet: + return "WS" + default: + return "" + } +} + +func isCompactEndpoint(endpoint string) bool { + path := strings.ToLower(strings.TrimSpace(endpoint)) + if fields := strings.Fields(path); len(fields) > 1 { + path = fields[len(fields)-1] + } + if query := strings.IndexByte(path, '?'); query >= 0 { + path = path[:query] + } + return strings.HasSuffix(strings.TrimRight(path, "/"), "/responses/compact") +} + func (a *App) Seen() uint64 { return a.seen.Load() } func (a *App) Shutdown() { a.mu.Lock() + defer a.mu.Unlock() + if a.closed { + return + } a.closed = true - a.mu.Unlock() + if a.usage != nil { + _ = a.usage.Close() + a.usage = nil + a.store = nil + } } diff --git a/internal/plugin/app_test.go b/internal/plugin/app_test.go index 9d12a34..0069190 100644 --- a/internal/plugin/app_test.go +++ b/internal/plugin/app_test.go @@ -3,6 +3,8 @@ package plugin import ( "encoding/base64" "encoding/json" + "fmt" + "path/filepath" "sync" "testing" ) @@ -16,9 +18,14 @@ func lifecycleRequest(t *testing.T, schema uint32, config string) []byte { return raw } +func testConfig(t *testing.T, config string) string { + t.Helper() + return fmt.Sprintf("database_path: %q\n%s", filepath.Join(t.TempDir(), "usage.db"), config) +} + func TestRegisterNegotiatesSchemaAndDeclaresOnlyUsage(t *testing.T) { app := NewApp() - raw, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, 99, "enabled: true\ncodex_only: true\n")) + raw, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, 99, testConfig(t, "enabled: true\ncodex_only: true\n"))) if err != nil { t.Fatal(err) } @@ -33,7 +40,7 @@ func TestRegisterNegotiatesSchemaAndDeclaresOnlyUsage(t *testing.T) { if err := json.Unmarshal(env.Result, &got); err != nil { t.Fatal(err) } - if got.SchemaVersion != SchemaVersion || !got.Capabilities.UsagePlugin { + if got.SchemaVersion != SchemaVersion || !got.Capabilities.RequestLifecyclePlugin || !got.Capabilities.UsagePlugin || !got.Capabilities.ManagementAPI { t.Fatalf("unexpected registration: %+v", got) } } @@ -53,7 +60,7 @@ func TestLifecycleConfigYAMLUsesBase64WireEncoding(t *testing.T) { func TestReconfigureIsAtomicAndFiltersUsage(t *testing.T) { app := NewApp() - if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, "enabled: true\ncodex_only: true\n")); err != nil { + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil { t.Fatal(err) } for _, record := range []UsageRecord{{Provider: "codex", Model: "gpt-5.5"}, {Provider: "gemini", Model: "gemini-pro"}} { @@ -77,6 +84,9 @@ func TestReconfigureIsAtomicAndFiltersUsage(t *testing.T) { func TestConcurrentUsage(t *testing.T) { app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil { + t.Fatal(err) + } raw, _ := json.Marshal(UsageRecord{Provider: "codex", Model: "gpt-5.5"}) var wg sync.WaitGroup for range 100 { @@ -107,3 +117,20 @@ func TestUnknownMethodReturnsErrorEnvelope(t *testing.T) { t.Fatalf("unexpected envelope: %s", raw) } } + +func TestUsageRequestTypeMatchesKeeperEndpointSemantics(t *testing.T) { + tests := []struct { + endpoint string + want string + }{ + {endpoint: "POST /v1/responses", want: "SSE"}, + {endpoint: "POST /v1/chat/completions", want: "SSE"}, + {endpoint: "GET /v1/responses", want: "WS"}, + {endpoint: "/v1/responses", want: ""}, + } + for _, test := range tests { + if got := usageRequestType(test.endpoint); got != test.want { + t.Fatalf("usageRequestType(%q) = %q, want %q", test.endpoint, got, test.want) + } + } +} diff --git a/internal/plugin/config.go b/internal/plugin/config.go index 2764268..2ba0e9a 100644 --- a/internal/plugin/config.go +++ b/internal/plugin/config.go @@ -8,12 +8,13 @@ import ( ) type Config struct { - Enabled bool `yaml:"enabled"` - CodexOnly bool `yaml:"codex_only"` + Enabled bool `yaml:"enabled"` + CodexOnly bool `yaml:"codex_only"` + DatabasePath string `yaml:"database_path"` } func defaultConfig() Config { - return Config{Enabled: true, CodexOnly: true} + return Config{Enabled: true, CodexOnly: true, DatabasePath: "data/cpa-ext.db"} } func decodeConfig(raw []byte) (Config, error) { @@ -24,6 +25,9 @@ func decodeConfig(raw []byte) (Config, error) { if err := yaml.Unmarshal(raw, &cfg); err != nil { return Config{}, fmt.Errorf("解析插件配置: %w", err) } + if strings.TrimSpace(cfg.DatabasePath) == "" { + return Config{}, fmt.Errorf("database_path 不能为空") + } return cfg, nil } diff --git a/internal/plugin/management.go b/internal/plugin/management.go new file mode 100644 index 0000000..33ab81d --- /dev/null +++ b/internal/plugin/management.go @@ -0,0 +1,210 @@ +package plugin + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "cpa-ext/internal/web" +) + +const ( + managementBase = "/v0/management/plugins/" + PluginName + resourceBase = "/v0/resource/plugins/" + PluginName + routeUsage = "/usage" + routePrices = "/prices" + resourceUI = "/ui" +) + +func managementRegistration() ManagementRegistrationResponse { + return ManagementRegistrationResponse{ + Routes: []ManagementRoute{ + {Method: http.MethodGet, Path: managementBase + routeUsage, Description: "查看最近的用量记录。"}, + {Method: http.MethodGet, Path: managementBase + routePrices, Description: "查看模型价格。"}, + {Method: http.MethodPut, Path: managementBase + routePrices, Description: "保存模型价格。"}, + {Method: http.MethodDelete, Path: managementBase + routePrices, Description: "删除模型价格。"}, + }, + Resources: []ResourceRoute{ + {Path: resourceBase + resourceUI, Menu: "用量记录", Description: "查看 CPA 最近收到的请求用量。"}, + }, + } +} + +func (a *App) handleManagement(raw []byte) ([]byte, error) { + var req ManagementRequest + if err := json.Unmarshal(raw, &req); err != nil { + return nil, fmt.Errorf("解析管理请求: %w", err) + } + path := strings.TrimRight(req.Path, "/") + if req.Method == http.MethodGet && path == resourceBase+resourceUI { + return OKEnvelope(ManagementResponse{ + StatusCode: http.StatusOK, + Headers: http.Header{ + "Content-Type": []string{"text/html; charset=utf-8"}, + "Cache-Control": []string{"no-store"}, + }, + Body: web.UI(), + }) + } + if req.Method == http.MethodGet && path == managementBase+routeUsage { + return OKEnvelope(a.usageResponse()) + } + if path == managementBase+routePrices { + switch req.Method { + case http.MethodGet: + return OKEnvelope(a.listPrices()) + case http.MethodPut: + return OKEnvelope(a.putPrice(req.Body)) + case http.MethodDelete: + return OKEnvelope(a.deletePrice(req.Body)) + } + } + return OKEnvelope(jsonManagementResponse(http.StatusNotFound, map[string]any{ + "error": map[string]string{ + "code": "not_found", + "message": "管理路由不存在: " + req.Method + " " + req.Path, + }, + })) +} + +type usageListResponse struct { + Records []usageListItem `json:"records"` + Retained int `json:"retained"` +} + +// usageListItem 是请求明细表的稳定接口。当前回调拿不到的字段保留为空值。 +type usageListItem struct { + RequestID string `json:"request_id"` + ExecutionID string `json:"execution_id"` + TraceID string `json:"trace_id"` + RequestedAt time.Time `json:"requested_at"` + APIKey string `json:"api_key"` + KeyAlias string `json:"key_alias"` + Model string `json:"model"` + ReasoningEffort string `json:"reasoning_effort"` + ServiceTier string `json:"service_tier"` + Speed string `json:"speed"` + Failed bool `json:"failed"` + Outcome string `json:"outcome"` + StatusCode int `json:"status_code"` + Error string `json:"error"` + ExecutorType string `json:"executor_type"` + RequestType string `json:"request_type"` + Endpoint string `json:"endpoint"` + TTFTMilliseconds int64 `json:"ttft_ms"` + SpeedTPS *float64 `json:"speed_tps"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + ReasoningTokens int64 `json:"reasoning_tokens"` + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheWriteTokens int64 `json:"cache_write_tokens"` + CacheRate *float64 `json:"cache_rate"` + TotalTokens int64 `json:"total_tokens"` + CostUSD *float64 `json:"cost_usd"` + CostAvailable bool `json:"cost_available"` + PriceTier string `json:"price_tier"` + FastRequested bool `json:"fast_requested"` + FastPricingApplied bool `json:"fast_pricing_applied"` + ClientIP string `json:"client_ip"` +} + +func (a *App) usageResponse() ManagementResponse { + a.mu.RLock() + defer a.mu.RUnlock() + if a.usage == nil { + return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{ + "error": map[string]string{"code": "database_unavailable", "message": "用量数据库尚未初始化"}, + }) + } + records, err := a.usage.Recent(context.Background(), 1000) + if err != nil { + return jsonManagementResponse(http.StatusInternalServerError, map[string]any{ + "error": map[string]string{"code": "database_error", "message": err.Error()}, + }) + } + items := make([]usageListItem, 0, len(records)) + for _, record := range records { + ttftMilliseconds := record.TTFT.Milliseconds() + speedTPS := usageSpeed(record.OutputTokens, record.TTFT, record.Latency) + if isCompactEndpoint(record.Endpoint) { + // Compact 返回的是一次性 JSON，首字延迟和生成速度没有可比较的含义。 + ttftMilliseconds = 0 + speedTPS = nil + } + items = append(items, usageListItem{ + RequestID: record.RequestID, + ExecutionID: record.ExecutionID, + TraceID: record.TraceID, + RequestedAt: record.RequestedAt, + APIKey: record.APIKey, + KeyAlias: record.KeyAlias, + Model: record.Model, + ReasoningEffort: record.ReasoningEffort, + ServiceTier: record.ServiceTier, + Speed: record.Speed, + Failed: record.Failed, + Outcome: record.Outcome, + StatusCode: record.StatusCode, + Error: record.Error, + ExecutorType: record.ExecutorType, + RequestType: record.RequestType, + Endpoint: record.Endpoint, + TTFTMilliseconds: ttftMilliseconds, + SpeedTPS: speedTPS, + InputTokens: record.InputTokens, + OutputTokens: record.OutputTokens, + ReasoningTokens: record.ReasoningTokens, + CacheReadTokens: record.CacheReadTokens, + CacheWriteTokens: record.CacheWriteTokens, + CacheRate: usageCacheRate(record.CacheReadTokens, record.InputTokens), + TotalTokens: record.TotalTokens, + CostUSD: microsToUSD(record.CostMicros), + CostAvailable: record.CostMicros != nil, + PriceTier: record.PriceTier, + FastRequested: record.FastRequested, + FastPricingApplied: record.FastPricingApplied, + ClientIP: record.ClientIP, + }) + } + return jsonManagementResponse(http.StatusOK, usageListResponse{Records: items, Retained: 1000}) +} + +func microsToUSD(micros *int64) *float64 { + if micros == nil { + return nil + } + value := float64(*micros) / 1_000_000 + return &value +} + +func usageSpeed(outputTokens int64, ttft, latency time.Duration) *float64 { + if outputTokens <= 0 || ttft <= 0 || latency <= ttft { + return nil + } + value := float64(outputTokens) / (latency - ttft).Seconds() + return &value +} + +func usageCacheRate(cacheReadTokens, inputTokens int64) *float64 { + if inputTokens <= 0 { + return nil + } + value := float64(cacheReadTokens) / float64(inputTokens) * 100 + return &value +} + +func jsonManagementResponse(status int, value any) ManagementResponse { + body, err := json.Marshal(value) + if err != nil { + body = []byte(`{"error":{"code":"marshal_error","message":"无法生成响应"}}`) + status = http.StatusInternalServerError + } + return ManagementResponse{ + StatusCode: status, + Headers: http.Header{"Content-Type": []string{"application/json; charset=utf-8"}}, + Body: body, + } +} diff --git a/internal/plugin/management_test.go b/internal/plugin/management_test.go new file mode 100644 index 0000000..fad760c --- /dev/null +++ b/internal/plugin/management_test.go @@ -0,0 +1,385 @@ +package plugin + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "testing" + "time" +) + +func managementCall(t *testing.T, app *App, method, path string) ManagementResponse { + return managementCallBody(t, app, method, path, nil) +} + +func managementCallBody(t *testing.T, app *App, method, path string, body []byte) ManagementResponse { + t.Helper() + request, err := json.Marshal(ManagementRequest{Method: method, Path: path, Body: body}) + if err != nil { + t.Fatal(err) + } + raw, err := app.HandleMethod(MethodManagementHandle, request) + if err != nil { + t.Fatal(err) + } + var envelope Envelope + if err := json.Unmarshal(raw, &envelope); err != nil { + t.Fatal(err) + } + var response ManagementResponse + if err := json.Unmarshal(envelope.Result, &response); err != nil { + t.Fatal(err) + } + return response +} + +func TestManagementRegistrationDeclaresUsageAPIAndUI(t *testing.T) { + raw, err := NewApp().HandleMethod(MethodManagementRegister, nil) + if err != nil { + t.Fatal(err) + } + var envelope Envelope + if err := json.Unmarshal(raw, &envelope); err != nil { + t.Fatal(err) + } + var registration ManagementRegistrationResponse + if err := json.Unmarshal(envelope.Result, ®istration); err != nil { + t.Fatal(err) + } + if len(registration.Routes) != 4 || registration.Routes[0].Path != managementBase+routeUsage || registration.Routes[1].Path != managementBase+routePrices { + t.Fatalf("unexpected management routes: %+v", registration.Routes) + } + if len(registration.Resources) != 1 || registration.Resources[0].Path != resourceBase+resourceUI { + t.Fatalf("unexpected resource routes: %+v", registration.Resources) + } +} + +func TestPriceManagementCalculatesUsageWithLongContextAndFastPolicy(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil { + t.Fatal(err) + } + priceBody := []byte(`{ + "model":"gpt-5.6-sol", + "base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"}, + "long_context":{"threshold_input_tokens":272000,"comparison":"gt","input_per_1m":"5","cache_read_per_1m":"0.5","cache_write_per_1m":"6.25","output_per_1m":"22.5"}, + "fast_pricing_enabled":true, + "fast_multiplier":"2.5" + }`) + response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody) + if response.StatusCode != http.StatusOK { + t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body) + } + + record := UsageRecord{ + Provider: "codex", Model: "gpt-5.6-sol", ServiceTier: "priority", RequestedAt: time.Now(), + Detail: UsageDetail{InputTokens: 272_001, OutputTokens: 10_000, TotalTokens: 282_001}, + } + raw, _ := json.Marshal(record) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + usageResponse := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(usageResponse.Body, &payload); err != nil { + t.Fatal(err) + } + got := payload.Records[0] + if !got.CostAvailable || got.CostUSD == nil || *got.CostUSD != 3.962513 || got.PriceTier != "long_context" || !got.FastRequested || !got.FastPricingApplied { + t.Fatalf("unexpected priced usage: %+v", got) + } + + pricesResponse := managementCall(t, app, http.MethodGet, managementBase+routePrices) + if pricesResponse.StatusCode != http.StatusOK || !strings.Contains(string(pricesResponse.Body), `"fast_multiplier":"2.5"`) { + t.Fatalf("unexpected prices response: %d %s", pricesResponse.StatusCode, pricesResponse.Body) + } +} + +func TestFastRequestCanUseStandardPricing(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil { + t.Fatal(err) + } + priceBody := []byte(`{"model":"gpt-5.6-sol","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":false,"fast_multiplier":"2.5"}`) + if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody); response.StatusCode != http.StatusOK { + t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body) + } + record := UsageRecord{Provider: "codex", Model: "gpt-5.6-sol", Speed: "fast", RequestedAt: time.Now(), Detail: UsageDetail{InputTokens: 100_000, TotalTokens: 100_000}} + raw, _ := json.Marshal(record) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + got := payload.Records[0] + if got.CostUSD == nil || *got.CostUSD != 0.25 || !got.FastRequested || got.FastPricingApplied { + t.Fatalf("unexpected standard-priced Fast usage: %+v", got) + } +} + +func TestCanceledRequestWithoutUsageRemainsVisible(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { + t.Fatal(err) + } + startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC) + completion := RequestCompletion{ + RequestID: "request-canceled", TraceID: "trace-canceled", SourceFormat: "openai-response", + Model: "deepseek-v4-flash", RequestedModel: "deepseek-flash", Stream: true, + Outcome: RequestCompletionCanceled, StatusCode: 499, Error: "context canceled", + StartedAt: startedAt, CompletedAt: startedAt.Add(2 * time.Second), + Metadata: map[string]any{"request_path": "/v1/responses"}, + } + raw, _ := json.Marshal(completion) + if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 1 { + t.Fatalf("records = %d, want 1", len(payload.Records)) + } + got := payload.Records[0] + if got.RequestID != "request-canceled" || got.TraceID != "trace-canceled" || got.Outcome != "canceled" || !got.Failed || got.StatusCode != 499 || got.RequestType != "SSE" || got.CostAvailable { + t.Fatalf("unexpected canceled request: %+v", got) + } +} + +func TestCompactUsageIsPricedAndUsesJSONMetrics(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { + t.Fatal(err) + } + priceBody := []byte(`{"model":"deepseek-v4-flash","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":true,"fast_multiplier":"2.5"}`) + if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody); response.StatusCode != http.StatusOK { + t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body) + } + record := UsageRecord{ + RequestID: "compact-success", ExecutionID: "compact-attempt", Provider: "openai", + Model: "deepseek-v4-flash", Endpoint: "POST /v1/responses/compact", RequestedAt: time.Now(), + Latency: time.Second, TTFT: 200 * time.Millisecond, + Detail: UsageDetail{InputTokens: 1_000, OutputTokens: 100, TotalTokens: 1_100}, + } + raw, _ := json.Marshal(record) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 1 { + t.Fatalf("records = %d, want 1", len(payload.Records)) + } + got := payload.Records[0] + if got.RequestType != "JSON" || got.Endpoint != "POST /v1/responses/compact" || got.TTFTMilliseconds != 0 || got.SpeedTPS != nil { + t.Fatalf("unexpected compact metadata: %+v", got) + } + if !got.CostAvailable || got.CostUSD == nil || *got.CostUSD != 0.004 { + t.Fatalf("unexpected compact cost: %+v", got) + } +} + +func TestCompactFailureAndCancellationRemainVisible(t *testing.T) { + for _, test := range []struct { + name string + outcome RequestCompletionOutcome + statusCode int + }{ + {name: "failed", outcome: RequestCompletionFailed, statusCode: 500}, + {name: "canceled", outcome: RequestCompletionCanceled, statusCode: 499}, + } { + t.Run(test.name, func(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { + t.Fatal(err) + } + startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC) + completion := RequestCompletion{ + RequestID: "compact-" + test.name, Model: "deepseek-v4-flash", Stream: false, + Outcome: test.outcome, StatusCode: test.statusCode, Error: test.name, + StartedAt: startedAt, CompletedAt: startedAt.Add(time.Second), + Metadata: map[string]any{"request_path": "/v1/responses/compact"}, + } + raw, _ := json.Marshal(completion) + if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 1 { + t.Fatalf("records = %d, want 1", len(payload.Records)) + } + got := payload.Records[0] + if got.Outcome != string(test.outcome) || !got.Failed || got.StatusCode != test.statusCode || got.RequestType != "JSON" || got.Endpoint != "/v1/responses/compact" || got.SpeedTPS != nil || got.CostAvailable { + t.Fatalf("unexpected compact terminal record: %+v", got) + } + }) + } +} + +func TestUsageAndLifecycleMergeInEitherOrderAndDeduplicateExecution(t *testing.T) { + for _, usageFirst := range []bool{true, false} { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { + t.Fatal(err) + } + startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC) + usage := UsageRecord{ + RequestID: "request-1", ExecutionID: "attempt-1", TraceID: "trace-1", + Provider: "openai", Model: "deepseek-v4-flash", RequestedAt: startedAt, + Endpoint: "POST /v1/responses", Detail: UsageDetail{InputTokens: 10, TotalTokens: 10}, + } + completion := RequestCompletion{ + RequestID: "request-1", TraceID: "trace-1", Model: "deepseek-v4-flash", + Outcome: RequestCompletionFailed, StatusCode: 500, Error: "upstream failed", + StartedAt: startedAt, CompletedAt: startedAt.Add(time.Second), + } + usageRaw, _ := json.Marshal(usage) + completionRaw, _ := json.Marshal(completion) + calls := []struct { + method string + raw []byte + }{{MethodUsageHandle, usageRaw}, {MethodRequestComplete, completionRaw}} + if !usageFirst { + calls[0], calls[1] = calls[1], calls[0] + } + for _, call := range calls { + if _, err := app.HandleMethod(call.method, call.raw); err != nil { + t.Fatal(err) + } + } + // 重复的 Usage 回调不能生成第二条执行记录。 + if _, err := app.HandleMethod(MethodUsageHandle, usageRaw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 1 || payload.Records[0].ExecutionID != "attempt-1" || payload.Records[0].Outcome != "failed" || payload.Records[0].StatusCode != 500 { + t.Fatalf("usageFirst=%v records=%+v", usageFirst, payload.Records) + } + } +} + +func TestRetriedExecutionsRemainSeparate(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { + t.Fatal(err) + } + startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC) + for index, failed := range []bool{true, false} { + usage := UsageRecord{ + RequestID: "request-retry", ExecutionID: fmt.Sprintf("attempt-%d", index+1), + Provider: "openai", Model: "deepseek-v4-flash", RequestedAt: startedAt.Add(time.Duration(index) * time.Second), + Failed: failed, Detail: UsageDetail{InputTokens: int64(10 + index), TotalTokens: int64(10 + index)}, + } + raw, _ := json.Marshal(usage) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + } + completion := RequestCompletion{RequestID: "request-retry", Outcome: RequestCompletionSucceeded, StartedAt: startedAt, CompletedAt: startedAt.Add(2 * time.Second)} + raw, _ := json.Marshal(completion) + if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil { + t.Fatal(err) + } + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 2 || payload.Records[0].ExecutionID != "attempt-2" || payload.Records[0].Failed || payload.Records[1].ExecutionID != "attempt-1" || !payload.Records[1].Failed { + t.Fatalf("unexpected retry records: %+v", payload.Records) + } +} + +func TestUsageManagementResponseContainsDisplayFields(t *testing.T) { + app := NewApp() + if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil { + t.Fatal(err) + } + record := UsageRecord{ + Provider: "codex", + APIKey: "test-key", + Model: "gpt-5.5", + ReasoningEffort: "high", + ServiceTier: "priority", + ExecutorType: "CodexExecutor", + Endpoint: "POST /v1/responses", + ClientIP: "192.0.2.10", + RequestedAt: time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC), + Latency: 1500 * time.Millisecond, + TTFT: 250 * time.Millisecond, + Detail: UsageDetail{ + InputTokens: 10, + OutputTokens: 5, + TotalTokens: 15, + CachedTokens: 4, + CacheReadTokens: 4, + CacheCreationTokens: 1, + ReasoningTokens: 2, + }, + } + raw, _ := json.Marshal(record) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + + response := managementCall(t, app, http.MethodGet, managementBase+routeUsage) + if response.StatusCode != http.StatusOK { + t.Fatalf("status = %d, body = %s", response.StatusCode, response.Body) + } + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if len(payload.Records) != 1 { + t.Fatalf("records = %d, want 1", len(payload.Records)) + } + got := payload.Records[0] + if got.APIKey != "test-key" || got.Model != "gpt-5.5" || got.ReasoningEffort != "high" || got.ServiceTier != "priority" || got.ExecutorType != "CodexExecutor" { + t.Fatalf("unexpected usage identity fields: %+v", got) + } + if got.RequestType != "SSE" || got.Endpoint != "POST /v1/responses" || got.ClientIP != "192.0.2.10" { + t.Fatalf("unexpected request metadata: %+v", got) + } + if got.TotalTokens != 15 || got.TTFTMilliseconds != 250 || got.CacheReadTokens != 4 || got.CacheWriteTokens != 1 { + t.Fatalf("unexpected usage counters: %+v", got) + } + if got.SpeedTPS == nil || *got.SpeedTPS != 4 || got.CacheRate == nil || *got.CacheRate != 40 { + t.Fatalf("unexpected derived display fields: %+v", got) + } + if got.KeyAlias != "" || got.CostUSD != nil || got.CostAvailable { + t.Fatalf("unexpected usage item: %+v", got) + } +} + +func TestUsageResourceServesTablePage(t *testing.T) { + response := managementCall(t, NewApp(), http.MethodGet, resourceBase+resourceUI) + page := string(response.Body) + if response.StatusCode != http.StatusOK || !strings.Contains(page, "最近用量记录") { + t.Fatalf("unexpected UI response: status=%d", response.StatusCode) + } + for _, column := range []string{"Key / 别名", "推理强度", "生成速度", "缓存写入", "总成本", "客户端 IP"} { + if !strings.Contains(page, column) { + t.Fatalf("UI does not contain column %q", column) + } + } + if !strings.Contains(page, `return "compact"`) || !strings.Contains(page, "isCompactEndpoint(record.endpoint)") { + t.Fatal("UI does not contain compact display rules") + } +} diff --git a/internal/plugin/pricing_management.go b/internal/plugin/pricing_management.go new file mode 100644 index 0000000..a024168 --- /dev/null +++ b/internal/plugin/pricing_management.go @@ -0,0 +1,222 @@ +package plugin + +import ( + "context" + "encoding/json" + "errors" + "math/big" + "net/http" + "strconv" + "strings" + + "cpa-ext/internal/pricing" +) + +type priceRatesDTO struct { + InputPer1M string `json:"input_per_1m"` + CacheReadPer1M string `json:"cache_read_per_1m"` + CacheWritePer1M string `json:"cache_write_per_1m"` + OutputPer1M string `json:"output_per_1m"` +} + +type longContextDTO struct { + ThresholdInputTokens int64 `json:"threshold_input_tokens"` + Comparison string `json:"comparison"` + priceRatesDTO +} + +type priceDTO struct { + Model string `json:"model"` + Base priceRatesDTO `json:"base"` + LongContext *longContextDTO `json:"long_context,omitempty"` + FastPricingEnabled bool `json:"fast_pricing_enabled"` + FastMultiplier string `json:"fast_multiplier"` +} + +type priceDeleteRequest struct { + Model string `json:"model"` +} + +func (a *App) listPrices() ManagementResponse { + a.mu.RLock() + defer a.mu.RUnlock() + if a.store == nil { + return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}}) + } + policies, err := a.store.ListPrices(context.Background()) + if err != nil { + return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}}) + } + prices := make([]priceDTO, 0, len(policies)) + for _, policy := range policies { + prices = append(prices, policyToDTO(policy)) + } + return jsonManagementResponse(http.StatusOK, map[string]any{"prices": prices}) +} + +func (a *App) putPrice(body []byte) ManagementResponse { + var dto priceDTO + if err := decodeJSONBody(body, &dto); err != nil { + return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}}) + } + policy, err := dtoToPolicy(dto) + if err != nil { + return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}}) + } + a.mu.Lock() + defer a.mu.Unlock() + if a.store == nil { + return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}}) + } + if err := a.store.UpsertPrice(context.Background(), policy); err != nil { + return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}}) + } + a.prices[policy.Model] = policy + if _, err := a.store.BackfillMissingCosts(context.Background(), policy); err != nil { + return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}}) + } + return jsonManagementResponse(http.StatusOK, policyToDTO(policy)) +} + +func (a *App) deletePrice(body []byte) ManagementResponse { + var request priceDeleteRequest + if err := decodeJSONBody(body, &request); err != nil { + return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": err.Error()}}) + } + request.Model = strings.TrimSpace(request.Model) + if request.Model == "" { + return jsonManagementResponse(http.StatusBadRequest, map[string]any{"error": map[string]string{"code": "invalid_price", "message": "model 不能为空"}}) + } + a.mu.Lock() + defer a.mu.Unlock() + if a.store == nil { + return jsonManagementResponse(http.StatusServiceUnavailable, map[string]any{"error": map[string]string{"code": "database_unavailable", "message": "价格数据库尚未初始化"}}) + } + if err := a.store.DeletePrice(context.Background(), request.Model); err != nil { + return jsonManagementResponse(http.StatusInternalServerError, map[string]any{"error": map[string]string{"code": "database_error", "message": err.Error()}}) + } + delete(a.prices, request.Model) + return jsonManagementResponse(http.StatusOK, map[string]any{"deleted": request.Model}) +} + +func dtoToPolicy(dto priceDTO) (pricing.Policy, error) { + base, err := dtoRates(dto.Base) + if err != nil { + return pricing.Policy{}, err + } + multiplierText := strings.TrimSpace(dto.FastMultiplier) + if multiplierText == "" { + multiplierText = "2.5" + } + multiplier, err := parseRatio(multiplierText) + if err != nil { + return pricing.Policy{}, errors.New("fast_multiplier 必须是正数，最多保留六位小数") + } + policy := pricing.Policy{Model: strings.TrimSpace(dto.Model), Base: base, FastPricingEnabled: dto.FastPricingEnabled, FastMultiplier: multiplier} + if dto.LongContext != nil { + rates, ratesErr := dtoRates(dto.LongContext.priceRatesDTO) + if ratesErr != nil { + return pricing.Policy{}, ratesErr + } + policy.LongContext = &pricing.LongContext{ThresholdInputTokens: dto.LongContext.ThresholdInputTokens, Comparison: dto.LongContext.Comparison, Rates: rates} + } + if err := policy.Validate(); err != nil { + return pricing.Policy{}, err + } + return policy, nil +} + +func dtoRates(dto priceRatesDTO) (pricing.Rates, error) { + values := []*int64{new(int64), new(int64), new(int64), new(int64)} + texts := []string{dto.InputPer1M, dto.CacheReadPer1M, dto.CacheWritePer1M, dto.OutputPer1M} + for index, text := range texts { + value, err := parseDecimalMicros(text) + if err != nil { + return pricing.Rates{}, errors.New("价格必须是非负数字，最多保留六位小数") + } + *values[index] = value + } + return pricing.Rates{InputMicrosPer1M: *values[0], CacheReadMicrosPer1M: *values[1], CacheWriteMicrosPer1M: *values[2], OutputMicrosPer1M: *values[3]}, nil +} + +func policyToDTO(policy pricing.Policy) priceDTO { + dto := priceDTO{Model: policy.Model, Base: ratesToDTO(policy.Base), FastPricingEnabled: policy.FastPricingEnabled, FastMultiplier: formatRatio(policy.FastMultiplier)} + if policy.LongContext != nil { + dto.LongContext = &longContextDTO{ThresholdInputTokens: policy.LongContext.ThresholdInputTokens, Comparison: policy.LongContext.Comparison, priceRatesDTO: ratesToDTO(policy.LongContext.Rates)} + } + return dto +} + +func ratesToDTO(rates pricing.Rates) priceRatesDTO { + return priceRatesDTO{InputPer1M: formatMicros(rates.InputMicrosPer1M), CacheReadPer1M: formatMicros(rates.CacheReadMicrosPer1M), CacheWritePer1M: formatMicros(rates.CacheWriteMicrosPer1M), OutputPer1M: formatMicros(rates.OutputMicrosPer1M)} +} + +func parseDecimalMicros(text string) (int64, error) { + text = strings.TrimSpace(text) + if text == "" || strings.HasPrefix(text, "-") || strings.Count(text, ".") > 1 { + return 0, errors.New("invalid decimal") + } + parts := strings.SplitN(text, ".", 2) + if parts[0] == "" { + parts[0] = "0" + } + fraction := "" + if len(parts) == 2 { + fraction = parts[1] + } + if len(fraction) > 6 { + return 0, errors.New("too many decimal places") + } + for len(fraction) < 6 { + fraction += "0" + } + whole := new(big.Int) + if _, ok := whole.SetString(parts[0]+fraction, 10); !ok || !whole.IsInt64() { + return 0, errors.New("invalid decimal") + } + return whole.Int64(), nil +} + +func parseRatio(text string) (pricing.Ratio, error) { + numerator, err := parseDecimalMicros(text) + if err != nil || numerator <= 0 { + return pricing.Ratio{}, errors.New("invalid ratio") + } + denominator := int64(1_000_000) + divisor := gcd(numerator, denominator) + return pricing.Ratio{Numerator: numerator / divisor, Denominator: denominator / divisor}, nil +} + +func gcd(a, b int64) int64 { + for b != 0 { + a, b = b, a%b + } + return a +} + +func formatMicros(value int64) string { + whole := value / 1_000_000 + fraction := strconv.FormatInt(value%1_000_000+1_000_000, 10)[1:] + fraction = strings.TrimRight(fraction, "0") + if fraction == "" { + return strconv.FormatInt(whole, 10) + } + return strconv.FormatInt(whole, 10) + "." + fraction +} + +func formatRatio(value pricing.Ratio) string { + if value.Denominator == 0 { + return "" + } + rational := new(big.Rat).SetFrac(big.NewInt(value.Numerator), big.NewInt(value.Denominator)) + return strings.TrimRight(strings.TrimRight(rational.FloatString(6), "0"), ".") +} + +func decodeJSONBody(body []byte, target any) error { + decoder := json.NewDecoder(strings.NewReader(string(body))) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return err + } + return nil +} diff --git a/internal/plugin/types.go b/internal/plugin/types.go index 4b9552d..c3d3f93 100644 --- a/internal/plugin/types.go +++ b/internal/plugin/types.go @@ -3,6 +3,7 @@ package plugin import ( "encoding/json" "net/http" + "net/url" "time" ) @@ -14,10 +15,13 @@ const ( ) const ( - MethodPluginRegister = "plugin.register" - MethodPluginReconfigure = "plugin.reconfigure" - MethodPluginShutdown = "plugin.shutdown" - MethodUsageHandle = "usage.handle" + MethodPluginRegister = "plugin.register" + MethodPluginReconfigure = "plugin.reconfigure" + MethodPluginShutdown = "plugin.shutdown" + MethodRequestComplete = "request.complete" + MethodUsageHandle = "usage.handle" + MethodManagementRegister = "management.register" + MethodManagementHandle = "management.handle" ) type Envelope struct { @@ -61,14 +65,77 @@ type ConfigField struct { } type Capabilities struct { - UsagePlugin bool `json:"usage_plugin"` + RequestLifecyclePlugin bool `json:"request_lifecycle_plugin"` + UsagePlugin bool `json:"usage_plugin"` + ManagementAPI bool `json:"management_api"` +} + +type RequestCompletionOutcome string + +const ( + RequestCompletionSucceeded RequestCompletionOutcome = "succeeded" + RequestCompletionFailed RequestCompletionOutcome = "failed" + RequestCompletionRejected RequestCompletionOutcome = "rejected" + RequestCompletionCanceled RequestCompletionOutcome = "canceled" +) + +// RequestCompletion 是一次模型执行的终态，Token 由 UsageRecord 单独补充。 +type RequestCompletion struct { + RequestID string + TraceID string + SourceFormat string + Model string + RequestedModel string + Stream bool + Outcome RequestCompletionOutcome + StatusCode int + Error string + StartedAt time.Time + CompletedAt time.Time + Metadata map[string]any +} + +type ManagementRegistrationResponse struct { + Routes []ManagementRoute `json:"routes,omitempty"` + Resources []ResourceRoute `json:"resources,omitempty"` +} + +type ManagementRoute struct { + Method string `json:"Method"` + Path string `json:"Path"` + Description string `json:"Description,omitempty"` +} + +type ResourceRoute struct { + Path string `json:"Path"` + Menu string `json:"Menu,omitempty"` + Description string `json:"Description,omitempty"` +} + +type ManagementRequest struct { + Method string `json:"Method"` + Path string `json:"Path"` + Headers http.Header `json:"Headers"` + Query url.Values `json:"Query"` + Body []byte `json:"Body"` +} + +type ManagementResponse struct { + StatusCode int `json:"StatusCode,omitempty"` + Headers http.Header `json:"Headers,omitempty"` + Body []byte `json:"Body,omitempty"` } // UsageRecord mirrors CLIProxyAPI sdk/pluginapi. Keep it in sync with the // target host because exported field names are part of the JSON wire contract. type UsageRecord struct { Provider string + RequestID string + ExecutionID string + TraceID string ExecutorType string + Endpoint string + ClientIP string Model string Alias string APIKey string @@ -78,6 +145,7 @@ type UsageRecord struct { Source string ReasoningEffort string ServiceTier string + Speed string Generate bool RequestedAt time.Time Latency time.Duration diff --git a/internal/pricing/pricing.go b/internal/pricing/pricing.go new file mode 100644 index 0000000..0a4ebab --- /dev/null +++ b/internal/pricing/pricing.go @@ -0,0 +1,163 @@ +// Package pricing implements deterministic model-price resolution and cost calculation. +package pricing + +import ( + "errors" + "fmt" + "math/big" + "strings" +) + +const perMillion = int64(1_000_000) + +// Rates stores micro-USD rates per one million tokens. +type Rates struct { + InputMicrosPer1M int64 + CacheReadMicrosPer1M int64 + CacheWriteMicrosPer1M int64 + OutputMicrosPer1M int64 +} + +// Ratio stores a multiplier exactly. Fast pricing is represented as 5/2. +type Ratio struct { + Numerator int64 + Denominator int64 +} + +// LongContext replaces all four rates after the input threshold is reached. +type LongContext struct { + ThresholdInputTokens int64 + Comparison string + Rates Rates +} + +// Policy is the complete pricing configuration for one exact model name. +type Policy struct { + Model string + Base Rates + LongContext *LongContext + FastPricingEnabled bool + FastMultiplier Ratio +} + +// Usage contains token buckets required by pricing. InputTokens includes cache buckets. +type Usage struct { + InputTokens int64 + CacheReadTokens int64 + CacheWriteTokens int64 + OutputTokens int64 + ServiceTier string + Speed string +} + +// Result records the applied policy facts together with the final amount. +type Result struct { + CostMicros int64 + PriceTier string + FastRequested bool + FastApplied bool + MultiplierNumerator int64 + MultiplierDenominator int64 +} + +// Validate rejects incomplete or ambiguous policies before they enter the runtime catalog. +func (p Policy) Validate() error { + if strings.TrimSpace(p.Model) == "" { + return errors.New("model is required") + } + if err := p.Base.validate(); err != nil { + return fmt.Errorf("base rates: %w", err) + } + if p.FastMultiplier.Numerator <= 0 || p.FastMultiplier.Denominator <= 0 { + return errors.New("fast multiplier must be positive") + } + if p.LongContext != nil { + if p.LongContext.ThresholdInputTokens < 0 { + return errors.New("long-context threshold cannot be negative") + } + if p.LongContext.Comparison != "gt" && p.LongContext.Comparison != "gte" { + return errors.New("long-context comparison must be gt or gte") + } + if err := p.LongContext.Rates.validate(); err != nil { + return fmt.Errorf("long-context rates: %w", err) + } + } + return nil +} + +func (r Rates) validate() error { + values := [...]int64{r.InputMicrosPer1M, r.CacheReadMicrosPer1M, r.CacheWriteMicrosPer1M, r.OutputMicrosPer1M} + for _, value := range values { + if value < 0 { + return errors.New("rates cannot be negative") + } + } + return nil +} + +// Calculate selects the long-context tier first, then applies Fast once to the total. +func Calculate(policy Policy, usage Usage) (Result, error) { + if err := policy.Validate(); err != nil { + return Result{}, err + } + if usage.InputTokens < 0 || usage.CacheReadTokens < 0 || usage.CacheWriteTokens < 0 || usage.OutputTokens < 0 { + return Result{}, errors.New("token counts cannot be negative") + } + if usage.CacheReadTokens > usage.InputTokens-usage.CacheWriteTokens { + return Result{}, errors.New("cache token counts exceed input tokens") + } + + rates := policy.Base + tier := "base" + if long := policy.LongContext; long != nil && thresholdMatched(usage.InputTokens, long.ThresholdInputTokens, long.Comparison) { + rates = long.Rates + tier = "long_context" + } + + multiplier := Ratio{Numerator: 1, Denominator: 1} + fastRequested := isFast(usage.ServiceTier, usage.Speed) + fastApplied := policy.FastPricingEnabled && fastRequested + if fastApplied { + multiplier = policy.FastMultiplier + } + + uncachedInput := usage.InputTokens - usage.CacheReadTokens - usage.CacheWriteTokens + numerator := new(big.Int) + addCostTerm(numerator, uncachedInput, rates.InputMicrosPer1M) + addCostTerm(numerator, usage.CacheReadTokens, rates.CacheReadMicrosPer1M) + addCostTerm(numerator, usage.CacheWriteTokens, rates.CacheWriteMicrosPer1M) + addCostTerm(numerator, usage.OutputTokens, rates.OutputMicrosPer1M) + numerator.Mul(numerator, big.NewInt(multiplier.Numerator)) + denominator := big.NewInt(perMillion * multiplier.Denominator) + // The complete request is rounded once using round-half-up. + numerator.Add(numerator, new(big.Int).Quo(new(big.Int).Set(denominator), big.NewInt(2))) + numerator.Quo(numerator, denominator) + if !numerator.IsInt64() { + return Result{}, errors.New("calculated cost exceeds int64") + } + return Result{ + CostMicros: numerator.Int64(), + PriceTier: tier, + FastRequested: fastRequested, + FastApplied: fastApplied, + MultiplierNumerator: multiplier.Numerator, + MultiplierDenominator: multiplier.Denominator, + }, nil +} + +func isFast(serviceTier, speed string) bool { + tier := strings.ToLower(strings.TrimSpace(serviceTier)) + return tier == "priority" || tier == "fast" || strings.EqualFold(strings.TrimSpace(speed), "fast") +} + +func thresholdMatched(input, threshold int64, comparison string) bool { + if comparison == "gte" { + return input >= threshold + } + return input > threshold +} + +func addCostTerm(total *big.Int, tokens, rate int64) { + term := new(big.Int).Mul(big.NewInt(tokens), big.NewInt(rate)) + total.Add(total, term) +} diff --git a/internal/pricing/pricing_test.go b/internal/pricing/pricing_test.go new file mode 100644 index 0000000..eb96242 --- /dev/null +++ b/internal/pricing/pricing_test.go @@ -0,0 +1,91 @@ +package pricing_test + +import ( + "testing" + + "cpa-ext/internal/pricing" +) + +func testPolicy() pricing.Policy { + return pricing.Policy{ + Model: "gpt-5.6-sol", + Base: pricing.Rates{ + InputMicrosPer1M: 2_500_000, + CacheReadMicrosPer1M: 250_000, + CacheWriteMicrosPer1M: 3_125_000, + OutputMicrosPer1M: 15_000_000, + }, + LongContext: &pricing.LongContext{ + ThresholdInputTokens: 272_000, + Comparison: "gt", + Rates: pricing.Rates{ + InputMicrosPer1M: 5_000_000, + CacheReadMicrosPer1M: 500_000, + CacheWriteMicrosPer1M: 6_250_000, + OutputMicrosPer1M: 22_500_000, + }, + }, + FastPricingEnabled: true, + FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2}, + } +} + +func TestCalculateUsesEachTokenBucketOnce(t *testing.T) { + result, err := pricing.Calculate(testPolicy(), pricing.Usage{ + InputTokens: 1_000_000, CacheReadTokens: 400_000, CacheWriteTokens: 100_000, OutputTokens: 500_000, + }) + if err != nil { + t.Fatal(err) + } + // The request is above the threshold, so all buckets use long-context rates. + if result.CostMicros != 14_575_000 || result.PriceTier != "long_context" || result.FastApplied { + t.Fatalf("unexpected result: %+v", result) + } +} + +func TestCalculateCanDisableFastPricingWithoutChangingRequestMode(t *testing.T) { + policy := testPolicy() + policy.FastPricingEnabled = false + result, err := pricing.Calculate(policy, pricing.Usage{InputTokens: 100_000, OutputTokens: 10_000, ServiceTier: "priority"}) + if err != nil { + t.Fatal(err) + } + if result.CostMicros != 400_000 || !result.FastRequested || result.FastApplied || result.MultiplierNumerator != 1 || result.MultiplierDenominator != 1 { + t.Fatalf("unexpected disabled Fast pricing result: %+v", result) + } +} + +func TestCalculateRecognizesPassedFastSpeed(t *testing.T) { + result, err := pricing.Calculate(testPolicy(), pricing.Usage{InputTokens: 100_000, Speed: "fast"}) + if err != nil { + t.Fatal(err) + } + if !result.FastRequested || !result.FastApplied || result.CostMicros != 625_000 { + t.Fatalf("unexpected speed=fast result: %+v", result) + } +} + +func TestCalculateUsesStrictThresholdAndAppliesFastOnce(t *testing.T) { + policy := testPolicy() + base, err := pricing.Calculate(policy, pricing.Usage{InputTokens: 272_000, OutputTokens: 10_000, ServiceTier: "priority"}) + if err != nil { + t.Fatal(err) + } + if base.PriceTier != "base" || base.CostMicros != 2_075_000 || !base.FastApplied || base.MultiplierNumerator != 5 || base.MultiplierDenominator != 2 { + t.Fatalf("unexpected threshold result: %+v", base) + } + long, err := pricing.Calculate(policy, pricing.Usage{InputTokens: 272_001, OutputTokens: 10_000, ServiceTier: "priority"}) + if err != nil { + t.Fatal(err) + } + if long.PriceTier != "long_context" || long.CostMicros != 3_962_513 { + t.Fatalf("unexpected long-context result: %+v", long) + } +} + +func TestCalculateRejectsOverlappingCacheCounts(t *testing.T) { + _, err := pricing.Calculate(testPolicy(), pricing.Usage{InputTokens: 10, CacheReadTokens: 8, CacheWriteTokens: 3}) + if err == nil { + t.Fatal("overlapping cache counts unexpectedly accepted") + } +} diff --git a/internal/repository/sqlite_pricing.go b/internal/repository/sqlite_pricing.go new file mode 100644 index 0000000..4588da9 --- /dev/null +++ b/internal/repository/sqlite_pricing.go @@ -0,0 +1,178 @@ +package repository + +import ( + "context" + "fmt" + "time" + + "cpa-ext/internal/pricing" +) + +// ListPrices returns every configured exact-model policy. +func (r *SQLiteUsageRepository) ListPrices(ctx context.Context) ([]pricing.Policy, error) { + rows, err := r.db.QueryContext(ctx, ` +SELECT model, input_rate_micros, cache_read_rate_micros, cache_write_rate_micros, output_rate_micros, + long_context_enabled, long_context_threshold, long_context_comparison, + long_input_rate_micros, long_cache_read_rate_micros, long_cache_write_rate_micros, long_output_rate_micros, + fast_pricing_enabled, fast_multiplier_numerator, fast_multiplier_denominator +FROM model_prices ORDER BY model`) + if err != nil { + return nil, fmt.Errorf("查询模型价格: %w", err) + } + defer rows.Close() + + var policies []pricing.Policy + for rows.Next() { + var policy pricing.Policy + var longEnabled bool + var long pricing.LongContext + if err := rows.Scan( + &policy.Model, + &policy.Base.InputMicrosPer1M, &policy.Base.CacheReadMicrosPer1M, + &policy.Base.CacheWriteMicrosPer1M, &policy.Base.OutputMicrosPer1M, + &longEnabled, &long.ThresholdInputTokens, &long.Comparison, + &long.Rates.InputMicrosPer1M, &long.Rates.CacheReadMicrosPer1M, + &long.Rates.CacheWriteMicrosPer1M, &long.Rates.OutputMicrosPer1M, + &policy.FastPricingEnabled, &policy.FastMultiplier.Numerator, &policy.FastMultiplier.Denominator, + ); err != nil { + return nil, fmt.Errorf("读取模型价格: %w", err) + } + if longEnabled { + policy.LongContext = &long + } + policies = append(policies, policy) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("遍历模型价格: %w", err) + } + return policies, nil +} + +// UpsertPrice atomically replaces one model policy. +func (r *SQLiteUsageRepository) UpsertPrice(ctx context.Context, policy pricing.Policy) error { + if err := policy.Validate(); err != nil { + return err + } + longEnabled := policy.LongContext != nil + long := pricing.LongContext{Comparison: "gt"} + if policy.LongContext != nil { + long = *policy.LongContext + } + _, err := r.db.ExecContext(ctx, ` +INSERT INTO model_prices ( + model, input_rate_micros, cache_read_rate_micros, cache_write_rate_micros, output_rate_micros, + long_context_enabled, long_context_threshold, long_context_comparison, + long_input_rate_micros, long_cache_read_rate_micros, long_cache_write_rate_micros, long_output_rate_micros, + fast_pricing_enabled, fast_multiplier_numerator, fast_multiplier_denominator, updated_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(model) DO UPDATE SET + input_rate_micros=excluded.input_rate_micros, + cache_read_rate_micros=excluded.cache_read_rate_micros, + cache_write_rate_micros=excluded.cache_write_rate_micros, + output_rate_micros=excluded.output_rate_micros, + long_context_enabled=excluded.long_context_enabled, + long_context_threshold=excluded.long_context_threshold, + long_context_comparison=excluded.long_context_comparison, + long_input_rate_micros=excluded.long_input_rate_micros, + long_cache_read_rate_micros=excluded.long_cache_read_rate_micros, + long_cache_write_rate_micros=excluded.long_cache_write_rate_micros, + long_output_rate_micros=excluded.long_output_rate_micros, + fast_pricing_enabled=excluded.fast_pricing_enabled, + fast_multiplier_numerator=excluded.fast_multiplier_numerator, + fast_multiplier_denominator=excluded.fast_multiplier_denominator, + updated_at=excluded.updated_at`, + policy.Model, + policy.Base.InputMicrosPer1M, policy.Base.CacheReadMicrosPer1M, + policy.Base.CacheWriteMicrosPer1M, policy.Base.OutputMicrosPer1M, + longEnabled, long.ThresholdInputTokens, long.Comparison, + long.Rates.InputMicrosPer1M, long.Rates.CacheReadMicrosPer1M, + long.Rates.CacheWriteMicrosPer1M, long.Rates.OutputMicrosPer1M, + policy.FastPricingEnabled, policy.FastMultiplier.Numerator, policy.FastMultiplier.Denominator, + time.Now().UTC().Format(time.RFC3339Nano), + ) + if err != nil { + return fmt.Errorf("保存模型价格: %w", err) + } + return nil +} + +// BackfillMissingCosts 只补算尚未定价的历史记录，已经保存的账单金额不会随价格修改而变化。 +func (r *SQLiteUsageRepository) BackfillMissingCosts(ctx context.Context, policy pricing.Policy) (int64, error) { + if err := policy.Validate(); err != nil { + return 0, err + } + rows, err := r.db.QueryContext(ctx, ` +SELECT id, input_tokens, cache_read_tokens, cache_write_tokens, output_tokens, service_tier, speed +FROM usage_records +WHERE model = ? AND cost_micros IS NULL`, policy.Model) + if err != nil { + return 0, fmt.Errorf("查询待补算用量: %w", err) + } + type pendingCost struct { + id int64 + inputTokens int64 + cacheReadTokens int64 + cacheWriteTokens int64 + outputTokens int64 + serviceTier, speed string + } + var pending []pendingCost + for rows.Next() { + var item pendingCost + if err := rows.Scan(&item.id, &item.inputTokens, &item.cacheReadTokens, &item.cacheWriteTokens, &item.outputTokens, &item.serviceTier, &item.speed); err != nil { + _ = rows.Close() + return 0, fmt.Errorf("读取待补算用量: %w", err) + } + pending = append(pending, item) + } + if err := rows.Err(); err != nil { + _ = rows.Close() + return 0, fmt.Errorf("遍历待补算用量: %w", err) + } + if err := rows.Close(); err != nil { + return 0, fmt.Errorf("关闭待补算查询: %w", err) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return 0, fmt.Errorf("开始补算事务: %w", err) + } + defer func() { _ = tx.Rollback() }() + var updated int64 + for _, item := range pending { + result, calculateErr := pricing.Calculate(policy, pricing.Usage{ + InputTokens: item.inputTokens, CacheReadTokens: item.cacheReadTokens, + CacheWriteTokens: item.cacheWriteTokens, OutputTokens: item.outputTokens, + ServiceTier: item.serviceTier, Speed: item.speed, + }) + if calculateErr != nil { + continue + } + change, updateErr := tx.ExecContext(ctx, ` +UPDATE usage_records +SET cost_micros = ?, price_tier = ?, fast_requested = ?, fast_pricing_applied = ?, + price_multiplier_numerator = ?, price_multiplier_denominator = ? +WHERE id = ? AND cost_micros IS NULL`, + result.CostMicros, result.PriceTier, result.FastRequested, result.FastApplied, + result.MultiplierNumerator, result.MultiplierDenominator, item.id) + if updateErr != nil { + return 0, fmt.Errorf("补算用量价格: %w", updateErr) + } + count, countErr := change.RowsAffected() + if countErr != nil { + return 0, fmt.Errorf("读取补算数量: %w", countErr) + } + updated += count + } + if err := tx.Commit(); err != nil { + return 0, fmt.Errorf("提交补算事务: %w", err) + } + return updated, nil +} + +func (r *SQLiteUsageRepository) DeletePrice(ctx context.Context, model string) error { + if _, err := r.db.ExecContext(ctx, `DELETE FROM model_prices WHERE model = ?`, model); err != nil { + return fmt.Errorf("删除模型价格: %w", err) + } + return nil +} diff --git a/internal/repository/sqlite_pricing_test.go b/internal/repository/sqlite_pricing_test.go new file mode 100644 index 0000000..95f50a3 --- /dev/null +++ b/internal/repository/sqlite_pricing_test.go @@ -0,0 +1,88 @@ +package repository_test + +import ( + "context" + "path/filepath" + "testing" + "time" + + "cpa-ext/internal/collection" + "cpa-ext/internal/pricing" + "cpa-ext/internal/repository" +) + +func TestSQLitePricingRoundTripAndDelete(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + policy := pricing.Policy{ + Model: "gpt-5.6-sol", + Base: pricing.Rates{InputMicrosPer1M: 2_500_000, CacheReadMicrosPer1M: 250_000, CacheWriteMicrosPer1M: 3_125_000, OutputMicrosPer1M: 15_000_000}, + LongContext: &pricing.LongContext{ThresholdInputTokens: 272_000, Comparison: "gt", Rates: pricing.Rates{InputMicrosPer1M: 5_000_000, OutputMicrosPer1M: 22_500_000}}, + FastPricingEnabled: true, + FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2}, + } + if err := store.UpsertPrice(context.Background(), policy); err != nil { + t.Fatal(err) + } + prices, err := store.ListPrices(context.Background()) + if err != nil { + t.Fatal(err) + } + if len(prices) != 1 || prices[0].Model != policy.Model || prices[0].LongContext == nil || !prices[0].FastPricingEnabled || prices[0].FastMultiplier != policy.FastMultiplier { + t.Fatalf("unexpected prices: %+v", prices) + } + if err := store.DeletePrice(context.Background(), policy.Model); err != nil { + t.Fatal(err) + } + prices, err = store.ListPrices(context.Background()) + if err != nil || len(prices) != 0 { + t.Fatalf("prices after delete: %+v, %v", prices, err) + } +} + +func TestBackfillMissingCostsDoesNotRewriteExistingCost(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + missing := collection.Record{RequestedAt: time.Now(), Model: "deepseek-v4-flash", InputTokens: 1_000_000, OutputTokens: 100_000} + existingCost := int64(7) + priced := collection.Record{RequestedAt: time.Now(), Model: "deepseek-v4-flash", InputTokens: 1_000_000, CostMicros: &existingCost} + if err := store.Insert(context.Background(), missing); err != nil { + t.Fatal(err) + } + if err := store.Insert(context.Background(), priced); err != nil { + t.Fatal(err) + } + policy := pricing.Policy{ + Model: "deepseek-v4-flash", Base: pricing.Rates{InputMicrosPer1M: 2_500_000, OutputMicrosPer1M: 15_000_000}, + FastMultiplier: pricing.Ratio{Numerator: 5, Denominator: 2}, + } + updated, err := store.BackfillMissingCosts(context.Background(), policy) + if err != nil { + t.Fatal(err) + } + if updated != 1 { + t.Fatalf("updated = %d, want 1", updated) + } + records, err := store.ListRecent(context.Background(), 10) + if err != nil { + t.Fatal(err) + } + if len(records) != 2 || records[0].CostMicros == nil || records[1].CostMicros == nil { + t.Fatalf("unexpected records: %+v", records) + } + got := map[int64]int{} + for _, record := range records { + got[*record.CostMicros]++ + } + if got[4_000_000] != 1 || got[7] != 1 { + t.Fatalf("costs = %#v", got) + } +} diff --git a/internal/repository/sqlite_usage.go b/internal/repository/sqlite_usage.go new file mode 100644 index 0000000..3f0d03c --- /dev/null +++ b/internal/repository/sqlite_usage.go @@ -0,0 +1,309 @@ +// Package repository 实现 cpa-ext 的本地持久化。 +package repository + +import ( + "context" + "database/sql" + "errors" + "fmt" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "cpa-ext/internal/collection" + _ "github.com/mattn/go-sqlite3" +) + +const usageSchema = ` +CREATE TABLE IF NOT EXISTS usage_records ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + request_id TEXT NOT NULL DEFAULT '', + execution_id TEXT NOT NULL DEFAULT '', + trace_id TEXT NOT NULL DEFAULT '', + requested_at TEXT NOT NULL, + api_key TEXT NOT NULL DEFAULT '', + key_alias TEXT NOT NULL DEFAULT '', + model TEXT NOT NULL DEFAULT '', + reasoning_effort TEXT NOT NULL DEFAULT '', + service_tier TEXT NOT NULL DEFAULT '', + speed TEXT NOT NULL DEFAULT '', + failed INTEGER NOT NULL, + executor_type TEXT NOT NULL DEFAULT '', + request_type TEXT NOT NULL DEFAULT '', + endpoint TEXT NOT NULL DEFAULT '', + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + reasoning_tokens INTEGER NOT NULL DEFAULT 0, + cached_tokens INTEGER NOT NULL DEFAULT 0, + cache_read_tokens INTEGER NOT NULL DEFAULT 0, + cache_write_tokens INTEGER NOT NULL DEFAULT 0, + total_tokens INTEGER NOT NULL DEFAULT 0, + ttft_ns INTEGER NOT NULL DEFAULT 0, + latency_ns INTEGER NOT NULL DEFAULT 0, + cost_micros INTEGER, + price_tier TEXT NOT NULL DEFAULT '', + fast_requested INTEGER NOT NULL DEFAULT 0, + fast_pricing_applied INTEGER NOT NULL DEFAULT 0, + price_multiplier_numerator INTEGER NOT NULL DEFAULT 1, + price_multiplier_denominator INTEGER NOT NULL DEFAULT 1, + client_ip TEXT NOT NULL DEFAULT '' +); +CREATE INDEX IF NOT EXISTS idx_usage_records_requested_at ON usage_records(requested_at DESC, id DESC); +CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_records_execution_id ON usage_records(execution_id) WHERE execution_id <> ''; + +CREATE TABLE IF NOT EXISTS request_records ( + request_id TEXT PRIMARY KEY, + trace_id TEXT NOT NULL DEFAULT '', + requested_at TEXT NOT NULL, + completed_at TEXT NOT NULL, + model TEXT NOT NULL DEFAULT '', + requested_model TEXT NOT NULL DEFAULT '', + source_format TEXT NOT NULL DEFAULT '', + stream INTEGER NOT NULL DEFAULT 0, + outcome TEXT NOT NULL, + status_code INTEGER NOT NULL DEFAULT 0, + error TEXT NOT NULL DEFAULT '', + endpoint TEXT NOT NULL DEFAULT '', + latency_ns INTEGER NOT NULL DEFAULT 0 +); +CREATE INDEX IF NOT EXISTS idx_request_records_requested_at ON request_records(requested_at DESC); + +CREATE TABLE IF NOT EXISTS model_prices ( + model TEXT PRIMARY KEY, + input_rate_micros INTEGER NOT NULL, + cache_read_rate_micros INTEGER NOT NULL, + cache_write_rate_micros INTEGER NOT NULL, + output_rate_micros INTEGER NOT NULL, + long_context_enabled INTEGER NOT NULL DEFAULT 0, + long_context_threshold INTEGER NOT NULL DEFAULT 0, + long_context_comparison TEXT NOT NULL DEFAULT 'gt', + long_input_rate_micros INTEGER NOT NULL DEFAULT 0, + long_cache_read_rate_micros INTEGER NOT NULL DEFAULT 0, + long_cache_write_rate_micros INTEGER NOT NULL DEFAULT 0, + long_output_rate_micros INTEGER NOT NULL DEFAULT 0, + fast_pricing_enabled INTEGER NOT NULL DEFAULT 0, + fast_multiplier_numerator INTEGER NOT NULL DEFAULT 5, + fast_multiplier_denominator INTEGER NOT NULL DEFAULT 2, + updated_at TEXT NOT NULL +); +` + +// SQLiteUsageRepository 使用单写连接保存全部用量记录。 +type SQLiteUsageRepository struct { + db *sql.DB +} + +func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) { + dsn, err := sqliteDSN(databasePath) + if err != nil { + return nil, err + } + db, err := sql.Open("sqlite3", dsn) + if err != nil { + return nil, fmt.Errorf("打开 SQLite: %w", err) + } + // 当前插件只需要一个同步 writer，避免并发写争用。 + db.SetMaxOpenConns(1) + if err := db.Ping(); err != nil { + _ = db.Close() + return nil, fmt.Errorf("连接 SQLite: %w", err) + } + if _, err := db.Exec(usageSchema); err != nil { + _ = db.Close() + return nil, fmt.Errorf("初始化 SQLite schema: %w", err) + } + return &SQLiteUsageRepository{db: db}, nil +} + +func sqliteDSN(databasePath string) (string, error) { + databasePath = strings.TrimSpace(databasePath) + if databasePath == "" { + return "", errors.New("database_path 不能为空") + } + if databasePath == ":memory:" { + return "file:cpa-ext-memory?mode=memory&cache=shared&_busy_timeout=5000&_foreign_keys=on", nil + } + if strings.HasPrefix(databasePath, "file:") { + separator := "?" + if strings.Contains(databasePath, "?") { + separator = "&" + } + return databasePath + separator + "_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil + } + absolutePath, err := filepath.Abs(databasePath) + if err != nil { + return "", fmt.Errorf("解析 database_path: %w", err) + } + if err := os.MkdirAll(filepath.Dir(absolutePath), 0o755); err != nil { + return "", fmt.Errorf("创建数据库目录: %w", err) + } + fileURL := (&url.URL{Scheme: "file", Path: filepath.ToSlash(absolutePath)}).String() + return fileURL + "?_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil +} + +func (r *SQLiteUsageRepository) Insert(ctx context.Context, record collection.Record) error { + _, err := r.db.ExecContext(ctx, ` +INSERT OR IGNORE INTO usage_records ( + request_id, execution_id, trace_id, requested_at, api_key, key_alias, model, reasoning_effort, service_tier, speed, + failed, executor_type, request_type, endpoint, input_tokens, output_tokens, reasoning_tokens, + cached_tokens, cache_read_tokens, cache_write_tokens, total_tokens, + ttft_ns, latency_ns, cost_micros, price_tier, fast_requested, fast_pricing_applied, + price_multiplier_numerator, price_multiplier_denominator, client_ip +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + record.RequestID, record.ExecutionID, record.TraceID, + record.RequestedAt.UTC().Format(time.RFC3339Nano), record.APIKey, record.KeyAlias, + record.Model, record.ReasoningEffort, record.ServiceTier, record.Speed, record.Failed, + record.ExecutorType, record.RequestType, record.Endpoint, record.InputTokens, record.OutputTokens, + record.ReasoningTokens, record.CachedTokens, record.CacheReadTokens, + record.CacheWriteTokens, record.TotalTokens, int64(record.TTFT), int64(record.Latency), + record.CostMicros, record.PriceTier, record.FastRequested, record.FastPricingApplied, + record.PriceMultiplierNumerator, record.PriceMultiplierDenominator, record.ClientIP, + ) + if err != nil { + return fmt.Errorf("写入用量记录: %w", err) + } + return nil +} + +func (r *SQLiteUsageRepository) UpsertRequest(ctx context.Context, record collection.RequestRecord) error { + requestID := strings.TrimSpace(record.RequestID) + if requestID == "" { + return errors.New("request_id 不能为空") + } + requestedAt := record.RequestedAt + completedAt := record.CompletedAt + if requestedAt.IsZero() { + requestedAt = completedAt + } + if completedAt.IsZero() { + completedAt = requestedAt + } + latency := completedAt.Sub(requestedAt) + if latency < 0 { + latency = 0 + } + _, err := r.db.ExecContext(ctx, ` +INSERT INTO request_records ( + request_id, trace_id, requested_at, completed_at, model, requested_model, + source_format, stream, outcome, status_code, error, endpoint, latency_ns +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON CONFLICT(request_id) DO UPDATE SET + trace_id=excluded.trace_id, + requested_at=excluded.requested_at, + completed_at=excluded.completed_at, + model=excluded.model, + requested_model=excluded.requested_model, + source_format=excluded.source_format, + stream=excluded.stream, + outcome=excluded.outcome, + status_code=excluded.status_code, + error=excluded.error, + endpoint=excluded.endpoint, + latency_ns=excluded.latency_ns`, + requestID, record.TraceID, requestedAt.UTC().Format(time.RFC3339Nano), + completedAt.UTC().Format(time.RFC3339Nano), record.Model, record.RequestedModel, + record.SourceFormat, record.Stream, record.Outcome, record.StatusCode, + record.Error, record.Endpoint, int64(latency), + ) + if err != nil { + return fmt.Errorf("写入请求终态: %w", err) + } + return nil +} + +func (r *SQLiteUsageRepository) ListRecent(ctx context.Context, limit int) ([]collection.Record, error) { + if limit < 1 { + return []collection.Record{}, nil + } + rows, err := r.db.QueryContext(ctx, ` +WITH combined AS ( + SELECT + u.request_id, u.execution_id, COALESCE(NULLIF(u.trace_id, ''), r.trace_id, '') AS trace_id, + u.requested_at, u.api_key, u.key_alias, u.model, u.reasoning_effort, u.service_tier, u.speed, + CASE + WHEN u.failed = 1 THEN 1 + WHEN r.outcome IN ('failed', 'rejected', 'canceled') THEN 1 + ELSE 0 + END AS failed, + u.executor_type, + CASE WHEN r.request_id IS NOT NULL THEN CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END ELSE u.request_type END AS request_type, + COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '') AS endpoint, + u.input_tokens, u.output_tokens, u.reasoning_tokens, u.cached_tokens, + u.cache_read_tokens, u.cache_write_tokens, u.total_tokens, + u.ttft_ns, u.latency_ns, u.cost_micros, u.price_tier, u.fast_requested, + u.fast_pricing_applied, u.price_multiplier_numerator, u.price_multiplier_denominator, + u.client_ip, COALESCE(r.outcome, '') AS outcome, COALESCE(r.status_code, 0) AS status_code, + COALESCE(r.error, '') AS error + FROM usage_records u + LEFT JOIN request_records r ON r.request_id = u.request_id + + UNION ALL + + SELECT + r.request_id, '', r.trace_id, r.requested_at, '', '', + COALESCE(NULLIF(r.model, ''), r.requested_model), '', '', '', + CASE WHEN r.outcome = 'succeeded' THEN 0 ELSE 1 END, + '', CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END, r.endpoint, + 0, 0, 0, 0, 0, 0, 0, 0, r.latency_ns, NULL, '', 0, 0, 1, 1, '', + r.outcome, r.status_code, r.error + FROM request_records r + WHERE NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id = r.request_id) +) +SELECT request_id, execution_id, trace_id, requested_at, api_key, key_alias, model, + reasoning_effort, service_tier, speed, failed, executor_type, request_type, endpoint, + input_tokens, output_tokens, reasoning_tokens, cached_tokens, cache_read_tokens, + cache_write_tokens, total_tokens, ttft_ns, latency_ns, cost_micros, price_tier, + fast_requested, fast_pricing_applied, price_multiplier_numerator, + price_multiplier_denominator, client_ip, outcome, status_code, error +FROM combined +ORDER BY requested_at DESC +LIMIT ?`, limit) + if err != nil { + return nil, fmt.Errorf("查询最近用量记录: %w", err) + } + defer rows.Close() + + records := make([]collection.Record, 0, limit) + for rows.Next() { + var record collection.Record + var requestedAt string + var failed bool + var ttftNS, latencyNS int64 + var cost sql.NullInt64 + if err := rows.Scan( + &record.RequestID, &record.ExecutionID, &record.TraceID, + &requestedAt, &record.APIKey, &record.KeyAlias, &record.Model, + &record.ReasoningEffort, &record.ServiceTier, &record.Speed, &failed, &record.ExecutorType, + &record.RequestType, &record.Endpoint, &record.InputTokens, &record.OutputTokens, + &record.ReasoningTokens, &record.CachedTokens, &record.CacheReadTokens, + &record.CacheWriteTokens, &record.TotalTokens, &ttftNS, &latencyNS, + &cost, &record.PriceTier, &record.FastRequested, &record.FastPricingApplied, + &record.PriceMultiplierNumerator, &record.PriceMultiplierDenominator, &record.ClientIP, + &record.Outcome, &record.StatusCode, &record.Error, + ); err != nil { + return nil, fmt.Errorf("读取用量记录: %w", err) + } + record.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt) + if err != nil { + return nil, fmt.Errorf("解析用量时间: %w", err) + } + record.Failed = failed + record.TTFT = time.Duration(ttftNS) + record.Latency = time.Duration(latencyNS) + if cost.Valid { + value := cost.Int64 + record.CostMicros = &value + } + records = append(records, record) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("遍历用量记录: %w", err) + } + return records, nil +} + +func (r *SQLiteUsageRepository) Close() error { + return r.db.Close() +} diff --git a/internal/repository/sqlite_usage_test.go b/internal/repository/sqlite_usage_test.go new file mode 100644 index 0000000..19b7e21 --- /dev/null +++ b/internal/repository/sqlite_usage_test.go @@ -0,0 +1,94 @@ +package repository_test + +import ( + "context" + "fmt" + "path/filepath" + "testing" + "time" + + "cpa-ext/internal/collection" + "cpa-ext/internal/repository" +) + +func TestSQLiteUsagePersistsAllRecordsAndLimitsQueries(t *testing.T) { + databasePath := filepath.Join(t.TempDir(), "usage.db") + store, err := repository.OpenSQLiteUsage(databasePath) + if err != nil { + t.Fatal(err) + } + + requestedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC) + for index := 1; index <= 1002; index++ { + record := collection.Record{ + RequestedAt: requestedAt.Add(time.Duration(index) * time.Second), + Model: fmt.Sprintf("model-%d", index), + TotalTokens: int64(index), + } + if err := store.Insert(context.Background(), record); err != nil { + t.Fatalf("insert record %d: %v", index, err) + } + } + + recent, err := store.ListRecent(context.Background(), 1000) + if err != nil { + t.Fatal(err) + } + if len(recent) != 1000 || recent[0].Model != "model-1002" || recent[999].Model != "model-3" { + t.Fatalf("unexpected recent records: len=%d first=%q last=%q", len(recent), recent[0].Model, recent[len(recent)-1].Model) + } + if err := store.Close(); err != nil { + t.Fatal(err) + } + + reopened, err := repository.OpenSQLiteUsage(databasePath) + if err != nil { + t.Fatal(err) + } + defer reopened.Close() + all, err := reopened.ListRecent(context.Background(), 2000) + if err != nil { + t.Fatal(err) + } + if len(all) != 1002 { + t.Fatalf("persisted records = %d, want 1002", len(all)) + } +} + +func TestSQLiteUsageRoundTripsNullableCostAndDurations(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + cost := int64(125_000) + want := collection.Record{ + RequestedAt: time.Date(2026, 8, 14, 12, 0, 0, 123, time.FixedZone("CST", 8*60*60)), + APIKey: "key", + Model: "model", + Failed: true, + RequestType: "SSE", + Endpoint: "POST /v1/responses", + InputTokens: 10, + CacheReadTokens: 8, + TTFT: 250 * time.Millisecond, + Latency: 2 * time.Second, + CostMicros: &cost, + PriceTier: "base", + PriceMultiplierNumerator: 1, + PriceMultiplierDenominator: 1, + ClientIP: "192.0.2.10", + } + if err := store.Insert(context.Background(), want); err != nil { + t.Fatal(err) + } + records, err := store.ListRecent(context.Background(), 1) + if err != nil { + t.Fatal(err) + } + got := records[0] + if !got.RequestedAt.Equal(want.RequestedAt) || got.APIKey != want.APIKey || got.RequestType != want.RequestType || got.Endpoint != want.Endpoint || got.ClientIP != want.ClientIP || got.TTFT != want.TTFT || got.Latency != want.Latency || got.CostMicros == nil || *got.CostMicros != cost || got.PriceTier != "base" { + t.Fatalf("round trip mismatch: %+v", got) + } +} diff --git a/internal/web/embed.go b/internal/web/embed.go new file mode 100644 index 0000000..bb703d9 --- /dev/null +++ b/internal/web/embed.go @@ -0,0 +1,12 @@ +// Package web 提供插件管理页面的静态资源。 +package web + +import _ "embed" + +//go:embed ui.html +var uiHTML []byte + +// UI 返回只读的管理页面内容。 +func UI() []byte { + return uiHTML +} diff --git a/internal/web/ui.html b/internal/web/ui.html new file mode 100644 index 0000000..4a0fd6a --- /dev/null +++ b/internal/web/ui.html @@ -0,0 +1,386 @@ + + + + + + CPA 用量记录 + + + +
+

最近用量记录

+
+ +
+ 选择列 +
+
+ +
+
+
+ 价格设置 +
+
+ + + + + + + + + + + + + + +
+
+
+
+
+
+
+
+ + + +
+
+ + +