hunjing/shenshi/internal/store/store.go

// Package store 神识自有库访问层：usage_logs / byok_keys / quota 余额。
//
// 双实现：PG（SHENSHI_DATABASE_URL，生产）与内存（开发/测试，未配置时回退）。
// 语义契约由 store_test.go 在内存实现上钉死；PG 实现为同接口直写 SQL。
package store

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"

	// PG 驱动（database/sql 标准接口）
	_ "github.com/jackc/pgx/v5/stdlib"
)

// UsageLog 一次端点尝试的完整计量记录（shenshi-gateway.md §3 usage_logs）。
type UsageLog struct {
	ID               string // 尝试唯一标识（run_id#endpoint#seq，配额结算幂等键）
	RunID            string
	CharacterID      string
	UserID           string
	Model            string
	Slot             string
	Endpoint         string
	Status           string // ok / error / first_byte_timeout
	ErrorCode        string
	FirstTokenMS     int64
	TotalMS          int64
	PromptTokens     int
	CompletionTokens int
	At               time.Time
}

// UsageAgg 用量聚合行（/me/usage）。
type UsageAgg struct {
	UserID           string `json:"user_id"`
	Model            string `json:"model"`
	Endpoint         string `json:"endpoint"`
	Requests         int    `json:"requests"`
	OK               int    `json:"ok"`
	PromptTokens     int    `json:"prompt_tokens"`
	CompletionTokens int    `json:"completion_tokens"`
}

// ByokKey 用户自带密钥（API 永远脱敏返回）。
type ByokKey struct {
	ID        string
	UserID    string
	Name      string
	Type      string // openai_compatible / anthropic / gemini（契约词典 §14.2）
	BaseURL   string
	APIKey    string // 加密存储；出站时解密
	Model     string
	CreatedAt time.Time
	LastUsed  *time.Time
	RevokedAt *time.Time
}

// QuotaBalance 用户配额账（三段式状态：Reserved 预扣 / Used 已结算）。
type QuotaBalance struct {
	UserID   string
	Reserved int64
	Used     int64
}

// ErrByokNotFound BYOK 密钥不存在或已撤回。
var ErrByokNotFound = errors.New("BYOK 密钥不存在或已撤回")

// Store 存储接口（全异步语义由调用方保证；实现自身须并发安全）。
type Store interface {
	InsertUsage(log UsageLog) error
	AggregateUsage(userID string) ([]UsageAgg, error)
	SampleUsage(limit int) ([]UsageLog, error) // 对账抽样：最近 N 条原始记录

	UpsertByok(key ByokKey) error
	ListByok(userID string) ([]ByokKey, error)
	GetByok(userID, id string) (*ByokKey, error)
	RevokeByok(userID, id string) error
	TouchByok(id string, at time.Time) error

	// Quota 三段式原子操作：Reserve 预扣（不足返回错误）；Settle 结算差值；Refund 幂等退款。
	Reserve(userID string, amount int64) error
	Settle(userID string, reserved, actual int64) error
	Refund(userID string, reserved int64) error
	Balance(userID string) (QuotaBalance, error)
}

// Open 按配置打开存储：DSN 非空时 PG（自动建表，含就绪重试），否则内存实现。
// 重试窗口 60 秒：容器编排下 PG 可能晚于神识就绪（跨 compose 文件无 depends_on）。
func Open(dsn string) (Store, error) {
	if dsn == "" {
		return NewMem(), nil
	}
	db, err := sql.Open("pgx", dsn)
	if err != nil {
		return nil, fmt.Errorf("打开神识库失败: %w", err)
	}
	deadline := time.Now().Add(60 * time.Second)
	for {
		err = migrate(db)
		if err == nil {
			return &pgStore{db: db}, nil
		}
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("神识库迁移失败（60 秒重试后放弃）: %w", err)
		}
		time.Sleep(2 * time.Second)
	}
}

func migrate(db *sql.DB) error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS usage_logs (
			id TEXT PRIMARY KEY,
			run_id TEXT NOT NULL,
			character_id TEXT NOT NULL DEFAULT '',
			user_id TEXT NOT NULL DEFAULT '',
			model TEXT NOT NULL,
			slot TEXT NOT NULL,
			endpoint TEXT NOT NULL,
			status TEXT NOT NULL,
			error_code TEXT NOT NULL DEFAULT '',
			first_token_ms BIGINT NOT NULL DEFAULT 0,
			total_ms BIGINT NOT NULL DEFAULT 0,
			prompt_tokens INT NOT NULL DEFAULT 0,
			completion_tokens INT NOT NULL DEFAULT 0,
			at TIMESTAMPTZ NOT NULL
		)`,
		`CREATE INDEX IF NOT EXISTS idx_usage_logs_user ON usage_logs(user_id, at DESC)`,
		`CREATE TABLE IF NOT EXISTS byok_keys (
			id TEXT PRIMARY KEY,
			user_id TEXT NOT NULL,
			name TEXT NOT NULL,
			type TEXT NOT NULL,
			base_url TEXT NOT NULL DEFAULT '',
			api_key TEXT NOT NULL,
			model TEXT NOT NULL DEFAULT '',
			created_at TIMESTAMPTZ NOT NULL,
			last_used TIMESTAMPTZ,
			revoked_at TIMESTAMPTZ
		)`,
		`CREATE INDEX IF NOT EXISTS idx_byok_user ON byok_keys(user_id)`,
		`CREATE TABLE IF NOT EXISTS quota_balances (
			user_id TEXT PRIMARY KEY,
			reserved BIGINT NOT NULL DEFAULT 0,
			used BIGINT NOT NULL DEFAULT 0
		)`,
	}
	for _, s := range stmts {
		if _, err := db.Exec(s); err != nil {
			return fmt.Errorf("建表失败: %w", err)
		}
	}
	return nil
}

// ── 内存实现 ──────────────────────────────────────────────

type memStore struct {
	mu    sync.Mutex
	usage []UsageLog
	byok  map[string][]ByokKey // user → keys
	quota map[string]*QuotaBalance
}

// NewMem 内存存储（保留最近 10 万条 usage，防无界增长）。
func NewMem() Store {
	return &memStore{byok: map[string][]ByokKey{}, quota: map[string]*QuotaBalance{}}
}

const memUsageLimit = 100_000

func (m *memStore) InsertUsage(log UsageLog) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if len(m.usage) >= memUsageLimit {
		m.usage = m.usage[1:]
	}
	m.usage = append(m.usage, log)
	return nil
}

func (m *memStore) AggregateUsage(userID string) ([]UsageAgg, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	agg := map[string]*UsageAgg{}
	for _, e := range m.usage {
		if userID != "" && e.UserID != userID {
			continue
		}
		key := e.UserID + "|" + e.Model + "|" + e.Endpoint
		row := agg[key]
		if row == nil {
			row = &UsageAgg{UserID: e.UserID, Model: e.Model, Endpoint: e.Endpoint}
			agg[key] = row
		}
		row.Requests++
		if e.Status == "ok" {
			row.OK++
			row.PromptTokens += e.PromptTokens
			row.CompletionTokens += e.CompletionTokens
		}
	}
	out := make([]UsageAgg, 0, len(agg))
	for _, row := range agg {
		out = append(out, *row)
	}
	return out, nil
}

func (m *memStore) SampleUsage(limit int) ([]UsageLog, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	n := limit
	if n > len(m.usage) {
		n = len(m.usage)
	}
	out := make([]UsageLog, n)
	copy(out, m.usage[len(m.usage)-n:])
	return out, nil
}

func (m *memStore) UpsertByok(key ByokKey) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	keys := m.byok[key.UserID]
	for i, k := range keys {
		if k.ID == key.ID {
			keys[i] = key
			return nil
		}
	}
	m.byok[key.UserID] = append(keys, key)
	return nil
}

func (m *memStore) ListByok(userID string) ([]ByokKey, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make([]ByokKey, 0, len(m.byok[userID]))
	for _, k := range m.byok[userID] {
		if k.RevokedAt == nil {
			out = append(out, k)
		}
	}
	return out, nil
}

func (m *memStore) GetByok(userID, id string) (*ByokKey, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, k := range m.byok[userID] {
		if k.ID == id && k.RevokedAt == nil {
			out := k
			return &out, nil
		}
	}
	return nil, nil
}

func (m *memStore) RevokeByok(userID, id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	now := time.Now().UTC()
	keys := m.byok[userID]
	for i, k := range keys {
		if k.ID == id && k.RevokedAt == nil {
			keys[i].RevokedAt = &now
			return nil
		}
	}
	return ErrByokNotFound
}

func (m *memStore) TouchByok(id string, at time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	for user, keys := range m.byok {
		for i, k := range keys {
			if k.ID == id {
				used := at
				keys[i].LastUsed = &used
				m.byok[user] = keys
				return nil
			}
		}
	}
	return nil
}

func (m *memStore) Reserve(userID string, amount int64) error {
	if amount <= 0 {
		return nil
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	b := m.quotaGet(userID)
	b.Reserved += amount
	return nil
}

func (m *memStore) Settle(userID string, reserved, actual int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	b := m.quotaGet(userID)
	b.Reserved -= reserved
	if b.Reserved < 0 {
		b.Reserved = 0
	}
	b.Used += actual
	return nil
}

func (m *memStore) Refund(userID string, reserved int64) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	b := m.quotaGet(userID)
	b.Reserved -= reserved
	if b.Reserved < 0 {
		b.Reserved = 0
	}
	return nil
}

func (m *memStore) Balance(userID string) (QuotaBalance, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	if b, ok := m.quota[userID]; ok {
		return *b, nil
	}
	return QuotaBalance{UserID: userID}, nil
}

func (m *memStore) quotaGet(userID string) *QuotaBalance {
	b, ok := m.quota[userID]
	if !ok {
		b = &QuotaBalance{UserID: userID}
		m.quota[userID] = b
	}
	return b
}