diff --git a/README.md b/README.md index a26d504..93056c5 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,8 @@ `cpa-ext` 是 CLIProxyAPI 的收敛式核心扩展，提供持久化用量与价格展示、下游 Key 管理、模型准入和上游凭证定向路由。一个 Key 代表一个调用用户；管理员可以创建、禁用或永久归档 Key，并按 Key 查看统计。 +每个 Key 同时拥有独立的美元额度账户。管理员可以分配本周期额度、设置日/周/月自动重置和并发上限，并查看不可变的扣费、额度调整与重置账目。额度采用请求结束后结算的软限制：余额为正时放行，实际 Usage 到达后扣费，最后一个或多个在途请求可能产生负余额，之后的新请求返回 429。 + ## 当前兼容目标 - CLIProxyAPI 源码：`CLIProxyAPI/`，检查时 revision 为 `f43aad7637ad813745bf7d341acb5663617570c5` @@ -51,12 +53,23 @@ go test ./... - `active` / `disabled` 状态切换，以及不可恢复但保留历史的归档； - CPA 自动选择上游，或严格绑定一个 OAuth/API Key 上游凭证； - 按客户端请求模型配置允许列表； +- 模型允许列表支持精确名称和 `*` 通配符，例如 `deepseek-*`； +- 分配美元额度、设置自动重置时间和每用户并发上限； +- 查看当前额度、已用、余额以及最近扣费账目； - 累计、今日和最近请求统计。 严格绑定的账号不可用或不支持目标模型时请求直接失败，不回退到其他账号。新 Key 在创建时复制 `default` 当时的路由与模型规则，之后独立维护。 +现有和新建 Key 的初始额度都是 `$0`，默认不自动重置、并发上限为 4。管理员分配额度后才能发起模型请求。金额以微美元整数保存；修改额度不会清空本周期已用金额，手动或自动重置不会结转旧余额。允许模型没有价格配置时，请求会在触达上游前以 503 拒绝。 + `config_yaml` 中包含宿主补充的 `enabled` 和 `priority`；插件会解析 `codex_only`、数据库路径和首次导入 Key。完整下游 Key 只通过受 CPA Management Key 保护的 Key 管理接口返回，不写入普通日志或错误消息；上游 Token、Cookie 和原始凭证不会由插件读取或保存。 +## 请求明细与汇总 + +请求明细由服务端分页，每页 100 条，默认按请求时间和稳定 ID 倒序展示。管理接口支持时间范围、用户 Key、模型、结果、上游 Auth ID、端点和完整 Request ID 筛选；连续翻页使用与筛选条件绑定的游标，数字页码跳转使用 SQLite 索引定位。用户页的今日汇总、各用户用量和近 7 日 Token 由数据库独立聚合，不受当前明细页影响。 + +升级已有数据库时，插件会从原始 Usage 和请求终态事实自动建立轻量查询投影。事实表、计费账目和历史统计保持不变；重试执行继续分别展示，没有 Usage 的取消或失败请求仍然可见，缺少 Request ID 的旧回调沿用双向唯一时间匹配规则。 + ## 工程布局 - `cmd/cpa-ext`：仅负责 C ABI、请求字节复制和 C 内存释放。 diff --git a/TODOS.md b/TODOS.md index e69de29..234aea3 100644 --- a/TODOS.md +++ b/TODOS.md @@ -0,0 +1,7 @@ +# 计划 + +- [x] 持久化用量、价格与管理面板 +- [x] 一个 Key 对应一个用户的认证、模型权限与上游路由 +- [x] 测试模式与真实数据隔离的紧凑管理台 +- [x] 美元额度、周期重置、并发限制与不可变账目 +- [x] 请求明细的服务端分页与筛选 diff --git a/internal/access/types.go b/internal/access/types.go index 5fe24b0..2d2f8ec 100644 --- a/internal/access/types.go +++ b/internal/access/types.go @@ -15,6 +15,11 @@ const ( RouteAuto = "auto" RouteStrict = "strict" + + ResetNone = "none" + ResetDaily = "daily" + ResetWeekly = "weekly" + ResetMonthly = "monthly" ) type ManagedKey struct { @@ -59,6 +64,40 @@ type UsageSummary struct { CostMicros int64 `json:"cost_micros"` } +type BillingState struct { + KeyID string `json:"key_id"` + QuotaMicros int64 `json:"-"` + SpentMicros int64 `json:"-"` + BalanceMicros int64 `json:"-"` + ResetPeriod string `json:"reset_period"` + NextResetAt *time.Time `json:"next_reset_at,omitempty"` + MaxConcurrency int `json:"max_concurrency"` + ActiveRequests int `json:"active_requests"` + CycleSequence int64 `json:"cycle_sequence"` + CycleStartedAt time.Time `json:"cycle_started_at"` +} + +type BillingSettings struct { + QuotaMicros int64 + ResetPeriod string + NextResetAt *time.Time + MaxConcurrency int +} + +type LedgerEntry struct { + ID int64 `json:"id"` + KeyID string `json:"key_id"` + CycleSequence int64 `json:"cycle_sequence"` + Kind string `json:"kind"` + AmountMicros int64 `json:"-"` + BalanceAfterMicros int64 `json:"-"` + RequestID string `json:"request_id,omitempty"` + ExecutionID string `json:"execution_id,omitempty"` + Model string `json:"model,omitempty"` + OccurredAt time.Time `json:"occurred_at"` + CreatedAt time.Time `json:"created_at"` +} + // CallerScope mirrors CLIProxyAPI's stable downstream principal namespace. func CallerScope(principal string) string { principal = strings.TrimSpace(principal) diff --git a/internal/collection/query.go b/internal/collection/query.go new file mode 100644 index 0000000..c56fe9b --- /dev/null +++ b/internal/collection/query.go @@ -0,0 +1,62 @@ +package collection + +import "time" + +const ( + UsageResultSucceeded = "succeeded" + UsageResultFailed = "failed" + UsageResultRejected = "rejected" + UsageResultCanceled = "canceled" +) + +// UsageQuery describes one stable page of request details. Cursor is opaque to +// callers and is bound to every filter in the query. +type UsageQuery struct { + Page int + PageSize int + Cursor string + From *time.Time + To *time.Time + KeyID string + Model string + Result string + AuthID string + Endpoint string + RequestID string +} + +type UsagePage struct { + Records []Record + Page int + PageSize int + Total int64 + TotalPages int + PreviousCursor string + NextCursor string +} + +type UsageDashboard struct { + Today UsageSummary + Users []UserUsageSummary + Days []DailyUsageSummary +} + +type UsageSummary struct { + Requests int64 + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + CostMicros int64 +} + +type UserUsageSummary struct { + KeyID string + KeyAlias string + Today UsageSummary + LastUsedAt *time.Time +} + +type DailyUsageSummary struct { + Date time.Time + TotalTokens int64 +} diff --git a/internal/collection/service.go b/internal/collection/service.go index 04f3be3..79ff0d5 100644 --- a/internal/collection/service.go +++ b/internal/collection/service.go @@ -1,12 +1,17 @@ package collection -import "context" +import ( + "context" + "time" +) // Repository 是采集模块需要的最小持久化接口。 type Repository interface { Insert(context.Context, Record) error UpsertRequest(context.Context, RequestRecord) error ListRecent(context.Context, int) ([]Record, error) + QueryUsage(context.Context, UsageQuery) (UsagePage, error) + UsageDashboard(context.Context, time.Time, int) (UsageDashboard, error) Close() error } @@ -31,6 +36,14 @@ func (s *Service) Recent(ctx context.Context, limit int) ([]Record, error) { return s.repository.ListRecent(ctx, limit) } +func (s *Service) Query(ctx context.Context, query UsageQuery) (UsagePage, error) { + return s.repository.QueryUsage(ctx, query) +} + +func (s *Service) Dashboard(ctx context.Context, today time.Time, days int) (UsageDashboard, error) { + return s.repository.UsageDashboard(ctx, today, days) +} + func (s *Service) Close() error { return s.repository.Close() } diff --git a/internal/plugin/access.go b/internal/plugin/access.go index b6abe28..6fc7741 100644 --- a/internal/plugin/access.go +++ b/internal/plugin/access.go @@ -10,6 +10,8 @@ import ( "time" managedaccess "cpa-ext/internal/access" + "cpa-ext/internal/pricing" + "cpa-ext/internal/repository" ) const ( @@ -93,8 +95,32 @@ func (a *App) interceptRequest(raw []byte, afterAuth bool) ([]byte, error) { if !managedKeyAllowsModel(key, requestedModel) { return OKEnvelope(denyRequest(http.StatusForbidden, "model_not_allowed", "该 Key 不允许使用此模型")) } - if req.RequestID != "" { - a.pending.Store(req.RequestID, key.ID) + if !afterAuth { + store, ok := a.currentStore() + if !ok { + return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_unavailable", "额度数据库不可用")) + } + _, authorizeErr := store.AuthorizeBilling(context.Background(), key.ID, req.RequestID, time.Now().UTC()) + switch { + case errors.Is(authorizeErr, repository.ErrBillingQuotaExhausted): + return OKEnvelope(denyRequest(http.StatusTooManyRequests, "billing_quota_exhausted", "该 Key 的可用余额已用尽")) + case errors.Is(authorizeErr, repository.ErrBillingConcurrency): + return OKEnvelope(denyRequest(http.StatusTooManyRequests, "billing_concurrency_exceeded", "该 Key 的并发请求已达到上限")) + case authorizeErr != nil: + return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_unavailable", "额度检查暂时不可用")) + } + if req.RequestID != "" { + a.pending.Store(req.RequestID, key.ID) + } + } + if afterAuth { + billingModel := strings.TrimSpace(req.Model) + if billingModel == "" { + billingModel = requestedModel + } + if _, found := a.priceForModel(billingModel); !found { + return OKEnvelope(denyRequest(http.StatusServiceUnavailable, "billing_price_unavailable", "该模型尚未配置价格")) + } } if afterAuth && key.RouteMode == managedaccess.RouteStrict { account, accountErr := a.upstreamAccount(key.UpstreamAccountID) @@ -174,16 +200,62 @@ func managedKeyAllowsModel(key managedaccess.ManagedKey, model string) bool { } model = strings.ToLower(strings.TrimSpace(model)) if model == "" { - return true + return false } for _, allowed := range key.Models { - if strings.EqualFold(strings.TrimSpace(allowed), model) { + if modelPatternMatches(allowed, model) { return true } } return false } +func modelPatternMatches(pattern, model string) bool { + pattern = normalizeModelName(pattern) + model = normalizeModelName(model) + if pattern == "" || model == "" { + return false + } + if pattern == "*" { + return true + } + parts := strings.Split(pattern, "*") + position := 0 + if !strings.HasPrefix(pattern, "*") { + if !strings.HasPrefix(model, parts[0]) { + return false + } + position = len(parts[0]) + parts = parts[1:] + } + for index, part := range parts { + if part == "" { + continue + } + last := index == len(parts)-1 && !strings.HasSuffix(pattern, "*") + if last { + return strings.HasSuffix(model[position:], part) + } + offset := strings.Index(model[position:], part) + if offset < 0 { + return false + } + position += offset + len(part) + } + return strings.HasSuffix(pattern, "*") || position == len(model) +} + +func normalizeModelName(model string) string { + return strings.ToLower(strings.TrimSpace(model)) +} + +func (a *App) priceForModel(model string) (pricing.Policy, bool) { + a.mu.RLock() + defer a.mu.RUnlock() + policy, found := a.prices[normalizeModelName(model)] + return policy, found +} + func denyRequest(status int, code, message string) RequestInterceptResponse { body, _ := json.Marshal(map[string]any{"error": map[string]string{"code": code, "message": message}}) return RequestInterceptResponse{ diff --git a/internal/plugin/access_test.go b/internal/plugin/access_test.go index 475b05a..b3b4d19 100644 --- a/internal/plugin/access_test.go +++ b/internal/plugin/access_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "net/http" + "strings" "testing" "time" @@ -62,6 +63,10 @@ func TestModelAllowlistAndStrictUpstreamRouting(t *testing.T) { if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil { t.Fatal(err) } + price := []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, price); response.StatusCode != http.StatusOK { + t.Fatalf("put price=%d %s", response.StatusCode, response.Body) + } store, _ := app.currentStore() if err := store.SyncUpstreamAccounts(context.Background(), []managedaccess.UpstreamAccount{{ CPAAuthID: "oauth-1", CPAAuthIndex: "index-1", Provider: "codex", DisplayName: "OAuth One", LastSeenAt: time.Now(), @@ -109,7 +114,7 @@ func TestModelAllowlistAndStrictUpstreamRouting(t *testing.T) { t.Fatalf("scheduler pick = %+v", pick) } afterRaw, _ := json.Marshal(RequestInterceptRequest{ - RequestID: "req-1", RequestedModel: "gpt-5.6-sol", + RequestID: "req-1", Model: "gpt-5.6-sol", RequestedModel: "gpt-5.6-sol", Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID), selectedAuthMetadata: "oauth-1"}, }) afterResponse, _ := app.HandleMethod(MethodRequestAfter, afterRaw) @@ -133,6 +138,20 @@ func TestModelAllowlistAndStrictUpstreamRouting(t *testing.T) { } } +func TestModelAllowlistSupportsFullNameWildcards(t *testing.T) { + key := managedaccess.ManagedKey{AllModels: false, Models: []string{"deepseek-*", "gpt-5.6-sol"}} + for _, model := range []string{"deepseek-v4-flash", "DEEPSEEK-SOURCE-B", "gpt-5.6-sol"} { + if !managedKeyAllowsModel(key, model) { + t.Fatalf("expected %q to match", model) + } + } + for _, model := range []string{"x-deepseek-v4-flash", "gpt-5.6-terra", ""} { + if managedKeyAllowsModel(key, model) { + t.Fatalf("expected %q to be denied", model) + } + } +} + func TestCreateKeyCopiesDefaultRuleSnapshot(t *testing.T) { app := NewApp() if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\n"))); err != nil { @@ -150,3 +169,70 @@ func TestCreateKeyCopiesDefaultRuleSnapshot(t *testing.T) { t.Fatalf("created = %+v", created) } } + +func TestBillingAdmissionConcurrencyAndUsageSettlement(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) + } + price := []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":false,"fast_multiplier":"2.5"}`) + if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, price); response.StatusCode != http.StatusOK { + t.Fatalf("put price=%d %s", response.StatusCode, response.Body) + } + key, _ := app.store.ManagedKeyByID(context.Background(), "key_default") + key.AllModels = false + key.Models = []string{"deepseek-*"} + if err := app.store.UpdateManagedKey(context.Background(), key); err != nil { + t.Fatal(err) + } + intercept := func(requestID string) RequestInterceptResponse { + raw, _ := json.Marshal(RequestInterceptRequest{RequestID: requestID, Model: "deepseek-v4-flash", RequestedModel: "deepseek-source-a", Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}}) + response, err := app.HandleMethod(MethodRequestBefore, raw) + if err != nil { + t.Fatal(err) + } + var envelope Envelope + _ = json.Unmarshal(response, &envelope) + var result RequestInterceptResponse + _ = json.Unmarshal(envelope.Result, &result) + return result + } + if denied := intercept("zero"); !denied.Terminate || denied.StatusCode != http.StatusTooManyRequests || !strings.Contains(string(denied.ResponseBody), "billing_quota_exhausted") { + t.Fatalf("zero quota response=%+v body=%s", denied, denied.ResponseBody) + } + patch := []byte(`{"id":"key_default","billing":{"quota_usd":"1","reset_period":"none","max_concurrency":1}}`) + if response := managementCallBody(t, app, http.MethodPatch, managementBase+routeKeys, patch); response.StatusCode != http.StatusOK { + t.Fatalf("update billing=%d %s", response.StatusCode, response.Body) + } + if allowed := intercept("one"); allowed.Terminate { + t.Fatalf("funded request denied: %+v", allowed) + } + afterRaw, _ := json.Marshal(RequestInterceptRequest{RequestID: "one", Model: "unpriced-resolved-model", RequestedModel: "deepseek-source-a", Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}}) + afterResponse, _ := app.HandleMethod(MethodRequestAfter, afterRaw) + var afterEnvelope Envelope + _ = json.Unmarshal(afterResponse, &afterEnvelope) + var afterResult RequestInterceptResponse + _ = json.Unmarshal(afterEnvelope.Result, &afterResult) + if !afterResult.Terminate || afterResult.StatusCode != http.StatusServiceUnavailable || !strings.Contains(string(afterResult.ResponseBody), "billing_price_unavailable") { + t.Fatalf("missing price response=%+v body=%s", afterResult, afterResult.ResponseBody) + } + if denied := intercept("two"); !denied.Terminate || !strings.Contains(string(denied.ResponseBody), "billing_concurrency_exceeded") { + t.Fatalf("concurrency response=%+v body=%s", denied, denied.ResponseBody) + } + completionRaw, _ := json.Marshal(RequestCompletion{RequestID: "one", StartedAt: time.Now().Add(-time.Second), CompletedAt: time.Now(), Outcome: RequestCompletionSucceeded, Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}}) + if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err != nil { + t.Fatal(err) + } + usage := UsageRecord{ExecutionID: "settled-once", APIKey: "000000", Model: "deepseek-v4-flash", RequestedAt: time.Now(), Detail: UsageDetail{OutputTokens: 10_000, TotalTokens: 10_000}} + raw, _ := json.Marshal(usage) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + state, err := app.store.BillingState(context.Background(), key.ID, time.Now()) + if err != nil || state.SpentMicros != 150_000 || state.BalanceMicros != 850_000 || state.ActiveRequests != 0 { + t.Fatalf("settled state=%+v err=%v", state, err) + } +} diff --git a/internal/plugin/app.go b/internal/plugin/app.go index 0655521..63b9b5f 100644 --- a/internal/plugin/app.go +++ b/internal/plugin/app.go @@ -8,6 +8,7 @@ import ( "strings" "sync" "sync/atomic" + "time" "cpa-ext/internal/collection" "cpa-ext/internal/pricing" @@ -100,7 +101,16 @@ func (a *App) configure(raw []byte) ([]byte, error) { } nextPrices := make(map[string]pricing.Policy, len(policies)) for _, policy := range policies { - nextPrices[policy.Model] = policy + nextPrices[normalizeModelName(policy.Model)] = policy + } + a.mu.RLock() + firstInitialization := a.usage == nil + a.mu.RUnlock() + if firstInitialization { + if err := usageRepository.ReleaseStaleAdmissions(context.Background(), time.Now().UTC()); err != nil { + _ = nextUsage.Close() + return nil, fmt.Errorf("释放遗留并发占用: %w", err) + } } a.mu.Lock() @@ -191,7 +201,7 @@ func (a *App) handleUsage(raw []byte) ([]byte, error) { observed.ManagedKeyID = key.ID observed.KeyAlias = key.Name } - if policy, found := a.prices[strings.TrimSpace(record.Model)]; found { + if policy, found := a.prices[normalizeModelName(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, @@ -227,6 +237,9 @@ func (a *App) handleRequestComplete(raw []byte) ([]byte, error) { if a.usage == nil { return nil, fmt.Errorf("用量数据库尚未初始化") } + if err := a.store.CompleteBillingAdmission(context.Background(), completion.RequestID, completion.CompletedAt); err != nil { + return nil, fmt.Errorf("释放并发占用: %w", err) + } endpoint, _ := completion.Metadata["request_path"].(string) managedKeyID := metadataString(completion.Metadata, callerScopeMetadata) if key, keyErr := a.store.ManagedKeyByReference(context.Background(), managedKeyID); keyErr == nil { diff --git a/internal/plugin/key_management.go b/internal/plugin/key_management.go index c922384..7df2dec 100644 --- a/internal/plugin/key_management.go +++ b/internal/plugin/key_management.go @@ -7,31 +7,70 @@ import ( "encoding/json" "errors" "net/http" + "net/url" + "strconv" "strings" "time" "unicode" managedaccess "cpa-ext/internal/access" + "cpa-ext/internal/collection" "cpa-ext/internal/repository" ) type createManagedKeyRequest struct { - Name string `json:"name"` - Secret string `json:"secret,omitempty"` - RouteMode string `json:"route_mode,omitempty"` - UpstreamAccountID string `json:"upstream_account_id,omitempty"` - AllModels *bool `json:"all_models,omitempty"` - Models *[]string `json:"models,omitempty"` + Name string `json:"name"` + Secret string `json:"secret,omitempty"` + RouteMode string `json:"route_mode,omitempty"` + UpstreamAccountID string `json:"upstream_account_id,omitempty"` + AllModels *bool `json:"all_models,omitempty"` + Models *[]string `json:"models,omitempty"` + Billing *billingSettingsRequest `json:"billing,omitempty"` } type updateManagedKeyRequest struct { - ID string `json:"id"` - Name *string `json:"name,omitempty"` - Status *string `json:"status,omitempty"` - RouteMode *string `json:"route_mode,omitempty"` - UpstreamAccountID *string `json:"upstream_account_id,omitempty"` - AllModels *bool `json:"all_models,omitempty"` - Models *[]string `json:"models,omitempty"` + ID string `json:"id"` + Name *string `json:"name,omitempty"` + Status *string `json:"status,omitempty"` + RouteMode *string `json:"route_mode,omitempty"` + UpstreamAccountID *string `json:"upstream_account_id,omitempty"` + AllModels *bool `json:"all_models,omitempty"` + Models *[]string `json:"models,omitempty"` + Billing *billingSettingsRequest `json:"billing,omitempty"` +} + +type billingSettingsRequest struct { + QuotaUSD string `json:"quota_usd"` + ResetPeriod string `json:"reset_period"` + NextResetAt *time.Time `json:"next_reset_at,omitempty"` + MaxConcurrency int `json:"max_concurrency"` +} + +type billingStateDTO struct { + QuotaUSD string `json:"quota_usd"` + SpentUSD string `json:"spent_usd"` + BalanceUSD string `json:"balance_usd"` + ResetPeriod string `json:"reset_period"` + NextResetAt *time.Time `json:"next_reset_at,omitempty"` + MaxConcurrency int `json:"max_concurrency"` + ActiveRequests int `json:"active_requests"` + CycleStartedAt time.Time `json:"cycle_started_at"` +} + +type managedKeyDTO struct { + managedaccess.ManagedKey + Billing billingStateDTO `json:"billing"` +} + +type ledgerEntryDTO struct { + ID int64 `json:"id"` + Kind string `json:"kind"` + AmountUSD string `json:"amount_usd"` + BalanceAfterUSD string `json:"balance_after_usd"` + RequestID string `json:"request_id,omitempty"` + ExecutionID string `json:"execution_id,omitempty"` + Model string `json:"model,omitempty"` + OccurredAt time.Time `json:"occurred_at"` } type archiveManagedKeyRequest struct { @@ -52,7 +91,15 @@ func (a *App) listManagedKeys(includeArchived bool) ManagementResponse { if err != nil { return managementError(http.StatusInternalServerError, "database_error", err.Error()) } - return jsonManagementResponse(http.StatusOK, map[string]any{"keys": keys}) + items := make([]managedKeyDTO, 0, len(keys)) + for _, key := range keys { + state, stateErr := store.BillingState(context.Background(), key.ID, time.Now().UTC()) + if stateErr != nil { + return managementError(http.StatusInternalServerError, "database_error", stateErr.Error()) + } + items = append(items, managedKeyResponse(key, state)) + } + return jsonManagementResponse(http.StatusOK, map[string]any{"keys": items}) } func (a *App) createManagedKey(body []byte) ManagementResponse { @@ -117,11 +164,25 @@ func (a *App) createManagedKey(body []byte) ManagementResponse { if err := validateManagedKeyRule(store, key); err != nil { return managementError(http.StatusBadRequest, "invalid_rule", err.Error()) } + var requestedBilling *managedaccess.BillingSettings + if req.Billing != nil { + settings, settingsErr := billingSettings(*req.Billing) + if settingsErr != nil { + return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error()) + } + requestedBilling = &settings + } if err := store.CreateManagedKey(context.Background(), key); err != nil { return managementError(http.StatusConflict, "key_conflict", err.Error()) } + if requestedBilling != nil { + if _, settingsErr := store.UpdateBilling(context.Background(), key.ID, *requestedBilling, time.Now().UTC()); settingsErr != nil { + return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error()) + } + } created, _ := store.ManagedKeyByID(context.Background(), key.ID) - return jsonManagementResponse(http.StatusCreated, created) + state, _ := store.BillingState(context.Background(), key.ID, time.Now().UTC()) + return jsonManagementResponse(http.StatusCreated, managedKeyResponse(created, state)) } func (a *App) updateManagedKey(body []byte) ManagementResponse { @@ -170,6 +231,14 @@ func (a *App) updateManagedKey(body []byte) ManagementResponse { if err := validateManagedKeyRule(store, key); err != nil { return managementError(http.StatusBadRequest, "invalid_rule", err.Error()) } + var requestedBilling *managedaccess.BillingSettings + if req.Billing != nil { + settings, settingsErr := billingSettings(*req.Billing) + if settingsErr != nil { + return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error()) + } + requestedBilling = &settings + } if err := store.UpdateManagedKey(context.Background(), key); err != nil { status := http.StatusConflict if errors.Is(err, repository.ErrManagedKeyNotFound) { @@ -177,8 +246,96 @@ func (a *App) updateManagedKey(body []byte) ManagementResponse { } return managementError(status, "update_failed", err.Error()) } + if requestedBilling != nil { + if _, settingsErr := store.UpdateBilling(context.Background(), key.ID, *requestedBilling, time.Now().UTC()); settingsErr != nil { + return managementError(http.StatusBadRequest, "invalid_billing", settingsErr.Error()) + } + } updated, _ := store.ManagedKeyByID(context.Background(), key.ID) - return jsonManagementResponse(http.StatusOK, updated) + state, _ := store.BillingState(context.Background(), key.ID, time.Now().UTC()) + return jsonManagementResponse(http.StatusOK, managedKeyResponse(updated, state)) +} + +func billingSettings(request billingSettingsRequest) (managedaccess.BillingSettings, error) { + quota, err := parseDecimalMicros(request.QuotaUSD) + if err != nil { + return managedaccess.BillingSettings{}, errors.New("额度必须是非负美元金额，最多六位小数") + } + period := strings.TrimSpace(request.ResetPeriod) + if period == "" { + period = managedaccess.ResetNone + } + if period != managedaccess.ResetNone && period != managedaccess.ResetDaily && period != managedaccess.ResetWeekly && period != managedaccess.ResetMonthly { + return managedaccess.BillingSettings{}, errors.New("重置周期只能是 none、daily、weekly 或 monthly") + } + if period != managedaccess.ResetNone && request.NextResetAt != nil && !request.NextResetAt.After(time.Now()) { + return managedaccess.BillingSettings{}, errors.New("下次重置时间必须晚于当前时间") + } + maxConcurrency := request.MaxConcurrency + if maxConcurrency == 0 { + maxConcurrency = 4 + } + if maxConcurrency < 1 || maxConcurrency > 64 { + return managedaccess.BillingSettings{}, errors.New("并发上限必须为 1-64") + } + return managedaccess.BillingSettings{QuotaMicros: quota, ResetPeriod: period, + NextResetAt: request.NextResetAt, MaxConcurrency: maxConcurrency}, nil +} + +func managedKeyResponse(key managedaccess.ManagedKey, state managedaccess.BillingState) managedKeyDTO { + return managedKeyDTO{ManagedKey: key, Billing: billingStateDTO{ + QuotaUSD: formatMicros(state.QuotaMicros), SpentUSD: formatMicros(state.SpentMicros), + BalanceUSD: formatMicros(state.BalanceMicros), ResetPeriod: state.ResetPeriod, + NextResetAt: state.NextResetAt, MaxConcurrency: state.MaxConcurrency, + ActiveRequests: state.ActiveRequests, CycleStartedAt: state.CycleStartedAt, + }} +} + +func (a *App) resetManagedKeyBilling(body []byte) ManagementResponse { + var request struct { + ID string `json:"id"` + } + if err := json.Unmarshal(body, &request); err != nil || strings.TrimSpace(request.ID) == "" { + return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID") + } + store, ok := a.currentStore() + if !ok { + return managementError(http.StatusServiceUnavailable, "database_unavailable", "额度数据库尚未初始化") + } + state, err := store.ResetBilling(context.Background(), request.ID, time.Now().UTC()) + if err != nil { + return managementError(http.StatusBadRequest, "reset_failed", err.Error()) + } + return jsonManagementResponse(http.StatusOK, billingStateDTO{ + QuotaUSD: formatMicros(state.QuotaMicros), SpentUSD: formatMicros(state.SpentMicros), + BalanceUSD: formatMicros(state.BalanceMicros), ResetPeriod: state.ResetPeriod, + NextResetAt: state.NextResetAt, MaxConcurrency: state.MaxConcurrency, + ActiveRequests: state.ActiveRequests, CycleStartedAt: state.CycleStartedAt, + }) +} + +func (a *App) managedKeyLedger(query url.Values) ManagementResponse { + keyID := strings.TrimSpace(query.Get("id")) + if keyID == "" { + return managementError(http.StatusBadRequest, "invalid_request", "缺少有效的 Key ID") + } + before, _ := strconv.ParseInt(query.Get("before"), 10, 64) + limit, _ := strconv.Atoi(query.Get("limit")) + store, ok := a.currentStore() + if !ok { + return managementError(http.StatusServiceUnavailable, "database_unavailable", "额度数据库尚未初始化") + } + entries, err := store.ListBillingLedger(context.Background(), keyID, before, limit) + if err != nil { + return managementError(http.StatusInternalServerError, "database_error", err.Error()) + } + items := make([]ledgerEntryDTO, 0, len(entries)) + for _, entry := range entries { + items = append(items, ledgerEntryDTO{ID: entry.ID, Kind: entry.Kind, + AmountUSD: formatMicros(entry.AmountMicros), BalanceAfterUSD: formatMicros(entry.BalanceAfterMicros), + RequestID: entry.RequestID, ExecutionID: entry.ExecutionID, Model: entry.Model, OccurredAt: entry.OccurredAt}) + } + return jsonManagementResponse(http.StatusOK, map[string]any{"entries": items}) } func (a *App) archiveManagedKey(body []byte) ManagementResponse { @@ -211,18 +368,13 @@ func (a *App) managedKeyStats(id string) ManagementResponse { if err != nil { return managementError(http.StatusInternalServerError, "database_error", err.Error()) } - records, err := store.ListRecent(context.Background(), 1000) + records, err := store.QueryUsage(context.Background(), collection.UsageQuery{Page: 1, PageSize: 50, KeyID: id}) if err != nil { return managementError(http.StatusInternalServerError, "database_error", err.Error()) } - recent := make([]usageListItem, 0, 50) - for _, record := range records { - if record.ManagedKeyID == id { - recent = append(recent, usageItem(record)) - if len(recent) == 50 { - break - } - } + recent := make([]usageListItem, 0, len(records.Records)) + for _, record := range records.Records { + recent = append(recent, usageItem(record)) } return jsonManagementResponse(http.StatusOK, managedKeyStatsResponse{Stats: stats, Recent: recent}) } @@ -282,6 +434,11 @@ func validateManagedKeyRule(store *repository.SQLiteUsageRepository, key managed if !key.AllModels && len(normalizeModels(key.Models)) == 0 { return errors.New("未允许全部模型时至少配置一个模型") } + for _, pattern := range key.Models { + if len(pattern) > 128 || strings.IndexFunc(pattern, func(value rune) bool { return unicode.IsSpace(value) || unicode.IsControl(value) }) >= 0 { + return errors.New("模型规则不能包含空白或控制字符，且最长 128 个字符") + } + } return nil } diff --git a/internal/plugin/management.go b/internal/plugin/management.go index 506eba5..9c227b7 100644 --- a/internal/plugin/management.go +++ b/internal/plugin/management.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "net/http" + "net/url" + "strconv" "strings" "time" @@ -13,21 +15,25 @@ import ( ) const ( - managementBase = "/v0/management/plugins/" + PluginName - resourceBase = "/v0/resource/plugins/" + PluginName - routeUsage = "/usage" - routePrices = "/prices" - routeKeys = "/keys" - routeKeyStats = "/key-stats" - routeUpstreams = "/upstreams" - routeModels = "/model-suggestions" - resourceUI = "/ui" + managementBase = "/v0/management/plugins/" + PluginName + resourceBase = "/v0/resource/plugins/" + PluginName + routeUsage = "/usage" + routeUsageSummary = "/usage-summary" + routePrices = "/prices" + routeKeys = "/keys" + routeKeyStats = "/key-stats" + routeUpstreams = "/upstreams" + routeModels = "/model-suggestions" + routeBillingReset = "/billing-reset" + routeBillingLedger = "/billing-ledger" + resourceUI = "/ui" ) func managementRegistration() ManagementRegistrationResponse { return ManagementRegistrationResponse{ Routes: []ManagementRoute{ {Method: http.MethodGet, Path: managementBase + routeUsage, Description: "查看最近的用量记录。"}, + {Method: http.MethodGet, Path: managementBase + routeUsageSummary, Description: "查看用户与每日用量汇总。"}, {Method: http.MethodGet, Path: managementBase + routePrices, Description: "查看模型价格。"}, {Method: http.MethodPut, Path: managementBase + routePrices, Description: "保存模型价格。"}, {Method: http.MethodDelete, Path: managementBase + routePrices, Description: "删除模型价格。"}, @@ -38,6 +44,8 @@ func managementRegistration() ManagementRegistrationResponse { {Method: http.MethodGet, Path: managementBase + routeKeyStats, Description: "查看 Key 统计。"}, {Method: http.MethodGet, Path: managementBase + routeUpstreams, Description: "同步并查看上游账号。"}, {Method: http.MethodGet, Path: managementBase + routeModels, Description: "查看模型建议。"}, + {Method: http.MethodPost, Path: managementBase + routeBillingReset, Description: "立即重置 Key 额度。"}, + {Method: http.MethodGet, Path: managementBase + routeBillingLedger, Description: "查看 Key 额度账目。"}, }, Resources: []ResourceRoute{ {Path: resourceBase + resourceUI, Menu: "用量记录", Description: "查看 CPA 最近收到的请求用量。"}, @@ -62,7 +70,10 @@ func (a *App) handleManagement(raw []byte) ([]byte, error) { }) } if req.Method == http.MethodGet && path == managementBase+routeUsage { - return OKEnvelope(a.usageResponse()) + return OKEnvelope(a.usageResponse(req.Query)) + } + if req.Method == http.MethodGet && path == managementBase+routeUsageSummary { + return OKEnvelope(a.usageDashboardResponse()) } if path == managementBase+routePrices { switch req.Method { @@ -95,6 +106,12 @@ func (a *App) handleManagement(raw []byte) ([]byte, error) { if req.Method == http.MethodGet && path == managementBase+routeModels { return OKEnvelope(a.modelSuggestions()) } + if req.Method == http.MethodPost && path == managementBase+routeBillingReset { + return OKEnvelope(a.resetManagedKeyBilling(req.Body)) + } + if req.Method == http.MethodGet && path == managementBase+routeBillingLedger { + return OKEnvelope(a.managedKeyLedger(req.Query)) + } return OKEnvelope(jsonManagementResponse(http.StatusNotFound, map[string]any{ "error": map[string]string{ "code": "not_found", @@ -104,8 +121,17 @@ func (a *App) handleManagement(raw []byte) ([]byte, error) { } type usageListResponse struct { - Records []usageListItem `json:"records"` - Retained int `json:"retained"` + Records []usageListItem `json:"records"` + Pagination usagePaginationResult `json:"pagination"` +} + +type usagePaginationResult struct { + Page int `json:"page"` + PageSize int `json:"page_size"` + Total int64 `json:"total"` + TotalPages int `json:"total_pages"` + PreviousCursor string `json:"previous_cursor,omitempty"` + NextCursor string `json:"next_cursor,omitempty"` } // usageListItem 是请求明细表的稳定接口。当前回调拿不到的字段保留为空值。 @@ -147,7 +173,11 @@ type usageListItem struct { ClientIP string `json:"client_ip"` } -func (a *App) usageResponse() ManagementResponse { +func (a *App) usageResponse(values url.Values) ManagementResponse { + query, err := parseUsageQuery(values) + if err != nil { + return managementError(http.StatusBadRequest, "invalid_query", err.Error()) + } a.mu.RLock() defer a.mu.RUnlock() if a.usage == nil { @@ -155,17 +185,90 @@ func (a *App) usageResponse() ManagementResponse { "error": map[string]string{"code": "database_unavailable", "message": "用量数据库尚未初始化"}, }) } - records, err := a.usage.Recent(context.Background(), 1000) + page, err := a.usage.Query(context.Background(), query) if err != nil { - return jsonManagementResponse(http.StatusInternalServerError, map[string]any{ - "error": map[string]string{"code": "database_error", "message": err.Error()}, + status := http.StatusInternalServerError + code := "database_error" + if strings.Contains(err.Error(), "cursor") || strings.Contains(err.Error(), "page_size") || strings.Contains(err.Error(), "from") || strings.Contains(err.Error(), "result") { + status = http.StatusBadRequest + code = "invalid_query" + } + return jsonManagementResponse(status, map[string]any{ + "error": map[string]string{"code": code, "message": err.Error()}, }) } - items := make([]usageListItem, 0, len(records)) - for _, record := range records { + items := make([]usageListItem, 0, len(page.Records)) + for _, record := range page.Records { items = append(items, usageItem(record)) } - return jsonManagementResponse(http.StatusOK, usageListResponse{Records: items, Retained: 1000}) + return jsonManagementResponse(http.StatusOK, usageListResponse{Records: items, Pagination: usagePaginationResult{ + Page: page.Page, PageSize: page.PageSize, Total: page.Total, TotalPages: page.TotalPages, + PreviousCursor: page.PreviousCursor, NextCursor: page.NextCursor, + }}) +} + +func parseUsageQuery(values url.Values) (collection.UsageQuery, error) { + query := collection.UsageQuery{Page: 1, PageSize: 100, Cursor: strings.TrimSpace(values.Get("cursor")), KeyID: values.Get("key_id"), Model: values.Get("model"), Result: values.Get("result"), AuthID: values.Get("auth_id"), Endpoint: values.Get("endpoint"), RequestID: values.Get("request_id")} + if raw := strings.TrimSpace(values.Get("page")); raw != "" { + value, err := strconv.Atoi(raw) + if err != nil || value < 1 { + return query, fmt.Errorf("page 必须是正整数") + } + query.Page = value + } + if raw := strings.TrimSpace(values.Get("page_size")); raw != "" { + value, err := strconv.Atoi(raw) + if err != nil || value < 1 || value > 100 { + return query, fmt.Errorf("page_size 必须在 1 到 100 之间") + } + query.PageSize = value + } + for name, target := range map[string]**time.Time{"from": &query.From, "to": &query.To} { + raw := strings.TrimSpace(values.Get(name)) + if raw == "" { + continue + } + parsed, err := time.Parse(time.RFC3339Nano, raw) + if err != nil { + return query, fmt.Errorf("%s 必须是 RFC3339 时间", name) + } + *target = &parsed + } + return query, nil +} + +type usageSummaryDTO struct { + Requests int64 `json:"requests"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + TotalTokens int64 `json:"total_tokens"` + CostUSD float64 `json:"cost_usd"` +} + +func usageSummaryItem(summary collection.UsageSummary) usageSummaryDTO { + return usageSummaryDTO{Requests: summary.Requests, InputTokens: summary.InputTokens, OutputTokens: summary.OutputTokens, TotalTokens: summary.TotalTokens, CostUSD: float64(summary.CostMicros) / 1_000_000} +} + +func (a *App) usageDashboardResponse() ManagementResponse { + a.mu.RLock() + defer a.mu.RUnlock() + if a.usage == nil { + return managementError(http.StatusServiceUnavailable, "database_unavailable", "用量数据库尚未初始化") + } + location, _ := time.LoadLocation("Asia/Shanghai") + dashboard, err := a.usage.Dashboard(context.Background(), time.Now().In(location), 7) + if err != nil { + return managementError(http.StatusInternalServerError, "database_error", err.Error()) + } + users := make([]map[string]any, 0, len(dashboard.Users)) + for _, user := range dashboard.Users { + users = append(users, map[string]any{"key_id": user.KeyID, "key_alias": user.KeyAlias, "today": usageSummaryItem(user.Today), "last_used_at": user.LastUsedAt}) + } + days := make([]map[string]any, 0, len(dashboard.Days)) + for _, day := range dashboard.Days { + days = append(days, map[string]any{"date": day.Date.Format("2006-01-02"), "total_tokens": day.TotalTokens}) + } + return jsonManagementResponse(http.StatusOK, map[string]any{"today": usageSummaryItem(dashboard.Today), "users": users, "days": days}) } func usageItem(record collection.Record) usageListItem { diff --git a/internal/plugin/management_test.go b/internal/plugin/management_test.go index f1a61ac..1e2dc0e 100644 --- a/internal/plugin/management_test.go +++ b/internal/plugin/management_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "net/http" + "net/url" "strings" "testing" "time" @@ -14,8 +15,12 @@ func managementCall(t *testing.T, app *App, method, path string) ManagementRespo } func managementCallBody(t *testing.T, app *App, method, path string, body []byte) ManagementResponse { + return managementCallRequest(t, app, ManagementRequest{Method: method, Path: path, Body: body}) +} + +func managementCallRequest(t *testing.T, app *App, requestValue ManagementRequest) ManagementResponse { t.Helper() - request, err := json.Marshal(ManagementRequest{Method: method, Path: path, Body: body}) + request, err := json.Marshal(requestValue) if err != nil { t.Fatal(err) } @@ -47,7 +52,7 @@ func TestManagementRegistrationDeclaresUsageAPIAndUI(t *testing.T) { if err := json.Unmarshal(envelope.Result, ®istration); err != nil { t.Fatal(err) } - if len(registration.Routes) != 11 || registration.Routes[0].Path != managementBase+routeUsage || registration.Routes[1].Path != managementBase+routePrices { + if len(registration.Routes) != 14 || registration.Routes[0].Path != managementBase+routeUsage || registration.Routes[1].Path != managementBase+routeUsageSummary || registration.Routes[2].Path != managementBase+routePrices { t.Fatalf("unexpected management routes: %+v", registration.Routes) } if len(registration.Resources) != 1 || registration.Resources[0].Path != resourceBase+resourceUI { @@ -55,6 +60,49 @@ func TestManagementRegistrationDeclaresUsageAPIAndUI(t *testing.T) { } } +func TestUsageManagementSupportsServerPaginationFiltersAndSummary(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) + } + base := time.Date(2026, 8, 15, 8, 0, 0, 0, time.UTC) + for index := 0; index < 205; index++ { + record := UsageRecord{RequestID: fmt.Sprintf("page-request-%03d", index), ExecutionID: fmt.Sprintf("page-execution-%03d", index), Provider: "openai", Model: fmt.Sprintf("model-%d", index%2), AuthID: fmt.Sprintf("auth-%d", index%2), RequestedAt: base.Add(time.Duration(index) * time.Second), Endpoint: "POST /v1/responses", Failed: index%10 == 0, Detail: UsageDetail{TotalTokens: int64(index + 1)}} + raw, _ := json.Marshal(record) + if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil { + t.Fatal(err) + } + } + response := managementCallRequest(t, app, ManagementRequest{Method: http.MethodGet, Path: managementBase + routeUsage, Query: url.Values{"page": {"2"}, "page_size": {"100"}, "model": {"model-1"}}}) + if response.StatusCode != http.StatusOK { + t.Fatalf("usage query status=%d body=%s", response.StatusCode, response.Body) + } + var payload usageListResponse + if err := json.Unmarshal(response.Body, &payload); err != nil { + t.Fatal(err) + } + if payload.Pagination.Total != 102 || payload.Pagination.Page != 2 || payload.Pagination.TotalPages != 2 || len(payload.Records) != 2 { + t.Fatalf("unexpected paginated response: %+v", payload.Pagination) + } + if payload.Records[0].Model != "model-1" || payload.Records[1].Model != "model-1" { + t.Fatalf("filter leaked records: %+v", payload.Records) + } + invalid := managementCallRequest(t, app, ManagementRequest{Method: http.MethodGet, Path: managementBase + routeUsage, Query: url.Values{"page_size": {"101"}}}) + if invalid.StatusCode != http.StatusBadRequest { + t.Fatalf("invalid page size status=%d", invalid.StatusCode) + } + summary := managementCall(t, app, http.MethodGet, managementBase+routeUsageSummary) + if summary.StatusCode != http.StatusOK || !strings.Contains(string(summary.Body), `"users"`) || !strings.Contains(string(summary.Body), `"days"`) { + t.Fatalf("unexpected summary: status=%d body=%s", summary.StatusCode, summary.Body) + } +} + +func TestFormatMicrosSupportsNegativeBalances(t *testing.T) { + if got := formatMicros(-125_000); got != "-0.125" { + t.Fatalf("formatMicros(-125000)=%q", got) + } +} + 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 { @@ -387,7 +435,7 @@ func TestUsageResourceServesTablePage(t *testing.T) { t.Fatalf("UI does not contain managed access feature %q", feature) } } - for _, feature := range []string{"测试模式", `data-mode="demo"`, `id="page-buttons"`, "const PAGE_SIZE = 100", "length: 2370", `id="user-usage"`, `id="daily-chart"`, `overflow-y: hidden`, "测试模式禁止访问真实管理接口"} { + for _, feature := range []string{"测试模式", `data-mode="demo"`, `id="page-buttons"`, "const PAGE_SIZE = 100", "length: 2370", `id="user-usage"`, `id="daily-chart"`, `id="usage-from"`, `id="usage-request-id"`, `id="apply-usage-filters"`, "SUMMARY_API", `overflow-y: hidden`, "测试模式禁止访问真实管理接口", `id="editor-quota"`, `id="editor-reset-period"`, `id="billing-ledger"`, "deepseek-*", "请求结束后按实际费用扣款"} { if !strings.Contains(page, feature) { t.Fatalf("UI does not contain workspace feature %q", feature) } diff --git a/internal/plugin/pricing_management.go b/internal/plugin/pricing_management.go index a024168..7a03b8b 100644 --- a/internal/plugin/pricing_management.go +++ b/internal/plugin/pricing_management.go @@ -71,7 +71,7 @@ func (a *App) putPrice(body []byte) ManagementResponse { 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 + a.prices[normalizeModelName(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()}}) } @@ -95,7 +95,7 @@ func (a *App) deletePrice(body []byte) ManagementResponse { 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) + delete(a.prices, normalizeModelName(request.Model)) return jsonManagementResponse(http.StatusOK, map[string]any{"deleted": request.Model}) } @@ -195,13 +195,18 @@ func gcd(a, b int64) int64 { } func formatMicros(value int64) string { + prefix := "" + if value < 0 { + prefix = "-" + value = -value + } 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 prefix + strconv.FormatInt(whole, 10) } - return strconv.FormatInt(whole, 10) + "." + fraction + return prefix + strconv.FormatInt(whole, 10) + "." + fraction } func formatRatio(value pricing.Ratio) string { diff --git a/internal/repository/sqlite_access.go b/internal/repository/sqlite_access.go index a5d7b7d..a92203c 100644 --- a/internal/repository/sqlite_access.go +++ b/internal/repository/sqlite_access.go @@ -10,7 +10,6 @@ import ( "time" managedaccess "cpa-ext/internal/access" - "cpa-ext/internal/collection" ) var ErrManagedKeyNotFound = errors.New("managed key not found") @@ -62,6 +61,16 @@ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`, key.ID, key.Name, key.Secret, key.Status, k return err } if _, err = tx.ExecContext(ctx, ` +INSERT INTO billing_accounts (managed_key_id, quota_micros, reset_period, max_concurrency, current_cycle_sequence, updated_at) +VALUES (?, 0, 'none', 4, 1, ?)`, key.ID, formatTime(key.CreatedAt)); err != nil { + return fmt.Errorf("创建 Key 额度账户: %w", err) + } + if _, err = tx.ExecContext(ctx, ` +INSERT INTO billing_cycles (managed_key_id, sequence, started_at, quota_micros, spent_micros) +VALUES (?, 1, ?, 0, 0)`, key.ID, formatTime(key.CreatedAt)); err != nil { + return fmt.Errorf("创建 Key 额度周期: %w", err) + } + if _, err = tx.ExecContext(ctx, ` UPDATE usage_records SET managed_key_id = ?, key_alias = ? WHERE managed_key_id = '' AND api_key = ?`, key.ID, key.Name, key.Secret); err != nil { return fmt.Errorf("关联历史用量: %w", err) @@ -326,74 +335,33 @@ SELECT id, cpa_auth_id, cpa_auth_index, provider, display_name, status, status_m func (r *SQLiteUsageRepository) KeyStats(ctx context.Context, keyID string, today time.Time) (managedaccess.KeyStats, error) { stats := managedaccess.KeyStats{KeyID: keyID} - facts, err := r.keyUsageFacts(ctx, keyID) + var last sql.NullString + err := r.db.QueryRowContext(ctx, ` +SELECT COUNT(*), + COALESCE(SUM(u.input_tokens),0), COALESCE(SUM(u.output_tokens),0), + COALESCE(SUM(u.total_tokens),0), COALESCE(SUM(u.cost_micros),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN 1 ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.input_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.output_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.total_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.cost_micros ELSE 0 END),0), + MAX(d.requested_at) +FROM request_detail_index d +LEFT JOIN usage_records u ON u.id=d.usage_id +WHERE d.managed_key_id=?`, formatTime(today), formatTime(today), formatTime(today), formatTime(today), formatTime(today), keyID).Scan( + &stats.Total.Requests, &stats.Total.InputTokens, &stats.Total.OutputTokens, &stats.Total.TotalTokens, &stats.Total.CostMicros, + &stats.Today.Requests, &stats.Today.InputTokens, &stats.Today.OutputTokens, &stats.Today.TotalTokens, &stats.Today.CostMicros, &last) if err != nil { - return stats, err + return stats, fmt.Errorf("汇总 Key 用量: %w", err) } - for _, fact := range mergeOrphanLifecycleRecords(facts) { - addUsageSummary(&stats.Total, fact) - if !fact.RequestedAt.Before(today) { - addUsageSummary(&stats.Today, fact) - } - if stats.LastUsedAt == nil || fact.RequestedAt.After(*stats.LastUsedAt) { - value := fact.RequestedAt - stats.LastUsedAt = &value + if last.Valid { + if parsed, parseErr := time.Parse(time.RFC3339Nano, last.String); parseErr == nil { + stats.LastUsedAt = &parsed } } return stats, nil } -func (r *SQLiteUsageRepository) keyUsageFacts(ctx context.Context, keyID string) ([]collection.Record, error) { - rows, err := r.db.QueryContext(ctx, ` -SELECT request_id, requested_at, model, api_key, executor_type, - input_tokens, output_tokens, total_tokens, cost_micros, '' -FROM usage_records WHERE managed_key_id=? -UNION ALL -SELECT request_id, requested_at, COALESCE(NULLIF(model, ''), requested_model), '', '', - 0, 0, 0, NULL, outcome -FROM request_records r -WHERE managed_key_id=? AND NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id=r.request_id) -ORDER BY requested_at DESC`, keyID, keyID) - if err != nil { - return nil, fmt.Errorf("查询 Key 用量事实: %w", err) - } - defer rows.Close() - facts := make([]collection.Record, 0) - for rows.Next() { - var fact collection.Record - var requestedAt string - var cost sql.NullInt64 - if err := rows.Scan(&fact.RequestID, &requestedAt, &fact.Model, &fact.APIKey, &fact.ExecutorType, - &fact.InputTokens, &fact.OutputTokens, &fact.TotalTokens, &cost, &fact.Outcome); err != nil { - return nil, fmt.Errorf("读取 Key 用量事实: %w", err) - } - fact.ManagedKeyID = keyID - fact.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt) - if err != nil { - return nil, fmt.Errorf("解析 Key 用量时间: %w", err) - } - if cost.Valid { - value := cost.Int64 - fact.CostMicros = &value - } - facts = append(facts, fact) - } - if err := rows.Err(); err != nil { - return nil, fmt.Errorf("遍历 Key 用量事实: %w", err) - } - return facts, nil -} - -func addUsageSummary(summary *managedaccess.UsageSummary, fact collection.Record) { - summary.Requests++ - summary.InputTokens += fact.InputTokens - summary.OutputTokens += fact.OutputTokens - summary.TotalTokens += fact.TotalTokens - if fact.CostMicros != nil { - summary.CostMicros += *fact.CostMicros - } -} - func (r *SQLiteUsageRepository) ModelSuggestions(ctx context.Context) ([]string, error) { rows, err := r.db.QueryContext(ctx, ` SELECT model FROM model_prices diff --git a/internal/repository/sqlite_billing.go b/internal/repository/sqlite_billing.go new file mode 100644 index 0000000..b6b586a --- /dev/null +++ b/internal/repository/sqlite_billing.go @@ -0,0 +1,400 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strings" + "time" + + managedaccess "cpa-ext/internal/access" +) + +var ( + ErrBillingQuotaExhausted = errors.New("billing quota exhausted") + ErrBillingConcurrency = errors.New("billing concurrency exceeded") +) + +func (r *SQLiteUsageRepository) BillingState(ctx context.Context, keyID string, now time.Time) (managedaccess.BillingState, error) { + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return managedaccess.BillingState{}, err + } + defer func() { _ = tx.Rollback() }() + if err := ensureBillingReset(ctx, tx, keyID, now); err != nil { + return managedaccess.BillingState{}, err + } + state, err := scanBillingState(ctx, tx, keyID) + if err != nil { + return state, err + } + if err := tx.Commit(); err != nil { + return state, err + } + return state, nil +} + +func (r *SQLiteUsageRepository) UpdateBilling(ctx context.Context, keyID string, settings managedaccess.BillingSettings, now time.Time) (managedaccess.BillingState, error) { + if settings.QuotaMicros < 0 { + return managedaccess.BillingState{}, errors.New("额度不能小于 0") + } + if settings.MaxConcurrency < 1 || settings.MaxConcurrency > 64 { + return managedaccess.BillingState{}, errors.New("并发上限必须为 1-64") + } + if !validResetPeriod(settings.ResetPeriod) { + return managedaccess.BillingState{}, errors.New("重置周期只能是 none、daily、weekly 或 monthly") + } + if now.IsZero() { + now = time.Now().UTC() + } + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return managedaccess.BillingState{}, err + } + defer func() { _ = tx.Rollback() }() + if err := ensureBillingReset(ctx, tx, keyID, now); err != nil { + return managedaccess.BillingState{}, err + } + current, err := scanBillingState(ctx, tx, keyID) + if err != nil { + return current, err + } + nextReset := settings.NextResetAt + anchorDay := 0 + if settings.ResetPeriod == managedaccess.ResetNone { + nextReset = nil + } else { + if nextReset == nil || nextReset.IsZero() { + var createdRaw string + if err := tx.QueryRowContext(ctx, `SELECT created_at FROM managed_keys WHERE id=?`, keyID).Scan(&createdRaw); err != nil { + return current, err + } + created, _ := time.Parse(time.RFC3339Nano, createdRaw) + value := nextResetAfter(created, settings.ResetPeriod, now) + nextReset = &value + } + if !nextReset.After(now) { + return current, errors.New("下次重置时间必须晚于当前时间") + } + anchorDay = nextReset.In(shanghaiLocation()).Day() + } + nextValue := any(nil) + if nextReset != nil { + nextValue = formatTime(*nextReset) + } + if _, err := tx.ExecContext(ctx, ` +UPDATE billing_accounts SET quota_micros=?, reset_period=?, next_reset_at=?, reset_anchor_day=?, + max_concurrency=?, updated_at=? WHERE managed_key_id=?`, settings.QuotaMicros, settings.ResetPeriod, + nextValue, anchorDay, settings.MaxConcurrency, formatTime(now), keyID); err != nil { + return current, fmt.Errorf("更新额度设置: %w", err) + } + if _, err := tx.ExecContext(ctx, ` +UPDATE billing_cycles SET quota_micros=? WHERE managed_key_id=? AND sequence=?`, + settings.QuotaMicros, keyID, current.CycleSequence); err != nil { + return current, fmt.Errorf("更新当前额度周期: %w", err) + } + if settings.QuotaMicros != current.QuotaMicros { + balance := settings.QuotaMicros - current.SpentMicros + _, err = tx.ExecContext(ctx, ` +INSERT INTO billing_ledger (event_key, managed_key_id, cycle_sequence, kind, amount_micros, + balance_after_micros, occurred_at, created_at) VALUES (?, ?, ?, 'quota_change', ?, ?, ?, ?)`, + fmt.Sprintf("quota:%s:%d", keyID, now.UnixNano()), keyID, current.CycleSequence, + settings.QuotaMicros-current.QuotaMicros, balance, formatTime(now), formatTime(now)) + if err != nil { + return current, fmt.Errorf("记录额度调整: %w", err) + } + } + state, err := scanBillingState(ctx, tx, keyID) + if err != nil { + return state, err + } + if err := tx.Commit(); err != nil { + return state, err + } + return state, nil +} + +func (r *SQLiteUsageRepository) ResetBilling(ctx context.Context, keyID string, now time.Time) (managedaccess.BillingState, error) { + if now.IsZero() { + now = time.Now().UTC() + } + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return managedaccess.BillingState{}, err + } + defer func() { _ = tx.Rollback() }() + var period string + var anchorDay, sequence int64 + var quota int64 + if err := tx.QueryRowContext(ctx, ` +SELECT reset_period, reset_anchor_day, current_cycle_sequence, quota_micros +FROM billing_accounts WHERE managed_key_id=?`, keyID).Scan(&period, &anchorDay, &sequence, "a); err != nil { + return managedaccess.BillingState{}, err + } + var next *time.Time + if period != managedaccess.ResetNone { + value := advanceReset(now, period, int(now.In(shanghaiLocation()).Day())) + next = &value + anchorDay = int64(now.In(shanghaiLocation()).Day()) + } + if err := resetBillingCycle(ctx, tx, keyID, sequence, quota, now, next, int(anchorDay)); err != nil { + return managedaccess.BillingState{}, err + } + state, err := scanBillingState(ctx, tx, keyID) + if err != nil { + return state, err + } + if err := tx.Commit(); err != nil { + return state, err + } + return state, nil +} + +func (r *SQLiteUsageRepository) AuthorizeBilling(ctx context.Context, keyID, requestID string, now time.Time) (managedaccess.BillingState, error) { + requestID = strings.TrimSpace(requestID) + if requestID == "" { + return managedaccess.BillingState{}, errors.New("request_id 不能为空") + } + if now.IsZero() { + now = time.Now().UTC() + } + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return managedaccess.BillingState{}, err + } + defer func() { _ = tx.Rollback() }() + if err := ensureBillingReset(ctx, tx, keyID, now); err != nil { + return managedaccess.BillingState{}, err + } + state, err := scanBillingState(ctx, tx, keyID) + if err != nil { + return state, err + } + var existingStatus string + err = tx.QueryRowContext(ctx, `SELECT status FROM billing_admissions WHERE request_id=?`, requestID).Scan(&existingStatus) + if err == nil { + if err := tx.Commit(); err != nil { + return state, err + } + return state, nil + } + if !errors.Is(err, sql.ErrNoRows) { + return state, err + } + if state.BalanceMicros <= 0 { + return state, ErrBillingQuotaExhausted + } + if state.ActiveRequests >= state.MaxConcurrency { + return state, ErrBillingConcurrency + } + if _, err := tx.ExecContext(ctx, ` +INSERT INTO billing_admissions (request_id, managed_key_id, cycle_sequence, opened_at, status) +VALUES (?, ?, ?, ?, 'open')`, requestID, keyID, state.CycleSequence, formatTime(now)); err != nil { + return state, fmt.Errorf("占用并发额度: %w", err) + } + state.ActiveRequests++ + if err := tx.Commit(); err != nil { + return state, err + } + return state, nil +} + +func (r *SQLiteUsageRepository) CompleteBillingAdmission(ctx context.Context, requestID string, now time.Time) error { + if strings.TrimSpace(requestID) == "" { + return nil + } + if now.IsZero() { + now = time.Now().UTC() + } + _, err := r.db.ExecContext(ctx, ` +UPDATE billing_admissions SET status='closed', closed_at=? WHERE request_id=? AND status='open'`, + formatTime(now), requestID) + return err +} + +func (r *SQLiteUsageRepository) ReleaseStaleAdmissions(ctx context.Context, now time.Time) error { + if now.IsZero() { + now = time.Now().UTC() + } + _, err := r.db.ExecContext(ctx, ` +UPDATE billing_admissions SET status='abandoned', closed_at=? WHERE status='open'`, formatTime(now)) + return err +} + +func (r *SQLiteUsageRepository) ListBillingLedger(ctx context.Context, keyID string, beforeID int64, limit int) ([]managedaccess.LedgerEntry, error) { + if limit < 1 || limit > 100 { + limit = 50 + } + query := ` +SELECT id, managed_key_id, cycle_sequence, kind, amount_micros, balance_after_micros, + request_id, execution_id, model, occurred_at, created_at +FROM billing_ledger WHERE managed_key_id=?` + args := []any{keyID} + if beforeID > 0 { + query += ` AND id < ?` + args = append(args, beforeID) + } + query += ` ORDER BY id DESC LIMIT ?` + args = append(args, limit) + rows, err := r.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + entries := make([]managedaccess.LedgerEntry, 0, limit) + for rows.Next() { + var entry managedaccess.LedgerEntry + var occurredRaw, createdRaw string + if err := rows.Scan(&entry.ID, &entry.KeyID, &entry.CycleSequence, &entry.Kind, + &entry.AmountMicros, &entry.BalanceAfterMicros, &entry.RequestID, &entry.ExecutionID, + &entry.Model, &occurredRaw, &createdRaw); err != nil { + return nil, err + } + entry.OccurredAt, _ = time.Parse(time.RFC3339Nano, occurredRaw) + entry.CreatedAt, _ = time.Parse(time.RFC3339Nano, createdRaw) + entries = append(entries, entry) + } + return entries, rows.Err() +} + +func scanBillingState(ctx context.Context, tx *sql.Tx, keyID string) (managedaccess.BillingState, error) { + var state managedaccess.BillingState + var nextRaw sql.NullString + var startedRaw string + err := tx.QueryRowContext(ctx, ` +SELECT a.managed_key_id, a.quota_micros, c.spent_micros, a.reset_period, a.next_reset_at, + a.max_concurrency, a.current_cycle_sequence, c.started_at, + (SELECT COUNT(*) FROM billing_admissions d WHERE d.managed_key_id=a.managed_key_id AND d.status='open') +FROM billing_accounts a JOIN billing_cycles c + ON c.managed_key_id=a.managed_key_id AND c.sequence=a.current_cycle_sequence +WHERE a.managed_key_id=?`, keyID).Scan(&state.KeyID, &state.QuotaMicros, &state.SpentMicros, + &state.ResetPeriod, &nextRaw, &state.MaxConcurrency, &state.CycleSequence, + &startedRaw, &state.ActiveRequests) + if err != nil { + return state, err + } + state.BalanceMicros = state.QuotaMicros - state.SpentMicros + state.CycleStartedAt, _ = time.Parse(time.RFC3339Nano, startedRaw) + if nextRaw.Valid && nextRaw.String != "" { + value, err := time.Parse(time.RFC3339Nano, nextRaw.String) + if err == nil { + state.NextResetAt = &value + } + } + return state, nil +} + +func ensureBillingReset(ctx context.Context, tx *sql.Tx, keyID string, now time.Time) error { + if now.IsZero() { + now = time.Now().UTC() + } + var period string + var nextRaw sql.NullString + var anchorDay int + var sequence, quota int64 + if err := tx.QueryRowContext(ctx, ` +SELECT reset_period, next_reset_at, reset_anchor_day, current_cycle_sequence, quota_micros +FROM billing_accounts WHERE managed_key_id=?`, keyID).Scan(&period, &nextRaw, &anchorDay, &sequence, "a); err != nil { + return err + } + if period == managedaccess.ResetNone || !nextRaw.Valid || nextRaw.String == "" { + return nil + } + next, err := time.Parse(time.RFC3339Nano, nextRaw.String) + if err != nil || now.Before(next) { + return err + } + boundary := next + future := advanceReset(next, period, anchorDay) + for !future.After(now) { + boundary = future + future = advanceReset(future, period, anchorDay) + } + return resetBillingCycle(ctx, tx, keyID, sequence, quota, boundary, &future, anchorDay) +} + +func resetBillingCycle(ctx context.Context, tx *sql.Tx, keyID string, sequence, quota int64, boundary time.Time, next *time.Time, anchorDay int) error { + if _, err := tx.ExecContext(ctx, ` +UPDATE billing_cycles SET ended_at=? WHERE managed_key_id=? AND sequence=?`, formatTime(boundary), keyID, sequence); err != nil { + return err + } + newSequence := sequence + 1 + if _, err := tx.ExecContext(ctx, ` +INSERT INTO billing_cycles (managed_key_id, sequence, started_at, quota_micros, spent_micros) +VALUES (?, ?, ?, ?, 0)`, keyID, newSequence, formatTime(boundary), quota); err != nil { + return err + } + nextValue := any(nil) + if next != nil { + nextValue = formatTime(*next) + } + if _, err := tx.ExecContext(ctx, ` +UPDATE billing_accounts SET current_cycle_sequence=?, next_reset_at=?, reset_anchor_day=?, updated_at=? +WHERE managed_key_id=?`, newSequence, nextValue, anchorDay, formatTime(boundary), keyID); err != nil { + return err + } + _, err := tx.ExecContext(ctx, ` +INSERT OR IGNORE INTO billing_ledger (event_key, managed_key_id, cycle_sequence, kind, amount_micros, + balance_after_micros, occurred_at, created_at) VALUES (?, ?, ?, 'cycle_reset', ?, ?, ?, ?)`, + fmt.Sprintf("reset:%s:%d", keyID, newSequence), keyID, newSequence, quota, quota, + formatTime(boundary), formatTime(time.Now().UTC())) + return err +} + +func validResetPeriod(period string) bool { + switch period { + case managedaccess.ResetNone, managedaccess.ResetDaily, managedaccess.ResetWeekly, managedaccess.ResetMonthly: + return true + default: + return false + } +} + +func shanghaiLocation() *time.Location { + location, err := time.LoadLocation("Asia/Shanghai") + if err != nil { + return time.FixedZone("Asia/Shanghai", 8*60*60) + } + return location +} + +func nextResetAfter(anchor time.Time, period string, now time.Time) time.Time { + if anchor.IsZero() { + anchor = now + } + anchorDay := anchor.In(shanghaiLocation()).Day() + next := advanceReset(anchor, period, anchorDay) + for !next.After(now) { + next = advanceReset(next, period, anchorDay) + } + return next +} + +func advanceReset(value time.Time, period string, anchorDay int) time.Time { + location := shanghaiLocation() + local := value.In(location) + switch period { + case managedaccess.ResetDaily: + return local.AddDate(0, 0, 1).UTC() + case managedaccess.ResetWeekly: + return local.AddDate(0, 0, 7).UTC() + case managedaccess.ResetMonthly: + year, month := local.Year(), local.Month()+1 + if month > 12 { + year++ + month = 1 + } + lastDay := time.Date(year, month+1, 0, local.Hour(), local.Minute(), local.Second(), local.Nanosecond(), location).Day() + day := anchorDay + if day < 1 { + day = local.Day() + } + if day > lastDay { + day = lastDay + } + return time.Date(year, month, day, local.Hour(), local.Minute(), local.Second(), local.Nanosecond(), location).UTC() + default: + return time.Time{} + } +} diff --git a/internal/repository/sqlite_billing_test.go b/internal/repository/sqlite_billing_test.go new file mode 100644 index 0000000..e24c468 --- /dev/null +++ b/internal/repository/sqlite_billing_test.go @@ -0,0 +1,97 @@ +package repository + +import ( + "context" + "errors" + "path/filepath" + "testing" + "time" + + managedaccess "cpa-ext/internal/access" + "cpa-ext/internal/collection" +) + +func TestBillingQuotaConcurrencySettlementAndReset(t *testing.T) { + ctx := context.Background() + store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "billing.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + if _, err := store.BootstrapManagedKey(ctx, "default", "000000"); err != nil { + t.Fatal(err) + } + now := time.Now().UTC().Add(time.Second) + state, err := store.BillingState(ctx, "key_default", now) + if err != nil || state.QuotaMicros != 0 || state.MaxConcurrency != 4 { + t.Fatalf("initial billing state=%+v err=%v", state, err) + } + if _, err := store.AuthorizeBilling(ctx, "key_default", "zero", now); !errors.Is(err, ErrBillingQuotaExhausted) { + t.Fatalf("zero quota authorization err=%v", err) + } + state, err = store.UpdateBilling(ctx, "key_default", managedaccess.BillingSettings{ + QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 2, + }, now) + if err != nil || state.BalanceMicros != 1_000_000 { + t.Fatalf("updated billing state=%+v err=%v", state, err) + } + for _, requestID := range []string{"one", "two"} { + if _, err := store.AuthorizeBilling(ctx, "key_default", requestID, now); err != nil { + t.Fatalf("authorize %s: %v", requestID, err) + } + } + if _, err := store.AuthorizeBilling(ctx, "key_default", "three", now); !errors.Is(err, ErrBillingConcurrency) { + t.Fatalf("third concurrent authorization err=%v", err) + } + if err := store.CompleteBillingAdmission(ctx, "one", now); err != nil { + t.Fatal(err) + } + if _, err := store.AuthorizeBilling(ctx, "key_default", "three", now); err != nil { + t.Fatalf("released slot not reusable: %v", err) + } + + cost := int64(600_000) + record := collection.Record{ManagedKeyID: "key_default", ExecutionID: "exec-1", RequestID: "one", + RequestedAt: now, Model: "deepseek-v4-flash", CostMicros: &cost} + if err := store.Insert(ctx, record); err != nil { + t.Fatal(err) + } + if err := store.Insert(ctx, record); err != nil { + t.Fatal(err) + } + state, _ = store.BillingState(ctx, "key_default", now) + if state.SpentMicros != 600_000 || state.BalanceMicros != 400_000 { + t.Fatalf("duplicate settlement changed balance: %+v", state) + } + cost = 500_000 + record.ExecutionID = "exec-2" + record.CostMicros = &cost + if err := store.Insert(ctx, record); err != nil { + t.Fatal(err) + } + state, _ = store.BillingState(ctx, "key_default", now) + if state.BalanceMicros != -100_000 { + t.Fatalf("soft quota did not preserve negative balance: %+v", state) + } + if _, err := store.AuthorizeBilling(ctx, "key_default", "blocked", now); !errors.Is(err, ErrBillingQuotaExhausted) { + t.Fatalf("negative balance authorization err=%v", err) + } + state, err = store.ResetBilling(ctx, "key_default", now.Add(time.Minute)) + if err != nil || state.SpentMicros != 0 || state.BalanceMicros != 1_000_000 || state.CycleSequence != 2 { + t.Fatalf("manual reset state=%+v err=%v", state, err) + } + entries, err := store.ListBillingLedger(ctx, "key_default", 0, 50) + if err != nil || len(entries) != 4 { + t.Fatalf("ledger=%+v err=%v", entries, err) + } +} + +func TestBillingAutomaticMonthlyResetPreservesAnchorDay(t *testing.T) { + location := shanghaiLocation() + anchor := time.Date(2026, time.January, 31, 10, 0, 0, 0, location) + february := advanceReset(anchor, managedaccess.ResetMonthly, 31).In(location) + march := advanceReset(february, managedaccess.ResetMonthly, 31).In(location) + if february.Day() != 28 || february.Month() != time.February || march.Day() != 31 || march.Month() != time.March { + t.Fatalf("monthly sequence: %s then %s", february, march) + } +} diff --git a/internal/repository/sqlite_projection.go b/internal/repository/sqlite_projection.go new file mode 100644 index 0000000..cfc4b57 --- /dev/null +++ b/internal/repository/sqlite_projection.go @@ -0,0 +1,350 @@ +package repository + +import ( + "context" + "database/sql" + "errors" + "fmt" + "strings" + "time" + + "cpa-ext/internal/collection" +) + +const requestDetailProjectionMigration = "request-detail-index-v1" + +func (r *SQLiteUsageRepository) ensureRequestDetailProjection(ctx context.Context) error { + var completed string + err := r.db.QueryRowContext(ctx, `SELECT completed_at FROM cpa_ext_migrations WHERE name=?`, requestDetailProjectionMigration).Scan(&completed) + if err == nil { + return nil + } + if !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("检查请求明细索引迁移: %w", err) + } + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("开始请求明细索引迁移: %w", err) + } + defer func() { _ = tx.Rollback() }() + if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index`); err != nil { + return fmt.Errorf("清理请求明细索引: %w", err) + } + if _, err := tx.ExecContext(ctx, usageProjectionInsertSQL+` +SELECT u.id, CASE WHEN r.request_id IS NULL THEN '' ELSE r.request_id END, + u.requested_at, COALESCE(NULLIF(u.managed_key_id, ''), r.managed_key_id, ''), + u.model, + CASE WHEN u.failed=1 THEN 'failed' + WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome + ELSE 'succeeded' END, + u.auth_id, `+endpointKindSQL("COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '')")+`, + COALESCE(NULLIF(u.request_id, ''), r.request_id, '') +FROM usage_records u +LEFT JOIN request_records r ON u.request_id <> '' AND r.request_id=u.request_id`); err != nil { + return fmt.Errorf("回填用量请求明细索引: %w", err) + } + if _, err := tx.ExecContext(ctx, lifecycleProjectionInsertSQL+` +SELECT NULL, r.request_id, r.requested_at, r.managed_key_id, + COALESCE(NULLIF(r.model, ''), r.requested_model), + CASE WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome ELSE 'succeeded' END, + '', `+endpointKindSQL("r.endpoint")+`, r.request_id +FROM request_records r +WHERE NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id=r.request_id)`); err != nil { + return fmt.Errorf("回填终态请求明细索引: %w", err) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("提交请求明细索引迁移: %w", err) + } + + rows, err := r.db.QueryContext(ctx, `SELECT requested_at FROM usage_records WHERE request_id='' ORDER BY requested_at`) + if err != nil { + return fmt.Errorf("扫描孤立用量时间: %w", err) + } + var orphanTimes []time.Time + for rows.Next() { + var raw string + if err := rows.Scan(&raw); err != nil { + _ = rows.Close() + return fmt.Errorf("读取孤立用量时间: %w", err) + } + parsed, parseErr := time.Parse(time.RFC3339Nano, raw) + if parseErr != nil { + _ = rows.Close() + return fmt.Errorf("解析孤立用量时间: %w", parseErr) + } + orphanTimes = append(orphanTimes, parsed) + } + if err := rows.Close(); err != nil { + return fmt.Errorf("关闭孤立用量扫描: %w", err) + } + var last time.Time + for _, candidate := range orphanTimes { + if !last.IsZero() && candidate.Sub(last) < orphanLifecycleMatchWindow { + continue + } + if err := r.reconcileOrphanProjection(ctx, candidate); err != nil { + return err + } + last = candidate + } + if _, err := r.db.ExecContext(ctx, `INSERT OR REPLACE INTO cpa_ext_migrations(name, completed_at) VALUES(?, ?)`, requestDetailProjectionMigration, formatTime(time.Now().UTC())); err != nil { + return fmt.Errorf("完成请求明细索引迁移: %w", err) + } + return nil +} + +const usageProjectionInsertSQL = ` +INSERT OR IGNORE INTO request_detail_index + (usage_id, lifecycle_request_id, requested_at, managed_key_id, model, result, auth_id, endpoint_kind, request_id) ` + +const lifecycleProjectionInsertSQL = ` +INSERT OR IGNORE INTO request_detail_index + (usage_id, lifecycle_request_id, requested_at, managed_key_id, model, result, auth_id, endpoint_kind, request_id) ` + +func endpointKindSQL(expression string) string { + return `CASE + WHEN LOWER(TRIM(` + expression + `)) LIKE '%/responses/compact' THEN 'compact' + WHEN LOWER(TRIM(` + expression + `)) LIKE '%/chat/completions' THEN 'chat' + WHEN LOWER(TRIM(` + expression + `)) LIKE '%/responses' THEN 'responses' + ELSE LOWER(TRIM(` + expression + `)) END` +} + +func endpointKind(value string) string { + path := strings.ToLower(strings.TrimSpace(value)) + if fields := strings.Fields(path); len(fields) > 1 { + path = fields[len(fields)-1] + } + if index := strings.IndexByte(path, '?'); index >= 0 { + path = path[:index] + } + path = strings.TrimRight(path, "/") + switch { + case strings.HasSuffix(path, "/responses/compact"): + return "compact" + case strings.HasSuffix(path, "/chat/completions"): + return "chat" + case strings.HasSuffix(path, "/responses"): + return "responses" + default: + return path + } +} + +func usageResult(failed bool, outcome string) string { + switch strings.ToLower(strings.TrimSpace(outcome)) { + case collection.UsageResultCanceled: + return collection.UsageResultCanceled + case collection.UsageResultRejected: + return collection.UsageResultRejected + case collection.UsageResultFailed: + return collection.UsageResultFailed + } + if failed { + return collection.UsageResultFailed + } + return collection.UsageResultSucceeded +} + +func insertUsageProjection(ctx context.Context, tx *sql.Tx, usageID int64) error { + _, err := tx.ExecContext(ctx, usageProjectionInsertSQL+` +SELECT u.id, CASE WHEN r.request_id IS NULL THEN '' ELSE r.request_id END, + u.requested_at, COALESCE(NULLIF(u.managed_key_id, ''), r.managed_key_id, ''), + u.model, + CASE WHEN u.failed=1 THEN 'failed' + WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome + ELSE 'succeeded' END, + u.auth_id, `+endpointKindSQL("COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '')")+`, + COALESCE(NULLIF(u.request_id, ''), r.request_id, '') +FROM usage_records u +LEFT JOIN request_records r ON u.request_id <> '' AND r.request_id=u.request_id +WHERE u.id=?`, usageID) + if err != nil { + return fmt.Errorf("写入请求明细索引: %w", err) + } + return nil +} + +func syncRequestProjection(ctx context.Context, tx *sql.Tx, requestID string) (bool, error) { + result, err := tx.ExecContext(ctx, ` +UPDATE request_detail_index +SET lifecycle_request_id=?, request_id=?, + managed_key_id=COALESCE(NULLIF((SELECT managed_key_id FROM usage_records WHERE id=request_detail_index.usage_id), ''), + (SELECT managed_key_id FROM request_records WHERE request_id=?), ''), + result=CASE + WHEN (SELECT failed FROM usage_records WHERE id=request_detail_index.usage_id)=1 THEN 'failed' + WHEN (SELECT outcome FROM request_records WHERE request_id=?) IN ('failed','rejected','canceled') + THEN (SELECT outcome FROM request_records WHERE request_id=?) + ELSE 'succeeded' END +WHERE usage_id IN (SELECT id FROM usage_records WHERE request_id=?)`, requestID, requestID, requestID, requestID, requestID, requestID) + if err != nil { + return false, fmt.Errorf("关联请求终态索引: %w", err) + } + linked, _ := result.RowsAffected() + if linked > 0 { + if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET endpoint_kind=`+endpointKindSQL(`COALESCE(NULLIF((SELECT endpoint FROM usage_records WHERE id=request_detail_index.usage_id), ''), (SELECT endpoint FROM request_records WHERE request_id=?), '')`)+` WHERE usage_id IN (SELECT id FROM usage_records WHERE request_id=?)`, requestID, requestID, requestID, requestID, requestID); err != nil { + return false, fmt.Errorf("更新请求端点索引: %w", err) + } + if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index WHERE usage_id IS NULL AND lifecycle_request_id=?`, requestID); err != nil { + return false, fmt.Errorf("删除重复请求终态索引: %w", err) + } + return true, nil + } + _, err = tx.ExecContext(ctx, lifecycleProjectionInsertSQL+` +SELECT NULL, r.request_id, r.requested_at, r.managed_key_id, + COALESCE(NULLIF(r.model, ''), r.requested_model), + CASE WHEN r.outcome IN ('failed','rejected','canceled') THEN r.outcome ELSE 'succeeded' END, + '', `+endpointKindSQL("r.endpoint")+`, r.request_id +FROM request_records r WHERE r.request_id=?`, requestID) + if err != nil { + return false, fmt.Errorf("写入终态请求明细索引: %w", err) + } + return false, nil +} + +type orphanUsageProjection struct { + ID int64 + ManagedKeyID string + RequestedAt time.Time + Model string + Failed bool + AuthID string + Endpoint string +} + +type orphanLifecycleProjection struct { + RequestID string + ManagedKeyID string + RequestedAt time.Time + Model string + Outcome string + Endpoint string +} + +func (r *SQLiteUsageRepository) reconcileOrphanProjection(ctx context.Context, at time.Time) error { + window := time.Second + from, to := formatTime(at.Add(-window)), formatTime(at.Add(window)) + usageRows, err := r.db.QueryContext(ctx, ` +SELECT id, managed_key_id, requested_at, model, failed, auth_id, endpoint +FROM usage_records WHERE request_id='' AND requested_at BETWEEN ? AND ?`, from, to) + if err != nil { + return fmt.Errorf("查询孤立用量候选: %w", err) + } + var usages []orphanUsageProjection + for usageRows.Next() { + var item orphanUsageProjection + var requestedAt string + if err := usageRows.Scan(&item.ID, &item.ManagedKeyID, &requestedAt, &item.Model, &item.Failed, &item.AuthID, &item.Endpoint); err != nil { + _ = usageRows.Close() + return fmt.Errorf("读取孤立用量候选: %w", err) + } + item.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt) + if err != nil { + _ = usageRows.Close() + return fmt.Errorf("解析孤立用量候选时间: %w", err) + } + usages = append(usages, item) + } + _ = usageRows.Close() + + lifecycleRows, err := r.db.QueryContext(ctx, ` +SELECT r.request_id, r.managed_key_id, r.requested_at, + COALESCE(NULLIF(r.model, ''), r.requested_model), r.outcome, r.endpoint +FROM request_records r +WHERE r.requested_at BETWEEN ? AND ? + AND NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id=r.request_id)`, from, to) + if err != nil { + return fmt.Errorf("查询孤立终态候选: %w", err) + } + var lifecycles []orphanLifecycleProjection + for lifecycleRows.Next() { + var item orphanLifecycleProjection + var requestedAt string + if err := lifecycleRows.Scan(&item.RequestID, &item.ManagedKeyID, &requestedAt, &item.Model, &item.Outcome, &item.Endpoint); err != nil { + _ = lifecycleRows.Close() + return fmt.Errorf("读取孤立终态候选: %w", err) + } + item.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt) + if err != nil { + _ = lifecycleRows.Close() + return fmt.Errorf("解析孤立终态候选时间: %w", err) + } + lifecycles = append(lifecycles, item) + } + _ = lifecycleRows.Close() + + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("开始重建孤立请求索引: %w", err) + } + defer func() { _ = tx.Rollback() }() + for _, usage := range usages { + if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET lifecycle_request_id='', request_id='', managed_key_id=?, model=?, result=?, auth_id=?, endpoint_kind=? WHERE usage_id=?`, usage.ManagedKeyID, usage.Model, usageResult(usage.Failed, ""), usage.AuthID, endpointKind(usage.Endpoint), usage.ID); err != nil { + return fmt.Errorf("重置孤立用量索引: %w", err) + } + } + for _, lifecycle := range lifecycles { + if _, err := tx.ExecContext(ctx, lifecycleProjectionInsertSQL+`VALUES(NULL, ?, ?, ?, ?, ?, '', ?, ?)`, lifecycle.RequestID, formatTime(lifecycle.RequestedAt), lifecycle.ManagedKeyID, lifecycle.Model, usageResult(false, lifecycle.Outcome), endpointKind(lifecycle.Endpoint), lifecycle.RequestID); err != nil { + return fmt.Errorf("恢复孤立终态索引: %w", err) + } + } + + usageMatches := make(map[int]lifecycleMatch) + requestMatches := make(map[int]lifecycleMatch) + for usageIndex, usage := range usages { + for requestIndex, lifecycle := range lifecycles { + if !sameOrphanProjection(usage, lifecycle) { + continue + } + distance := usage.RequestedAt.Sub(lifecycle.RequestedAt) + if distance < 0 { + distance = -distance + } + if distance > orphanLifecycleMatchWindow { + continue + } + if current, ok := usageMatches[usageIndex]; ok { + usageMatches[usageIndex] = betterLifecycleMatch(current, requestIndex, distance) + } else { + usageMatches[usageIndex] = lifecycleMatch{index: requestIndex, distance: distance} + } + if current, ok := requestMatches[requestIndex]; ok { + requestMatches[requestIndex] = betterLifecycleMatch(current, usageIndex, distance) + } else { + requestMatches[requestIndex] = lifecycleMatch{index: usageIndex, distance: distance} + } + } + } + for usageIndex, requestMatch := range usageMatches { + usageMatch, ok := requestMatches[requestMatch.index] + if requestMatch.ambiguous || !ok || usageMatch.ambiguous || usageMatch.index != usageIndex { + continue + } + usage, lifecycle := usages[usageIndex], lifecycles[requestMatch.index] + managedKeyID := usage.ManagedKeyID + if managedKeyID == "" { + managedKeyID = lifecycle.ManagedKeyID + } + endpoint := usage.Endpoint + if strings.TrimSpace(endpoint) == "" { + endpoint = lifecycle.Endpoint + } + if _, err := tx.ExecContext(ctx, `UPDATE request_detail_index SET lifecycle_request_id=?, request_id=?, managed_key_id=?, result=?, endpoint_kind=? WHERE usage_id=?`, lifecycle.RequestID, lifecycle.RequestID, managedKeyID, usageResult(usage.Failed, lifecycle.Outcome), endpointKind(endpoint), usage.ID); err != nil { + return fmt.Errorf("合并孤立请求索引: %w", err) + } + if _, err := tx.ExecContext(ctx, `DELETE FROM request_detail_index WHERE usage_id IS NULL AND lifecycle_request_id=?`, lifecycle.RequestID); err != nil { + return fmt.Errorf("删除已合并终态索引: %w", err) + } + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("提交孤立请求索引: %w", err) + } + return nil +} + +func sameOrphanProjection(usage orphanUsageProjection, lifecycle orphanLifecycleProjection) bool { + if usage.ManagedKeyID != "" && lifecycle.ManagedKeyID != "" { + return usage.ManagedKeyID == lifecycle.ManagedKeyID + } + return strings.TrimSpace(usage.Model) != "" && strings.EqualFold(strings.TrimSpace(usage.Model), strings.TrimSpace(lifecycle.Model)) +} diff --git a/internal/repository/sqlite_query.go b/internal/repository/sqlite_query.go new file mode 100644 index 0000000..695197c --- /dev/null +++ b/internal/repository/sqlite_query.go @@ -0,0 +1,280 @@ +package repository + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/base64" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "math" + "strings" + "time" + + "cpa-ext/internal/collection" +) + +type usageCursor struct { + Version int `json:"v"` + RequestedAt string `json:"t"` + ID int64 `json:"id"` + Direction string `json:"d"` + Page int `json:"p"` + Signature string `json:"s"` +} + +type detailKey struct { + ID int64 + RequestedAt string +} + +func normalizeUsageQuery(query collection.UsageQuery) (collection.UsageQuery, error) { + if query.Page < 1 { + query.Page = 1 + } + if query.PageSize < 1 { + query.PageSize = 100 + } + if query.PageSize > 100 { + return query, errors.New("page_size 不能超过 100") + } + query.KeyID = strings.TrimSpace(query.KeyID) + query.Model = strings.TrimSpace(query.Model) + query.Result = strings.ToLower(strings.TrimSpace(query.Result)) + query.AuthID = strings.TrimSpace(query.AuthID) + query.Endpoint = endpointKind(query.Endpoint) + query.RequestID = strings.TrimSpace(query.RequestID) + if query.From != nil && query.To != nil && !query.From.Before(*query.To) { + return query, errors.New("from 必须早于 to") + } + if query.Result != "" && query.Result != collection.UsageResultSucceeded && query.Result != collection.UsageResultFailed && query.Result != collection.UsageResultRejected && query.Result != collection.UsageResultCanceled { + return query, errors.New("result 必须是 succeeded、failed、rejected 或 canceled") + } + return query, nil +} + +func usageQuerySignature(query collection.UsageQuery) string { + payload := struct { + PageSize int + From string + To string + KeyID string + Model string + Result string + AuthID string + Endpoint string + RequestID string + }{PageSize: query.PageSize, KeyID: query.KeyID, Model: strings.ToLower(query.Model), Result: query.Result, AuthID: query.AuthID, Endpoint: query.Endpoint, RequestID: query.RequestID} + if query.From != nil { + payload.From = formatTime(*query.From) + } + if query.To != nil { + payload.To = formatTime(*query.To) + } + raw, _ := json.Marshal(payload) + sum := sha256.Sum256(raw) + return hex.EncodeToString(sum[:]) +} + +func encodeUsageCursor(cursor usageCursor) string { + raw, _ := json.Marshal(cursor) + return base64.RawURLEncoding.EncodeToString(raw) +} + +func decodeUsageCursor(raw, signature string) (usageCursor, error) { + decoded, err := base64.RawURLEncoding.DecodeString(strings.TrimSpace(raw)) + if err != nil { + return usageCursor{}, errors.New("cursor 无效") + } + var cursor usageCursor + if err := json.Unmarshal(decoded, &cursor); err != nil || cursor.Version != 1 || cursor.ID < 1 || cursor.RequestedAt == "" || cursor.Page < 1 || (cursor.Direction != "next" && cursor.Direction != "previous") { + return usageCursor{}, errors.New("cursor 无效") + } + if cursor.Signature != signature { + return usageCursor{}, errors.New("cursor 与当前筛选条件不匹配") + } + return cursor, nil +} + +func usageWhere(query collection.UsageQuery) (string, []any) { + conditions := []string{"1=1"} + args := make([]any, 0, 8) + if query.From != nil { + conditions = append(conditions, "d.requested_at>=?") + args = append(args, formatTime(*query.From)) + } + if query.To != nil { + conditions = append(conditions, "d.requested_at 0 { + totalPages = int(math.Ceil(float64(total) / float64(query.PageSize))) + } + signature := usageQuerySignature(query) + page := query.Page + order := "d.requested_at DESC, d.id DESC" + limitArgs := append([]any{}, args...) + if query.Cursor != "" { + cursor, cursorErr := decodeUsageCursor(query.Cursor, signature) + if cursorErr != nil { + return collection.UsagePage{}, cursorErr + } + page = cursor.Page + if cursor.Direction == "next" { + where += " AND (d.requested_at? OR (d.requested_at=? AND d.id>?))" + limitArgs = append(limitArgs, cursor.RequestedAt, cursor.RequestedAt, cursor.ID) + order = "d.requested_at ASC, d.id ASC" + } + } else { + limitArgs = append(limitArgs, query.PageSize, (page-1)*query.PageSize) + } + + statement := usageDetailSelectSQL + " WHERE " + where + " ORDER BY " + order + if query.Cursor != "" { + statement += " LIMIT ?" + limitArgs = append(limitArgs, query.PageSize) + } else { + statement += " LIMIT ? OFFSET ?" + } + rows, err := r.db.QueryContext(ctx, statement, limitArgs...) + if err != nil { + return collection.UsagePage{}, fmt.Errorf("查询请求明细: %w", err) + } + records, keys, err := scanUsageDetailRows(rows) + if err != nil { + return collection.UsagePage{}, err + } + if order == "d.requested_at ASC, d.id ASC" { + reverseRecords(records) + reverseKeys(keys) + } + result := collection.UsagePage{Records: records, Page: page, PageSize: query.PageSize, Total: total, TotalPages: totalPages} + if len(keys) > 0 { + if page > 1 { + result.PreviousCursor = encodeUsageCursor(usageCursor{Version: 1, RequestedAt: keys[0].RequestedAt, ID: keys[0].ID, Direction: "previous", Page: page - 1, Signature: signature}) + } + if page < totalPages { + last := keys[len(keys)-1] + result.NextCursor = encodeUsageCursor(usageCursor{Version: 1, RequestedAt: last.RequestedAt, ID: last.ID, Direction: "next", Page: page + 1, Signature: signature}) + } + } + return result, nil +} + +const usageDetailSelectSQL = ` +SELECT d.id, d.requested_at, + d.managed_key_id, d.request_id, COALESCE(u.execution_id, ''), + COALESCE(NULLIF(u.trace_id, ''), r.trace_id, ''), COALESCE(u.api_key, ''), + COALESCE(k.name, u.key_alias, ''), COALESCE(u.auth_id, ''), COALESCE(u.auth_index, ''), COALESCE(u.auth_type, ''), + COALESCE(NULLIF(u.model, ''), NULLIF(r.model, ''), r.requested_model, d.model, ''), + COALESCE(u.reasoning_effort, ''), COALESCE(u.service_tier, ''), COALESCE(u.speed, ''), + CASE WHEN d.result='succeeded' THEN 0 ELSE 1 END, + COALESCE(u.executor_type, ''), + CASE WHEN r.request_id IS NOT NULL THEN CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END ELSE COALESCE(u.request_type, '') END, + COALESCE(NULLIF(u.endpoint, ''), r.endpoint, ''), + COALESCE(u.input_tokens, 0), COALESCE(u.output_tokens, 0), COALESCE(u.reasoning_tokens, 0), COALESCE(u.cached_tokens, 0), + COALESCE(u.cache_read_tokens, 0), COALESCE(u.cache_write_tokens, 0), COALESCE(u.total_tokens, 0), + COALESCE(u.ttft_ns, 0), COALESCE(u.latency_ns, r.latency_ns, 0), u.cost_micros, + COALESCE(u.price_tier, ''), COALESCE(u.fast_requested, 0), COALESCE(u.fast_pricing_applied, 0), + COALESCE(u.price_multiplier_numerator, 1), COALESCE(u.price_multiplier_denominator, 1), COALESCE(u.client_ip, ''), + COALESCE(r.outcome, ''), COALESCE(r.status_code, 0), COALESCE(r.error, '') +FROM request_detail_index d +LEFT JOIN usage_records u ON u.id=d.usage_id +LEFT JOIN request_records r ON r.request_id=d.lifecycle_request_id +LEFT JOIN managed_keys k ON k.id=d.managed_key_id` + +func scanUsageDetailRows(rows *sql.Rows) ([]collection.Record, []detailKey, error) { + defer rows.Close() + var records []collection.Record + var keys []detailKey + for rows.Next() { + var record collection.Record + var key detailKey + var failed bool + var ttftNS, latencyNS int64 + var cost sql.NullInt64 + if err := rows.Scan(&key.ID, &key.RequestedAt, &record.ManagedKeyID, &record.RequestID, &record.ExecutionID, + &record.TraceID, &record.APIKey, &record.KeyAlias, &record.AuthID, &record.AuthIndex, &record.AuthType, + &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, nil, fmt.Errorf("读取请求明细: %w", err) + } + requestedAt, err := time.Parse(time.RFC3339Nano, key.RequestedAt) + if err != nil { + return nil, nil, fmt.Errorf("解析请求明细时间: %w", err) + } + record.RequestedAt = requestedAt + 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) + keys = append(keys, key) + } + if err := rows.Err(); err != nil { + return nil, nil, fmt.Errorf("遍历请求明细: %w", err) + } + return records, keys, nil +} + +func reverseRecords(values []collection.Record) { + for left, right := 0, len(values)-1; left < right; left, right = left+1, right-1 { + values[left], values[right] = values[right], values[left] + } +} + +func reverseKeys(values []detailKey) { + for left, right := 0, len(values)-1; left < right; left, right = left+1, right-1 { + values[left], values[right] = values[right], values[left] + } +} diff --git a/internal/repository/sqlite_query_migration_test.go b/internal/repository/sqlite_query_migration_test.go new file mode 100644 index 0000000..befa61c --- /dev/null +++ b/internal/repository/sqlite_query_migration_test.go @@ -0,0 +1,47 @@ +package repository + +import ( + "context" + "path/filepath" + "testing" + "time" + + "cpa-ext/internal/collection" +) + +func TestRequestDetailProjectionBackfillsExistingFacts(t *testing.T) { + path := filepath.Join(t.TempDir(), "usage.db") + store, err := OpenSQLiteUsage(path) + if err != nil { + t.Fatal(err) + } + startedAt := time.Date(2026, 8, 15, 12, 0, 0, 0, time.UTC) + if err := store.Insert(context.Background(), collection.Record{RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: "000000", Model: "deepseek-v4-flash", TotalTokens: 99}); err != nil { + t.Fatal(err) + } + if err := store.UpsertRequest(context.Background(), collection.RequestRecord{RequestID: "migrated-request", RequestedAt: startedAt, CompletedAt: startedAt.Add(time.Second), Model: "deepseek-v4-flash", Outcome: "succeeded", StatusCode: 200, Endpoint: "/v1/responses"}); err != nil { + t.Fatal(err) + } + if _, err := store.db.Exec(`DELETE FROM request_detail_index`); err != nil { + t.Fatal(err) + } + if _, err := store.db.Exec(`DELETE FROM cpa_ext_migrations WHERE name=?`, requestDetailProjectionMigration); err != nil { + t.Fatal(err) + } + if err := store.Close(); err != nil { + t.Fatal(err) + } + + reopened, err := OpenSQLiteUsage(path) + if err != nil { + t.Fatal(err) + } + defer reopened.Close() + page, err := reopened.QueryUsage(context.Background(), collection.UsageQuery{RequestID: "migrated-request", PageSize: 100}) + if err != nil { + t.Fatal(err) + } + if page.Total != 1 || len(page.Records) != 1 || page.Records[0].TotalTokens != 99 || page.Records[0].StatusCode != 200 { + t.Fatalf("backfilled projection mismatch: %+v", page) + } +} diff --git a/internal/repository/sqlite_query_scale_test.go b/internal/repository/sqlite_query_scale_test.go new file mode 100644 index 0000000..5400f51 --- /dev/null +++ b/internal/repository/sqlite_query_scale_test.go @@ -0,0 +1,68 @@ +package repository + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "cpa-ext/internal/collection" +) + +func TestSQLiteUsageQueryAtMillionRows(t *testing.T) { + if os.Getenv("CPA_EXT_MILLION_TEST") != "1" { + t.Skip("set CPA_EXT_MILLION_TEST=1 to run the million-row query check") + } + store, err := OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + started := time.Now() + _, err = store.db.Exec(` +WITH RECURSIVE sequence(value) AS ( + SELECT 1 UNION ALL SELECT value+1 FROM sequence WHERE value<1000000 +) +INSERT INTO request_detail_index + (lifecycle_request_id, requested_at, managed_key_id, model, result, auth_id, endpoint_kind, request_id) +SELECT 'load-' || value, + printf('2026-08-%02dT%02d:%02d:%02d.000000000Z', 1 + ((value / 86400) % 15), (value / 3600) % 24, (value / 60) % 60, value % 60), + 'key-' || (value % 5), 'model-' || (value % 3), + CASE WHEN value % 17=0 THEN 'failed' ELSE 'succeeded' END, + 'auth-' || (value % 2), 'responses', 'load-' || value +FROM sequence`) + if err != nil { + t.Fatal(err) + } + t.Logf("seeded one million indexed details in %s", time.Since(started)) + + started = time.Now() + page, err := store.QueryUsage(context.Background(), collection.UsageQuery{Page: 600, PageSize: 100, KeyID: "key-1", Model: "model-1"}) + if err != nil { + t.Fatal(err) + } + if page.Total == 0 || len(page.Records) == 0 { + t.Fatalf("deep filtered page is empty: %+v", page) + } + t.Logf("counted and loaded deep filtered page in %s", time.Since(started)) + + rows, err := store.db.Query(`EXPLAIN QUERY PLAN SELECT id FROM request_detail_index d WHERE d.managed_key_id=? ORDER BY d.requested_at DESC, d.id DESC LIMIT 100 OFFSET 900000`, "key-1") + if err != nil { + t.Fatal(err) + } + defer rows.Close() + var plan strings.Builder + for rows.Next() { + var id, parent, unused int + var detail string + if err := rows.Scan(&id, &parent, &unused, &detail); err != nil { + t.Fatal(err) + } + plan.WriteString(detail) + } + if !strings.Contains(plan.String(), "idx_request_detail_key") { + t.Fatalf("key query did not use projection index: %s", plan.String()) + } +} diff --git a/internal/repository/sqlite_query_test.go b/internal/repository/sqlite_query_test.go new file mode 100644 index 0000000..036f309 --- /dev/null +++ b/internal/repository/sqlite_query_test.go @@ -0,0 +1,154 @@ +package repository_test + +import ( + "context" + "fmt" + "path/filepath" + "testing" + "time" + + "cpa-ext/internal/collection" + "cpa-ext/internal/repository" +) + +func TestSQLiteUsageQueryPaginatesAndFilters(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + base := time.Date(2026, 8, 15, 0, 0, 0, 0, time.UTC) + for index := 0; index < 235; index++ { + failed := index%11 == 0 + record := collection.Record{ + ManagedKeyID: fmt.Sprintf("key-%d", index%3), RequestID: fmt.Sprintf("request-%03d", index), + ExecutionID: fmt.Sprintf("execution-%03d", index), RequestedAt: base.Add(time.Duration(index) * time.Second), + Model: fmt.Sprintf("model-%d", index%2), Failed: failed, AuthID: fmt.Sprintf("auth-%d", index%2), + Endpoint: "/v1/responses", TotalTokens: int64(index + 1), + } + if err := store.Insert(context.Background(), record); err != nil { + t.Fatalf("insert %d: %v", index, err) + } + } + + first, err := store.QueryUsage(context.Background(), collection.UsageQuery{Page: 1, PageSize: 100}) + if err != nil { + t.Fatal(err) + } + if len(first.Records) != 100 || first.Total != 235 || first.TotalPages != 3 || first.Page != 1 || first.NextCursor == "" || first.Records[0].RequestID != "request-234" { + t.Fatalf("unexpected first page: %+v first=%+v", first, first.Records[0]) + } + second, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: first.NextCursor}) + if err != nil { + t.Fatal(err) + } + if len(second.Records) != 100 || second.Page != 2 || second.PreviousCursor == "" || second.Records[0].RequestID != "request-134" { + t.Fatalf("unexpected second page: %+v first=%+v", second, second.Records[0]) + } + previous, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: second.PreviousCursor}) + if err != nil { + t.Fatal(err) + } + if len(previous.Records) != 100 || previous.Records[0].RequestID != first.Records[0].RequestID || previous.Records[99].RequestID != first.Records[99].RequestID { + t.Fatalf("previous page did not return the original first page") + } + + filters := []struct { + name string + query collection.UsageQuery + check func(collection.Record) bool + }{ + {name: "key", query: collection.UsageQuery{KeyID: "key-1"}, check: func(record collection.Record) bool { return record.ManagedKeyID == "key-1" }}, + {name: "model", query: collection.UsageQuery{Model: "MODEL-1"}, check: func(record collection.Record) bool { return record.Model == "model-1" }}, + {name: "failed", query: collection.UsageQuery{Result: collection.UsageResultFailed}, check: func(record collection.Record) bool { return record.Failed }}, + {name: "auth", query: collection.UsageQuery{AuthID: "auth-1"}, check: func(record collection.Record) bool { return record.AuthID == "auth-1" }}, + {name: "endpoint", query: collection.UsageQuery{Endpoint: "responses"}, check: func(record collection.Record) bool { return record.Endpoint == "/v1/responses" }}, + {name: "request", query: collection.UsageQuery{RequestID: "request-123"}, check: func(record collection.Record) bool { return record.RequestID == "request-123" }}, + } + for _, test := range filters { + t.Run(test.name, func(t *testing.T) { + test.query.PageSize = 100 + page, err := store.QueryUsage(context.Background(), test.query) + if err != nil { + t.Fatal(err) + } + if len(page.Records) == 0 { + t.Fatal("filter returned no records") + } + for _, record := range page.Records { + if !test.check(record) { + t.Fatalf("unexpected filtered record: %+v", record) + } + } + }) + } + + if _, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 100, Cursor: first.NextCursor, Model: "model-1"}); err == nil { + t.Fatal("cursor from another filter unexpectedly accepted") + } +} + +func TestSQLiteUsageQueryKeepsCursorStableAndProjectsOrphans(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + + base := time.Date(2026, 8, 15, 12, 0, 0, 0, time.UTC) + for index := 0; index < 4; index++ { + if err := store.Insert(context.Background(), collection.Record{RequestID: fmt.Sprintf("request-%d", index), ExecutionID: fmt.Sprintf("execution-%d", index), RequestedAt: base.Add(time.Duration(index) * time.Second), Model: "model", TotalTokens: 10}); err != nil { + t.Fatal(err) + } + } + first, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 2}) + if err != nil { + t.Fatal(err) + } + if err := store.Insert(context.Background(), collection.Record{RequestID: "newest", ExecutionID: "newest", RequestedAt: base.Add(10 * time.Second), Model: "model"}); err != nil { + t.Fatal(err) + } + second, err := store.QueryUsage(context.Background(), collection.UsageQuery{PageSize: 2, Cursor: first.NextCursor}) + if err != nil { + t.Fatal(err) + } + if second.Records[0].RequestID != "request-1" || second.Records[1].RequestID != "request-0" { + t.Fatalf("cursor shifted after a new insert: %+v", second.Records) + } + + startedAt := base.Add(time.Minute) + if err := store.Insert(context.Background(), collection.Record{RequestedAt: startedAt.Add(5 * time.Millisecond), APIKey: "000000", Model: "deepseek-v4-flash", TotalTokens: 99}); err != nil { + t.Fatal(err) + } + if err := store.UpsertRequest(context.Background(), collection.RequestRecord{RequestID: "orphan-request", RequestedAt: startedAt, CompletedAt: startedAt.Add(time.Second), Model: "deepseek-v4-flash", Stream: true, Outcome: "canceled", StatusCode: 499, Endpoint: "/v1/responses"}); err != nil { + t.Fatal(err) + } + page, err := store.QueryUsage(context.Background(), collection.UsageQuery{RequestID: "orphan-request", PageSize: 100}) + if err != nil { + t.Fatal(err) + } + if page.Total != 1 || len(page.Records) != 1 || page.Records[0].TotalTokens != 99 || page.Records[0].Outcome != "canceled" || page.Records[0].StatusCode != 499 { + t.Fatalf("orphan projection mismatch: %+v", page) + } +} + +func TestSQLiteUsageDashboardUsesAllIndexedHistory(t *testing.T) { + store, err := repository.OpenSQLiteUsage(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer store.Close() + location, _ := time.LoadLocation("Asia/Shanghai") + today := time.Date(2026, 8, 15, 12, 0, 0, 0, location) + if err := store.Insert(context.Background(), collection.Record{ManagedKeyID: "key-a", RequestedAt: today.Add(-time.Hour), Model: "model", InputTokens: 10, OutputTokens: 2, TotalTokens: 12, CostMicros: nil}); err != nil { + t.Fatal(err) + } + dashboard, err := store.UsageDashboard(context.Background(), today, 7) + if err != nil { + t.Fatal(err) + } + if dashboard.Today.Requests != 1 || dashboard.Today.TotalTokens != 12 || len(dashboard.Users) != 1 || len(dashboard.Days) != 7 { + t.Fatalf("unexpected dashboard: %+v", dashboard) + } +} diff --git a/internal/repository/sqlite_summary.go b/internal/repository/sqlite_summary.go new file mode 100644 index 0000000..8a8d070 --- /dev/null +++ b/internal/repository/sqlite_summary.go @@ -0,0 +1,78 @@ +package repository + +import ( + "context" + "fmt" + "time" + + "cpa-ext/internal/collection" +) + +func (r *SQLiteUsageRepository) UsageDashboard(ctx context.Context, today time.Time, days int) (collection.UsageDashboard, error) { + if days < 1 { + days = 7 + } + start := time.Date(today.Year(), today.Month(), today.Day(), 0, 0, 0, 0, today.Location()) + dashboard := collection.UsageDashboard{} + if err := r.db.QueryRowContext(ctx, ` +SELECT COUNT(*), COALESCE(SUM(u.input_tokens),0), COALESCE(SUM(u.output_tokens),0), + COALESCE(SUM(u.total_tokens),0), COALESCE(SUM(u.cost_micros),0) +FROM request_detail_index d LEFT JOIN usage_records u ON u.id=d.usage_id +WHERE d.requested_at>=?`, formatTime(start)).Scan(&dashboard.Today.Requests, &dashboard.Today.InputTokens, &dashboard.Today.OutputTokens, &dashboard.Today.TotalTokens, &dashboard.Today.CostMicros); err != nil { + return dashboard, fmt.Errorf("汇总今日用量: %w", err) + } + rows, err := r.db.QueryContext(ctx, ` +SELECT d.managed_key_id, COALESCE(k.name, '未识别'), + SUM(CASE WHEN d.requested_at>=? THEN 1 ELSE 0 END), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.input_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.output_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.total_tokens ELSE 0 END),0), + COALESCE(SUM(CASE WHEN d.requested_at>=? THEN u.cost_micros ELSE 0 END),0), + MAX(d.requested_at) +FROM request_detail_index d +LEFT JOIN usage_records u ON u.id=d.usage_id +LEFT JOIN managed_keys k ON k.id=d.managed_key_id +GROUP BY d.managed_key_id, COALESCE(k.name, '未识别') +ORDER BY 6 DESC, 2`, formatTime(start), formatTime(start), formatTime(start), formatTime(start), formatTime(start)) + if err != nil { + return dashboard, fmt.Errorf("汇总用户用量: %w", err) + } + for rows.Next() { + var item collection.UserUsageSummary + var last string + if err := rows.Scan(&item.KeyID, &item.KeyAlias, &item.Today.Requests, &item.Today.InputTokens, &item.Today.OutputTokens, &item.Today.TotalTokens, &item.Today.CostMicros, &last); err != nil { + _ = rows.Close() + return dashboard, fmt.Errorf("读取用户用量汇总: %w", err) + } + if parsed, parseErr := time.Parse(time.RFC3339Nano, last); parseErr == nil { + item.LastUsedAt = &parsed + } + dashboard.Users = append(dashboard.Users, item) + } + _ = rows.Close() + + dailyStart := start.AddDate(0, 0, -(days - 1)) + rows, err = r.db.QueryContext(ctx, ` +SELECT DATE(d.requested_at, '+8 hours'), COALESCE(SUM(u.total_tokens),0) +FROM request_detail_index d LEFT JOIN usage_records u ON u.id=d.usage_id +WHERE d.requested_at>=? GROUP BY DATE(d.requested_at, '+8 hours') ORDER BY 1`, formatTime(dailyStart)) + if err != nil { + return dashboard, fmt.Errorf("汇总每日用量: %w", err) + } + byDate := make(map[string]int64) + for rows.Next() { + var date string + var tokens int64 + if err := rows.Scan(&date, &tokens); err != nil { + _ = rows.Close() + return dashboard, fmt.Errorf("读取每日用量: %w", err) + } + byDate[date] = tokens + } + _ = rows.Close() + for index := 0; index < days; index++ { + date := dailyStart.AddDate(0, 0, index) + dashboard.Days = append(dashboard.Days, collection.DailyUsageSummary{Date: date, TotalTokens: byDate[date.Format("2006-01-02")]}) + } + return dashboard, nil +} diff --git a/internal/repository/sqlite_usage.go b/internal/repository/sqlite_usage.go index cf754c1..99a61eb 100644 --- a/internal/repository/sqlite_usage.go +++ b/internal/repository/sqlite_usage.go @@ -3,7 +3,9 @@ package repository import ( "context" + "crypto/sha256" "database/sql" + "encoding/hex" "errors" "fmt" "net/url" @@ -56,6 +58,7 @@ CREATE TABLE IF NOT EXISTS usage_records ( ); 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 INDEX IF NOT EXISTS idx_usage_records_model ON usage_records(model COLLATE NOCASE) WHERE model <> ''; CREATE TABLE IF NOT EXISTS request_records ( request_id TEXT PRIMARY KEY, @@ -74,6 +77,44 @@ CREATE TABLE IF NOT EXISTS request_records ( latency_ns INTEGER NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_request_records_requested_at ON request_records(requested_at DESC); +CREATE INDEX IF NOT EXISTS idx_usage_records_request_id ON usage_records(request_id) WHERE request_id <> ''; + +CREATE TABLE IF NOT EXISTS request_detail_index ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + usage_id INTEGER, + lifecycle_request_id TEXT NOT NULL DEFAULT '', + requested_at TEXT NOT NULL, + managed_key_id TEXT NOT NULL DEFAULT '', + model TEXT NOT NULL DEFAULT '' COLLATE NOCASE, + result TEXT NOT NULL DEFAULT 'succeeded', + auth_id TEXT NOT NULL DEFAULT '', + endpoint_kind TEXT NOT NULL DEFAULT '', + request_id TEXT NOT NULL DEFAULT '', + FOREIGN KEY(usage_id) REFERENCES usage_records(id) ON DELETE CASCADE +); +CREATE UNIQUE INDEX IF NOT EXISTS idx_request_detail_usage +ON request_detail_index(usage_id) WHERE usage_id IS NOT NULL; +CREATE UNIQUE INDEX IF NOT EXISTS idx_request_detail_lifecycle_only +ON request_detail_index(lifecycle_request_id) WHERE usage_id IS NULL; +CREATE INDEX IF NOT EXISTS idx_request_detail_time +ON request_detail_index(requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_key +ON request_detail_index(managed_key_id, requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_model +ON request_detail_index(model COLLATE NOCASE, requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_result +ON request_detail_index(result, requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_auth +ON request_detail_index(auth_id, requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_endpoint +ON request_detail_index(endpoint_kind, requested_at DESC, id DESC); +CREATE INDEX IF NOT EXISTS idx_request_detail_request +ON request_detail_index(request_id, requested_at DESC, id DESC); + +CREATE TABLE IF NOT EXISTS cpa_ext_migrations ( + name TEXT PRIMARY KEY, + completed_at TEXT NOT NULL +); CREATE TABLE IF NOT EXISTS managed_keys ( id TEXT PRIMARY KEY, @@ -127,6 +168,59 @@ CREATE TABLE IF NOT EXISTS model_prices ( fast_multiplier_denominator INTEGER NOT NULL DEFAULT 2, updated_at TEXT NOT NULL ); + +CREATE TABLE IF NOT EXISTS billing_accounts ( + managed_key_id TEXT PRIMARY KEY, + quota_micros INTEGER NOT NULL DEFAULT 0 CHECK(quota_micros >= 0), + reset_period TEXT NOT NULL DEFAULT 'none' CHECK(reset_period IN ('none', 'daily', 'weekly', 'monthly')), + next_reset_at TEXT, + reset_anchor_day INTEGER NOT NULL DEFAULT 0, + max_concurrency INTEGER NOT NULL DEFAULT 4 CHECK(max_concurrency BETWEEN 1 AND 64), + current_cycle_sequence INTEGER NOT NULL DEFAULT 1, + updated_at TEXT NOT NULL, + FOREIGN KEY(managed_key_id) REFERENCES managed_keys(id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS billing_cycles ( + managed_key_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + started_at TEXT NOT NULL, + ended_at TEXT, + quota_micros INTEGER NOT NULL DEFAULT 0, + spent_micros INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY(managed_key_id, sequence), + FOREIGN KEY(managed_key_id) REFERENCES managed_keys(id) ON DELETE CASCADE +); + +CREATE TABLE IF NOT EXISTS billing_admissions ( + request_id TEXT PRIMARY KEY, + managed_key_id TEXT NOT NULL, + cycle_sequence INTEGER NOT NULL, + opened_at TEXT NOT NULL, + closed_at TEXT, + status TEXT NOT NULL DEFAULT 'open', + FOREIGN KEY(managed_key_id) REFERENCES managed_keys(id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_billing_admissions_active +ON billing_admissions(managed_key_id, status); + +CREATE TABLE IF NOT EXISTS billing_ledger ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + event_key TEXT NOT NULL UNIQUE, + managed_key_id TEXT NOT NULL, + cycle_sequence INTEGER NOT NULL, + kind TEXT NOT NULL, + amount_micros INTEGER NOT NULL, + balance_after_micros INTEGER NOT NULL, + request_id TEXT NOT NULL DEFAULT '', + execution_id TEXT NOT NULL DEFAULT '', + model TEXT NOT NULL DEFAULT '', + occurred_at TEXT NOT NULL, + created_at TEXT NOT NULL, + FOREIGN KEY(managed_key_id) REFERENCES managed_keys(id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_billing_ledger_key +ON billing_ledger(managed_key_id, id DESC); ` const ( @@ -170,12 +264,30 @@ func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) { `ALTER TABLE usage_records ADD COLUMN auth_id TEXT NOT NULL DEFAULT ''`, `ALTER TABLE usage_records ADD COLUMN auth_index TEXT NOT NULL DEFAULT ''`, `ALTER TABLE usage_records ADD COLUMN auth_type TEXT NOT NULL DEFAULT ''`, + `ALTER TABLE usage_records ADD COLUMN billing_event_key TEXT NOT NULL DEFAULT ''`, } { if _, err := db.Exec(migration); err != nil && !strings.Contains(strings.ToLower(err.Error()), "duplicate column") { _ = db.Close() return nil, fmt.Errorf("迁移 SQLite schema: %w", err) } } + if _, err := db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_records_billing_event ON usage_records(billing_event_key) WHERE billing_event_key <> ''`); err != nil { + _ = db.Close() + return nil, fmt.Errorf("创建用量计费幂等索引: %w", err) + } + now := formatTime(time.Now().UTC()) + if _, err := db.Exec(` +INSERT OR IGNORE INTO billing_accounts (managed_key_id, updated_at) +SELECT id, ? FROM managed_keys`, now); err != nil { + _ = db.Close() + return nil, fmt.Errorf("迁移 Key 额度账户: %w", err) + } + if _, err := db.Exec(` +INSERT OR IGNORE INTO billing_cycles (managed_key_id, sequence, started_at, quota_micros, spent_micros) +SELECT k.id, 1, k.created_at, 0, 0 FROM managed_keys k`); err != nil { + _ = db.Close() + return nil, fmt.Errorf("迁移 Key 额度周期: %w", err) + } if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_usage_records_managed_key ON usage_records(managed_key_id, requested_at DESC)`); err != nil { _ = db.Close() return nil, fmt.Errorf("创建 Key 用量索引: %w", err) @@ -184,7 +296,12 @@ func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) { _ = db.Close() return nil, fmt.Errorf("创建 Key 请求索引: %w", err) } - return &SQLiteUsageRepository{db: db}, nil + repository := &SQLiteUsageRepository{db: db} + if err := repository.ensureRequestDetailProjection(context.Background()); err != nil { + _ = db.Close() + return nil, err + } + return repository, nil } func sqliteDSN(databasePath string) (string, error) { @@ -214,16 +331,22 @@ func sqliteDSN(databasePath string) (string, error) { } func (r *SQLiteUsageRepository) Insert(ctx context.Context, record collection.Record) error { - _, err := r.db.ExecContext(ctx, ` + eventKey := usageBillingEventKey(record) + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("开始写入用量: %w", err) + } + defer func() { _ = tx.Rollback() }() + result, err := tx.ExecContext(ctx, ` INSERT OR IGNORE INTO usage_records ( - managed_key_id, request_id, execution_id, trace_id, requested_at, api_key, key_alias, auth_id, auth_index, auth_type, + managed_key_id, request_id, execution_id, trace_id, billing_event_key, requested_at, api_key, key_alias, auth_id, auth_index, auth_type, 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.ManagedKeyID, record.RequestID, record.ExecutionID, record.TraceID, +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + record.ManagedKeyID, record.RequestID, record.ExecutionID, record.TraceID, eventKey, record.RequestedAt.UTC().Format(time.RFC3339Nano), record.APIKey, record.KeyAlias, record.AuthID, record.AuthIndex, record.AuthType, record.Model, record.ReasoningEffort, record.ServiceTier, record.Speed, record.Failed, @@ -236,6 +359,86 @@ INSERT OR IGNORE INTO usage_records ( if err != nil { return fmt.Errorf("写入用量记录: %w", err) } + inserted, _ := result.RowsAffected() + var usageID int64 + if inserted > 0 { + usageID, _ = result.LastInsertId() + if err := insertUsageProjection(ctx, tx, usageID); err != nil { + return err + } + if record.RequestID != "" { + if _, err := syncRequestProjection(ctx, tx, record.RequestID); err != nil { + return err + } + } + } + if inserted > 0 && record.ManagedKeyID != "" && record.CostMicros != nil { + if err := chargeUsage(ctx, tx, record, eventKey); err != nil { + return err + } + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("提交用量和账目: %w", err) + } + if inserted > 0 && record.RequestID == "" { + if err := r.reconcileOrphanProjection(ctx, record.RequestedAt); err != nil { + return err + } + } + return nil +} + +func usageBillingEventKey(record collection.Record) string { + if executionID := strings.TrimSpace(record.ExecutionID); executionID != "" { + return "execution:" + record.ManagedKeyID + ":" + executionID + } + hash := sha256.New() + _, _ = fmt.Fprintf(hash, "%s\x00%s\x00%s\x00%s\x00%s\x00%s\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%d\x00%t", + record.ManagedKeyID, record.RequestID, record.AuthID, record.AuthIndex, record.Model, + record.RequestedAt.UTC().Format(time.RFC3339Nano), record.InputTokens, record.OutputTokens, + record.CacheReadTokens, record.CacheWriteTokens, record.TotalTokens, record.TTFT, + record.Latency, record.StatusCode, record.Failed) + return "usage:" + hex.EncodeToString(hash.Sum(nil)) +} + +func chargeUsage(ctx context.Context, tx *sql.Tx, record collection.Record, eventKey string) error { + requestedAt := record.RequestedAt + if requestedAt.IsZero() { + requestedAt = time.Now().UTC() + } + var sequence, quota, spent int64 + err := tx.QueryRowContext(ctx, ` +SELECT sequence, quota_micros, spent_micros FROM billing_cycles +WHERE managed_key_id=? AND started_at <= ? AND (ended_at IS NULL OR ? < ended_at) +ORDER BY sequence DESC LIMIT 1`, record.ManagedKeyID, formatTime(requestedAt), formatTime(requestedAt)).Scan(&sequence, "a, &spent) + if errors.Is(err, sql.ErrNoRows) { + err = tx.QueryRowContext(ctx, ` +SELECT c.sequence, c.quota_micros, c.spent_micros FROM billing_accounts a +JOIN billing_cycles c ON c.managed_key_id=a.managed_key_id AND c.sequence=a.current_cycle_sequence +WHERE a.managed_key_id=?`, record.ManagedKeyID).Scan(&sequence, "a, &spent) + } + if err != nil { + return fmt.Errorf("定位用量额度周期: %w", err) + } + cost := *record.CostMicros + balance := quota - spent - cost + result, err := tx.ExecContext(ctx, ` +INSERT OR IGNORE INTO billing_ledger (event_key, managed_key_id, cycle_sequence, kind, amount_micros, + balance_after_micros, request_id, execution_id, model, occurred_at, created_at) +VALUES (?, ?, ?, 'charge', ?, ?, ?, ?, ?, ?, ?)`, "charge:"+eventKey, record.ManagedKeyID, + sequence, -cost, balance, record.RequestID, record.ExecutionID, record.Model, + formatTime(requestedAt), formatTime(time.Now().UTC())) + if err != nil { + return fmt.Errorf("记录用量扣费: %w", err) + } + inserted, _ := result.RowsAffected() + if inserted > 0 { + if _, err := tx.ExecContext(ctx, ` +UPDATE billing_cycles SET spent_micros=spent_micros+? WHERE managed_key_id=? AND sequence=?`, + cost, record.ManagedKeyID, sequence); err != nil { + return fmt.Errorf("更新额度消费: %w", err) + } + } return nil } @@ -256,7 +459,12 @@ func (r *SQLiteUsageRepository) UpsertRequest(ctx context.Context, record collec if latency < 0 { latency = 0 } - _, err := r.db.ExecContext(ctx, ` + tx, err := r.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("开始写入请求终态: %w", err) + } + defer func() { _ = tx.Rollback() }() + _, err = tx.ExecContext(ctx, ` INSERT INTO request_records ( request_id, managed_key_id, trace_id, requested_at, completed_at, model, requested_model, source_format, stream, outcome, status_code, error, endpoint, latency_ns @@ -283,6 +491,18 @@ ON CONFLICT(request_id) DO UPDATE SET if err != nil { return fmt.Errorf("写入请求终态: %w", err) } + linked, err := syncRequestProjection(ctx, tx, requestID) + if err != nil { + return err + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("提交请求终态和索引: %w", err) + } + if !linked { + if err := r.reconcileOrphanProjection(ctx, requestedAt); err != nil { + return err + } + } return nil } diff --git a/internal/web/ui.html b/internal/web/ui.html index 5495f62..79dfe63 100644 --- a/internal/web/ui.html +++ b/internal/web/ui.html @@ -139,12 +139,26 @@ .ok { color: #34d399; } .failed { color: #f87171; } .drawer-overlay { position: fixed; z-index: 40; inset: 0; background: #020617a6; backdrop-filter: blur(2px); } - .drawer { position: fixed; z-index: 50; top: 0; right: 0; display: flex; flex-direction: column; width: min(430px, 100vw); height: 100vh; border-left: 1px solid var(--border); background: #0f1621; box-shadow: -20px 0 60px #0008; } + .drawer { position: fixed; z-index: 50; top: 0; right: 0; display: flex; flex-direction: column; width: min(540px, 100vw); height: 100vh; border-left: 1px solid var(--border); background: #0f1621; box-shadow: -20px 0 60px #0008; } .drawer-header { display: flex; align-items: flex-start; justify-content: space-between; gap: 14px; padding: 17px 18px 13px; border-bottom: 1px solid var(--border-soft); } .drawer-header h2 { margin-bottom: 3px; } .drawer-header p { margin-bottom: 0; color: var(--muted); font-size: 11px; } .drawer-body { flex: 1; overflow: auto; padding: 16px 18px; } .drawer-form { display: grid; gap: 13px; } + .drawer-section { display: grid; gap: 10px; padding: 12px; border: 1px solid var(--border-soft); border-radius: 8px; background: #0c131e; } + .drawer-section h3 { margin: 0; font-size: 11px; } + .field-hint { margin: -2px 0 0; color: #64748b; font-size: 9px; } + .billing-preview { display: grid; grid-template-columns: repeat(3, 1fr); gap: 7px; } + .billing-preview div { padding: 9px; border: 1px solid var(--border-soft); border-radius: 7px; background: #111925; } + .billing-preview span { display: block; color: var(--muted); font-size: 9px; } + .billing-preview strong { display: block; margin-top: 3px; font-size: 14px; } + .billing-tools { display: flex; align-items: center; justify-content: space-between; gap: 8px; } + .ledger-list { display: grid; max-height: 210px; overflow: auto; gap: 5px; } + .ledger-item { display: grid; grid-template-columns: 1fr auto; gap: 3px 10px; padding: 8px; border: 1px solid var(--border-soft); border-radius: 7px; background: #101824; } + .ledger-item strong { font-size: 10px; } + .ledger-item .positive { color: var(--success); } + .ledger-item .negative { color: var(--danger); } + .ledger-item small { grid-column: 1 / -1; overflow: hidden; color: var(--muted); font-size: 9px; text-overflow: ellipsis; white-space: nowrap; } .field-row { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; } .drawer-footer { display: flex; justify-content: space-between; gap: 8px; padding: 12px 18px; border-top: 1px solid var(--border-soft); } .drawer-footer div { display: flex; gap: 7px; } @@ -163,8 +177,14 @@ .page-buttons button.active { border-color: #0e7490; background: #0e7490; color: white; } .page-buttons button:disabled { cursor: default; opacity: .4; } .page-ellipsis { width: 22px; color: var(--muted); text-align: center; font-size: 10px; } + .usage-filters { display: grid; grid-template-columns: repeat(6, minmax(120px, 1fr)); gap: 9px; margin-bottom: 10px; padding: 11px; } + .usage-filter { display: grid; gap: 5px; min-width: 0; color: var(--muted); font-size: 9px; } + .usage-filter input, .usage-filter select { height: 32px; padding: 6px 8px; font-size: 11px; } + .usage-filter.request-filter { grid-column: span 2; } + .usage-filter-actions { display: flex; align-items: end; justify-content: flex-end; gap: 7px; grid-column: span 2; } + @media (max-width: 1100px) { .usage-filters { grid-template-columns: repeat(3, minmax(140px, 1fr)); } } @media (max-width: 900px) { .price-layout { grid-template-columns: 220px minmax(0, 1fr); } .price-grid, .price-grid.long-grid { grid-template-columns: repeat(2, minmax(130px, 1fr)); } .summary-grid { grid-template-columns: repeat(3, minmax(130px, 1fr)); overflow-x: auto; } .chart-grid { grid-template-columns: 1fr; } } - @media (max-width: 680px) { .topbar { align-items: flex-start; height: auto; padding: 11px 12px; } .brand span, .auth-box label { display: none; } .auth-box { flex-wrap: wrap; justify-content: flex-end; } #key { width: 128px; } .demo-banner { padding-right: 12px; padding-left: 12px; } .workspace-nav { padding-right: 12px; padding-left: 12px; } main { padding: 14px 12px 22px; } .section-heading { align-items: stretch; flex-direction: column; } .section-heading .tools { justify-content: flex-start; } .price-layout { grid-template-columns: 1fr; } .price-sidebar { border-right: 0; border-bottom: 1px solid var(--border-soft); } .price-grid, .price-grid.long-grid, .field-row, .fast-row { grid-template-columns: 1fr; } .drawer { width: 100vw; } .summary-grid { margin-right: -12px; padding-right: 12px; } .pagination { align-items: flex-start; flex-direction: column; } } + @media (max-width: 680px) { .topbar { align-items: flex-start; height: auto; padding: 11px 12px; } .brand span, .auth-box label { display: none; } .auth-box { flex-wrap: wrap; justify-content: flex-end; } #key { width: 128px; } .demo-banner { padding-right: 12px; padding-left: 12px; } .workspace-nav { padding-right: 12px; padding-left: 12px; } main { padding: 14px 12px 22px; } .section-heading { align-items: stretch; flex-direction: column; } .section-heading .tools { justify-content: flex-start; } .price-layout { grid-template-columns: 1fr; } .price-sidebar { border-right: 0; border-bottom: 1px solid var(--border-soft); } .price-grid, .price-grid.long-grid, .field-row, .fast-row { grid-template-columns: 1fr; } .drawer { width: 100vw; } .summary-grid { margin-right: -12px; padding-right: 12px; } .pagination { align-items: flex-start; flex-direction: column; } .usage-filters { grid-template-columns: 1fr 1fr; } .usage-filter.request-filter, .usage-filter-actions { grid-column: span 2; } } @@ -189,13 +209,24 @@

今日用户用量

近 7 日 Token

-
点击“管理”修改路由与权限
-
用户完整 Key状态路由允许模型操作
+
点击“管理”修改访问、路由与额度
+
用户完整 Key状态余额 / 额度并发路由允许模型操作
选择列
+
+ + + + + + + + +
+
@@ -225,13 +256,22 @@

管理 Key

修改用户状态、路由和模型权限

- - +

用户

+
+

访问与路由

- +

多个规则用逗号分隔；* 匹配任意字符，未匹配的模型会被拒绝。

+

额度与周期

+
额度$0.0000
已用$0.0000
余额$0.0000
+
+
+

请求结束后按实际费用扣款；最后一个并发请求可能产生负余额。

+
+
+

最近账目

@@ -240,11 +280,14 @@