763 lines
30 KiB
Go
763 lines
30 KiB
Go
// Package repository 实现 cpa-ext 的本地持久化。
|
|
package repository
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"database/sql"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
|
|
"cpa-ext/internal/collection"
|
|
_ "github.com/mattn/go-sqlite3"
|
|
)
|
|
|
|
const usageSchema = `
|
|
CREATE TABLE IF NOT EXISTS usage_records (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
managed_key_id TEXT NOT NULL DEFAULT '',
|
|
request_id TEXT NOT NULL DEFAULT '',
|
|
execution_id TEXT NOT NULL DEFAULT '',
|
|
trace_id TEXT NOT NULL DEFAULT '',
|
|
requested_at TEXT NOT NULL,
|
|
api_key TEXT NOT NULL DEFAULT '',
|
|
key_alias TEXT NOT NULL DEFAULT '',
|
|
auth_id TEXT NOT NULL DEFAULT '',
|
|
auth_index TEXT NOT NULL DEFAULT '',
|
|
auth_type TEXT NOT NULL DEFAULT '',
|
|
model TEXT NOT NULL DEFAULT '',
|
|
reasoning_effort TEXT NOT NULL DEFAULT '',
|
|
service_tier TEXT NOT NULL DEFAULT '',
|
|
speed TEXT NOT NULL DEFAULT '',
|
|
failed INTEGER NOT NULL,
|
|
executor_type TEXT NOT NULL DEFAULT '',
|
|
request_type TEXT NOT NULL DEFAULT '',
|
|
endpoint TEXT NOT NULL DEFAULT '',
|
|
input_tokens INTEGER NOT NULL DEFAULT 0,
|
|
output_tokens INTEGER NOT NULL DEFAULT 0,
|
|
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_read_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cache_write_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
ttft_ns INTEGER NOT NULL DEFAULT 0,
|
|
latency_ns INTEGER NOT NULL DEFAULT 0,
|
|
cost_micros INTEGER,
|
|
price_tier TEXT NOT NULL DEFAULT '',
|
|
fast_requested INTEGER NOT NULL DEFAULT 0,
|
|
fast_pricing_applied INTEGER NOT NULL DEFAULT 0,
|
|
price_multiplier_numerator INTEGER NOT NULL DEFAULT 1,
|
|
price_multiplier_denominator INTEGER NOT NULL DEFAULT 1,
|
|
client_ip TEXT NOT NULL DEFAULT ''
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_usage_records_requested_at ON usage_records(requested_at DESC, id DESC);
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_usage_records_execution_id ON usage_records(execution_id) WHERE execution_id <> '';
|
|
CREATE INDEX IF NOT EXISTS idx_usage_records_model ON usage_records(model COLLATE NOCASE) WHERE model <> '';
|
|
CREATE INDEX IF NOT EXISTS idx_usage_records_missing_cost_model
|
|
ON usage_records(model COLLATE NOCASE, id) WHERE cost_micros IS NULL;
|
|
|
|
CREATE TABLE IF NOT EXISTS request_records (
|
|
request_id TEXT PRIMARY KEY,
|
|
managed_key_id TEXT NOT NULL DEFAULT '',
|
|
trace_id TEXT NOT NULL DEFAULT '',
|
|
requested_at TEXT NOT NULL,
|
|
completed_at TEXT NOT NULL,
|
|
model TEXT NOT NULL DEFAULT '',
|
|
requested_model TEXT NOT NULL DEFAULT '',
|
|
source_format TEXT NOT NULL DEFAULT '',
|
|
stream INTEGER NOT NULL DEFAULT 0,
|
|
outcome TEXT NOT NULL,
|
|
status_code INTEGER NOT NULL DEFAULT 0,
|
|
error TEXT NOT NULL DEFAULT '',
|
|
endpoint TEXT NOT NULL DEFAULT '',
|
|
latency_ns INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_request_records_requested_at ON request_records(requested_at DESC);
|
|
CREATE 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,
|
|
name TEXT NOT NULL COLLATE NOCASE UNIQUE,
|
|
secret TEXT NOT NULL UNIQUE,
|
|
status TEXT NOT NULL CHECK(status IN ('active', 'disabled', 'archived')),
|
|
route_mode TEXT NOT NULL CHECK(route_mode IN ('auto', 'strict')),
|
|
upstream_account_id TEXT NOT NULL DEFAULT '',
|
|
all_models INTEGER NOT NULL DEFAULT 1,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS managed_key_models (
|
|
managed_key_id TEXT NOT NULL,
|
|
model TEXT NOT NULL COLLATE NOCASE,
|
|
PRIMARY KEY(managed_key_id, model),
|
|
FOREIGN KEY(managed_key_id) REFERENCES managed_keys(id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS upstream_accounts (
|
|
id TEXT PRIMARY KEY,
|
|
cpa_auth_id TEXT NOT NULL UNIQUE,
|
|
cpa_auth_index TEXT NOT NULL DEFAULT '',
|
|
provider TEXT NOT NULL DEFAULT '',
|
|
display_name TEXT NOT NULL DEFAULT '',
|
|
status TEXT NOT NULL DEFAULT '',
|
|
status_message TEXT NOT NULL DEFAULT '',
|
|
disabled INTEGER NOT NULL DEFAULT 0,
|
|
unavailable INTEGER NOT NULL DEFAULT 0,
|
|
priority INTEGER NOT NULL DEFAULT 0,
|
|
last_seen_at TEXT NOT NULL
|
|
);
|
|
CREATE UNIQUE INDEX IF NOT EXISTS idx_upstream_auth_index ON upstream_accounts(cpa_auth_index) WHERE cpa_auth_index <> '';
|
|
|
|
CREATE TABLE IF NOT EXISTS model_prices (
|
|
model TEXT PRIMARY KEY,
|
|
input_rate_micros INTEGER NOT NULL,
|
|
cache_read_rate_micros INTEGER NOT NULL,
|
|
cache_write_rate_micros INTEGER NOT NULL,
|
|
output_rate_micros INTEGER NOT NULL,
|
|
long_context_enabled INTEGER NOT NULL DEFAULT 0,
|
|
long_context_threshold INTEGER NOT NULL DEFAULT 0,
|
|
long_context_comparison TEXT NOT NULL DEFAULT 'gt',
|
|
long_input_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_cache_read_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_cache_write_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
long_output_rate_micros INTEGER NOT NULL DEFAULT 0,
|
|
fast_pricing_enabled INTEGER NOT NULL DEFAULT 0,
|
|
fast_multiplier_numerator INTEGER NOT NULL DEFAULT 5,
|
|
fast_multiplier_denominator INTEGER NOT NULL DEFAULT 2,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
|
|
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 (
|
|
// CLIProxyAPI UsageRecord currently has no RequestID. The usage and terminal
|
|
// callbacks nevertheless carry request-start timestamps from the same host
|
|
// execution, normally only a few milliseconds apart. This window is used
|
|
// only for the read projection and never changes persisted billing facts.
|
|
orphanLifecycleMatchWindow = 250 * time.Millisecond
|
|
// Near-equal candidates are deliberately left separate instead of risking a
|
|
// cross-request association under concurrent same-model traffic.
|
|
orphanLifecycleTieWindow = time.Millisecond
|
|
)
|
|
|
|
// SQLiteUsageRepository 使用单写连接保存事实，并用独立连接池承载只读查询。
|
|
type SQLiteUsageRepository struct {
|
|
db *sql.DB
|
|
readDB *sql.DB
|
|
}
|
|
|
|
func OpenSQLiteUsage(databasePath string) (*SQLiteUsageRepository, error) {
|
|
dsn, err := sqliteDSN(databasePath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
db, err := sql.Open("sqlite3", dsn)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("打开 SQLite: %w", err)
|
|
}
|
|
// 当前插件只需要一个同步 writer，避免并发写争用。
|
|
db.SetMaxOpenConns(1)
|
|
if err := db.Ping(); err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("连接 SQLite: %w", err)
|
|
}
|
|
if _, err := db.Exec(usageSchema); err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("初始化 SQLite schema: %w", err)
|
|
}
|
|
for _, migration := range []string{
|
|
`ALTER TABLE usage_records ADD COLUMN managed_key_id TEXT NOT NULL DEFAULT ''`,
|
|
`ALTER TABLE request_records ADD COLUMN managed_key_id TEXT NOT NULL DEFAULT ''`,
|
|
`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)
|
|
}
|
|
if _, err := db.Exec(`CREATE INDEX IF NOT EXISTS idx_request_records_managed_key ON request_records(managed_key_id, requested_at DESC)`); err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("创建 Key 请求索引: %w", err)
|
|
}
|
|
repository := &SQLiteUsageRepository{db: db}
|
|
if err := repository.ensureRequestDetailProjection(context.Background()); err != nil {
|
|
_ = db.Close()
|
|
return nil, err
|
|
}
|
|
readDB, err := sql.Open("sqlite3", dsn)
|
|
if err != nil {
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("打开 SQLite 只读连接池: %w", err)
|
|
}
|
|
readDB.SetMaxOpenConns(4)
|
|
readDB.SetMaxIdleConns(4)
|
|
if err := readDB.Ping(); err != nil {
|
|
_ = readDB.Close()
|
|
_ = db.Close()
|
|
return nil, fmt.Errorf("连接 SQLite 只读连接池: %w", err)
|
|
}
|
|
repository.readDB = readDB
|
|
return repository, nil
|
|
}
|
|
|
|
func sqliteDSN(databasePath string) (string, error) {
|
|
databasePath = strings.TrimSpace(databasePath)
|
|
if databasePath == "" {
|
|
return "", errors.New("database_path 不能为空")
|
|
}
|
|
if databasePath == ":memory:" {
|
|
return "file:cpa-ext-memory?mode=memory&cache=shared&_busy_timeout=5000&_foreign_keys=on", nil
|
|
}
|
|
if strings.HasPrefix(databasePath, "file:") {
|
|
separator := "?"
|
|
if strings.Contains(databasePath, "?") {
|
|
separator = "&"
|
|
}
|
|
return databasePath + separator + "_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil
|
|
}
|
|
absolutePath, err := filepath.Abs(databasePath)
|
|
if err != nil {
|
|
return "", fmt.Errorf("解析 database_path: %w", err)
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(absolutePath), 0o755); err != nil {
|
|
return "", fmt.Errorf("创建数据库目录: %w", err)
|
|
}
|
|
fileURL := (&url.URL{Scheme: "file", Path: filepath.ToSlash(absolutePath)}).String()
|
|
return fileURL + "?_busy_timeout=5000&_foreign_keys=on&_journal_mode=WAL&_synchronous=NORMAL", nil
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) Insert(ctx context.Context, record collection.Record) error {
|
|
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, 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, 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,
|
|
record.ExecutorType, record.RequestType, record.Endpoint, record.InputTokens, record.OutputTokens,
|
|
record.ReasoningTokens, record.CachedTokens, record.CacheReadTokens,
|
|
record.CacheWriteTokens, record.TotalTokens, int64(record.TTFT), int64(record.Latency),
|
|
record.CostMicros, record.PriceTier, record.FastRequested, record.FastPricingApplied,
|
|
record.PriceMultiplierNumerator, record.PriceMultiplierDenominator, record.ClientIP,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("写入用量记录: %w", err)
|
|
}
|
|
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
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) UpsertRequest(ctx context.Context, record collection.RequestRecord) error {
|
|
requestID := strings.TrimSpace(record.RequestID)
|
|
if requestID == "" {
|
|
return errors.New("request_id 不能为空")
|
|
}
|
|
requestedAt := record.RequestedAt
|
|
completedAt := record.CompletedAt
|
|
if requestedAt.IsZero() {
|
|
requestedAt = completedAt
|
|
}
|
|
if completedAt.IsZero() {
|
|
completedAt = requestedAt
|
|
}
|
|
latency := completedAt.Sub(requestedAt)
|
|
if latency < 0 {
|
|
latency = 0
|
|
}
|
|
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
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(request_id) DO UPDATE SET
|
|
managed_key_id=CASE WHEN excluded.managed_key_id <> '' THEN excluded.managed_key_id ELSE request_records.managed_key_id END,
|
|
trace_id=excluded.trace_id,
|
|
requested_at=excluded.requested_at,
|
|
completed_at=excluded.completed_at,
|
|
model=excluded.model,
|
|
requested_model=excluded.requested_model,
|
|
source_format=excluded.source_format,
|
|
stream=excluded.stream,
|
|
outcome=excluded.outcome,
|
|
status_code=excluded.status_code,
|
|
error=excluded.error,
|
|
endpoint=excluded.endpoint,
|
|
latency_ns=excluded.latency_ns`,
|
|
requestID, record.ManagedKeyID, record.TraceID, requestedAt.UTC().Format(time.RFC3339Nano),
|
|
completedAt.UTC().Format(time.RFC3339Nano), record.Model, record.RequestedModel,
|
|
record.SourceFormat, record.Stream, record.Outcome, record.StatusCode,
|
|
record.Error, record.Endpoint, int64(latency),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("写入请求终态: %w", err)
|
|
}
|
|
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
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) ListRecent(ctx context.Context, limit int) ([]collection.Record, error) {
|
|
if limit < 1 {
|
|
return []collection.Record{}, nil
|
|
}
|
|
// A logical request can occupy two source rows until projection time, so read
|
|
// ahead before applying the caller's logical-request limit.
|
|
scanLimit := limit * 2
|
|
if scanLimit < limit {
|
|
scanLimit = limit
|
|
}
|
|
rows, err := r.readDB.QueryContext(ctx, `
|
|
WITH combined AS (
|
|
SELECT
|
|
u.managed_key_id, u.request_id, u.execution_id, COALESCE(NULLIF(u.trace_id, ''), r.trace_id, '') AS trace_id,
|
|
u.requested_at, u.api_key, COALESCE(k.name, u.key_alias) AS key_alias,
|
|
u.auth_id, u.auth_index, u.auth_type, u.model, u.reasoning_effort, u.service_tier, u.speed,
|
|
CASE
|
|
WHEN u.failed = 1 THEN 1
|
|
WHEN r.outcome IN ('failed', 'rejected', 'canceled') THEN 1
|
|
ELSE 0
|
|
END AS failed,
|
|
u.executor_type,
|
|
CASE WHEN r.request_id IS NOT NULL THEN CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END ELSE u.request_type END AS request_type,
|
|
COALESCE(NULLIF(u.endpoint, ''), r.endpoint, '') AS endpoint,
|
|
u.input_tokens, u.output_tokens, u.reasoning_tokens, u.cached_tokens,
|
|
u.cache_read_tokens, u.cache_write_tokens, u.total_tokens,
|
|
u.ttft_ns, u.latency_ns, u.cost_micros, u.price_tier, u.fast_requested,
|
|
u.fast_pricing_applied, u.price_multiplier_numerator, u.price_multiplier_denominator,
|
|
u.client_ip, COALESCE(r.outcome, '') AS outcome, COALESCE(r.status_code, 0) AS status_code,
|
|
COALESCE(r.error, '') AS error
|
|
FROM usage_records u
|
|
LEFT JOIN request_records r ON r.request_id = u.request_id
|
|
LEFT JOIN managed_keys k ON k.id = u.managed_key_id
|
|
|
|
UNION ALL
|
|
|
|
SELECT
|
|
r.managed_key_id, r.request_id, '', r.trace_id, r.requested_at, '', COALESCE(k.name, '') AS key_alias,
|
|
'', '', '',
|
|
COALESCE(NULLIF(r.model, ''), r.requested_model), '', '', '',
|
|
CASE WHEN r.outcome = 'succeeded' THEN 0 ELSE 1 END,
|
|
'', CASE WHEN r.stream THEN 'SSE' ELSE 'JSON' END, r.endpoint,
|
|
0, 0, 0, 0, 0, 0, 0, 0, r.latency_ns, NULL, '', 0, 0, 1, 1, '',
|
|
r.outcome, r.status_code, r.error
|
|
FROM request_records r
|
|
LEFT JOIN managed_keys k ON k.id = r.managed_key_id
|
|
WHERE NOT EXISTS (SELECT 1 FROM usage_records u WHERE u.request_id = r.request_id)
|
|
)
|
|
SELECT managed_key_id, request_id, execution_id, trace_id, 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, outcome, status_code, error
|
|
FROM combined
|
|
ORDER BY requested_at DESC
|
|
LIMIT ?`, scanLimit)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("查询最近用量记录: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
records := make([]collection.Record, 0, scanLimit)
|
|
for rows.Next() {
|
|
var record collection.Record
|
|
var requestedAt string
|
|
var failed bool
|
|
var ttftNS, latencyNS int64
|
|
var cost sql.NullInt64
|
|
if err := rows.Scan(
|
|
&record.ManagedKeyID, &record.RequestID, &record.ExecutionID, &record.TraceID,
|
|
&requestedAt, &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, fmt.Errorf("读取用量记录: %w", err)
|
|
}
|
|
record.RequestedAt, err = time.Parse(time.RFC3339Nano, requestedAt)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("解析用量时间: %w", err)
|
|
}
|
|
record.Failed = failed
|
|
record.TTFT = time.Duration(ttftNS)
|
|
record.Latency = time.Duration(latencyNS)
|
|
if cost.Valid {
|
|
value := cost.Int64
|
|
record.CostMicros = &value
|
|
}
|
|
records = append(records, record)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, fmt.Errorf("遍历用量记录: %w", err)
|
|
}
|
|
records = mergeOrphanLifecycleRecords(records)
|
|
if len(records) > limit {
|
|
records = records[:limit]
|
|
}
|
|
return records, nil
|
|
}
|
|
|
|
type lifecycleMatch struct {
|
|
index int
|
|
distance time.Duration
|
|
ambiguous bool
|
|
}
|
|
|
|
func mergeOrphanLifecycleRecords(records []collection.Record) []collection.Record {
|
|
usageMatches := make(map[int]lifecycleMatch)
|
|
requestMatches := make(map[int]lifecycleMatch)
|
|
|
|
for usageIndex := range records {
|
|
if !isOrphanUsage(records[usageIndex]) {
|
|
continue
|
|
}
|
|
for requestIndex := range records {
|
|
if !isLifecycleOnly(records[requestIndex]) || !sameProjectionModel(records[usageIndex], records[requestIndex]) {
|
|
continue
|
|
}
|
|
distance := records[usageIndex].RequestedAt.Sub(records[requestIndex].RequestedAt)
|
|
if distance < 0 {
|
|
distance = -distance
|
|
}
|
|
if distance > orphanLifecycleMatchWindow {
|
|
continue
|
|
}
|
|
if current, exists := usageMatches[usageIndex]; exists {
|
|
usageMatches[usageIndex] = betterLifecycleMatch(current, requestIndex, distance)
|
|
} else {
|
|
usageMatches[usageIndex] = lifecycleMatch{index: requestIndex, distance: distance}
|
|
}
|
|
if current, exists := requestMatches[requestIndex]; exists {
|
|
requestMatches[requestIndex] = betterLifecycleMatch(current, usageIndex, distance)
|
|
} else {
|
|
requestMatches[requestIndex] = lifecycleMatch{index: usageIndex, distance: distance}
|
|
}
|
|
}
|
|
}
|
|
|
|
remove := make(map[int]struct{})
|
|
for usageIndex, requestMatch := range usageMatches {
|
|
if requestMatch.ambiguous {
|
|
continue
|
|
}
|
|
usageMatch, ok := requestMatches[requestMatch.index]
|
|
if !ok || usageMatch.ambiguous || usageMatch.index != usageIndex {
|
|
continue
|
|
}
|
|
records[usageIndex] = mergeLifecycleIntoUsage(records[usageIndex], records[requestMatch.index])
|
|
remove[requestMatch.index] = struct{}{}
|
|
}
|
|
|
|
if len(remove) == 0 {
|
|
return records
|
|
}
|
|
merged := make([]collection.Record, 0, len(records)-len(remove))
|
|
for index, record := range records {
|
|
if _, drop := remove[index]; !drop {
|
|
merged = append(merged, record)
|
|
}
|
|
}
|
|
return merged
|
|
}
|
|
|
|
func betterLifecycleMatch(current lifecycleMatch, candidate int, distance time.Duration) lifecycleMatch {
|
|
if distance+orphanLifecycleTieWindow < current.distance {
|
|
return lifecycleMatch{index: candidate, distance: distance}
|
|
}
|
|
if current.distance+orphanLifecycleTieWindow < distance {
|
|
return current
|
|
}
|
|
current.ambiguous = true
|
|
return current
|
|
}
|
|
|
|
func isOrphanUsage(record collection.Record) bool {
|
|
if strings.TrimSpace(record.RequestID) != "" {
|
|
return false
|
|
}
|
|
return record.CostMicros != nil || record.TotalTokens != 0 || record.InputTokens != 0 ||
|
|
record.OutputTokens != 0 || record.CacheReadTokens != 0 || record.CacheWriteTokens != 0 ||
|
|
strings.TrimSpace(record.APIKey) != "" || strings.TrimSpace(record.ExecutorType) != ""
|
|
}
|
|
|
|
func isLifecycleOnly(record collection.Record) bool {
|
|
return strings.TrimSpace(record.RequestID) != "" && strings.TrimSpace(record.ExecutionID) == "" &&
|
|
strings.TrimSpace(record.APIKey) == "" && record.CostMicros == nil && record.InputTokens == 0 &&
|
|
record.OutputTokens == 0 && record.CacheReadTokens == 0 && record.CacheWriteTokens == 0 &&
|
|
record.TotalTokens == 0 && strings.TrimSpace(record.Outcome) != ""
|
|
}
|
|
|
|
func sameProjectionModel(left, right collection.Record) bool {
|
|
if left.ManagedKeyID != "" && right.ManagedKeyID != "" && left.ManagedKeyID != right.ManagedKeyID {
|
|
return false
|
|
}
|
|
// The lifecycle callback keeps the client-facing alias while usage reports the
|
|
// resolved upstream model. A shared managed Key plus the mutual nearest-time
|
|
// match is therefore the stable identity when both sides are authenticated.
|
|
if left.ManagedKeyID != "" && left.ManagedKeyID == right.ManagedKeyID {
|
|
return true
|
|
}
|
|
leftModel := strings.ToLower(strings.TrimSpace(left.Model))
|
|
rightModel := strings.ToLower(strings.TrimSpace(right.Model))
|
|
return leftModel != "" && leftModel == rightModel
|
|
}
|
|
|
|
func mergeLifecycleIntoUsage(usage, lifecycle collection.Record) collection.Record {
|
|
if usage.ManagedKeyID == "" {
|
|
usage.ManagedKeyID = lifecycle.ManagedKeyID
|
|
}
|
|
usage.RequestID = lifecycle.RequestID
|
|
if usage.TraceID == "" {
|
|
usage.TraceID = lifecycle.TraceID
|
|
}
|
|
usage.Failed = usage.Failed || lifecycle.Failed
|
|
usage.RequestType = lifecycle.RequestType
|
|
usage.Endpoint = lifecycle.Endpoint
|
|
usage.Outcome = lifecycle.Outcome
|
|
usage.StatusCode = lifecycle.StatusCode
|
|
usage.Error = lifecycle.Error
|
|
if usage.Latency == 0 {
|
|
usage.Latency = lifecycle.Latency
|
|
}
|
|
return usage
|
|
}
|
|
|
|
func (r *SQLiteUsageRepository) Close() error {
|
|
var readErr error
|
|
if r.readDB != nil {
|
|
readErr = r.readDB.Close()
|
|
}
|
|
return errors.Join(readErr, r.db.Close())
|
|
}
|