AskaEth
be38651ba4
- model 层新增可选 UsageSink（nil ⇒ 逐字不变，既有用例零改动） - 组合根注入 UsageWriter：有界队列 1024 + 单写者 + 批量事务（128/批、2s 兜底）， 队列满丢弃 + 限流 WARN，绝不阻塞模型调用；sink 挂 RuntimeLLM（热重载不断） - 迁移 0017：表 llm_usage（13 列，session_id 可空）+ idx_created_at + idx_session_id - 新端点 GET /api/v1/admin/llm/usage（owner|admin 读；缺省近 24h、上限 90 天、 超限 400 不静默裁剪；hitRate = cacheHit/(cacheHit+cacheMiss)，分母 0 ⇒ 0； 绝不含凭据与正文；未接线/查库失败 ⇒ 500 usage_unavailable） - LLM_USAGE_RETENTION（默认 720h，≤0 = 关闭清理）：启动清一次 + 每小时 - 面板 LLM 页「用量（近 24h）」小卡 + Docs 两份同步 v0.64 - 已知缺口（如实登记）：session_id 恒 NULL，需 loop 侧在请求上打标，留列后接
321 lines
12 KiB
Go
321 lines
12 KiB
Go
package gateway
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"encoding/json"
|
||
"net/http"
|
||
"path/filepath"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/auth"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
|
||
)
|
||
|
||
// 本文件钉住 LLM 用量汇总只读端点（GET /api/v1/admin/llm/usage，切片 U1，
|
||
// 契约 v0.64）：数值正确（含 hitRate 公式）、时间窗与用途过滤生效、
|
||
// 超上限窗口 400、鉴权（401 / viewer 403）、不回显凭据。
|
||
|
||
// llmUsageTotalsDTO 与 internal/services/gateway/admin_llm_usage.go 的字段逐字
|
||
// 对齐；字段名写错就解不出来。
|
||
type llmUsageTotalsDTO struct {
|
||
Calls int64 `json:"calls"`
|
||
Prompt int64 `json:"prompt"`
|
||
Completion int64 `json:"completion"`
|
||
CacheHit int64 `json:"cacheHit"`
|
||
CacheMiss int64 `json:"cacheMiss"`
|
||
HitRate float64 `json:"hitRate"`
|
||
}
|
||
|
||
type llmUsageResponseDTO struct {
|
||
Window struct {
|
||
Since string `json:"since"`
|
||
Until string `json:"until"`
|
||
Purpose string `json:"purpose"`
|
||
} `json:"window"`
|
||
Totals llmUsageTotalsDTO `json:"totals"`
|
||
ByPurpose []struct {
|
||
Purpose string `json:"purpose"`
|
||
llmUsageTotalsDTO
|
||
} `json:"byPurpose"`
|
||
ByModel []struct {
|
||
Model string `json:"model"`
|
||
llmUsageTotalsDTO
|
||
} `json:"byModel"`
|
||
}
|
||
|
||
// llmUsageTestDB 在 dir 下打开一个真实 SQLite 文件库并跑迁移（端点用的就是这条
|
||
// 活连接：Deps.Usage）。
|
||
func llmUsageTestDB(t *testing.T, dir string) *sql.DB {
|
||
t.Helper()
|
||
ctx := context.Background()
|
||
conn, err := db.Open(ctx, filepath.Join(dir, "cyrene.db"), gwTestLogger())
|
||
if err != nil {
|
||
t.Fatalf("open db: %v", err)
|
||
}
|
||
t.Cleanup(func() { _ = conn.Close() })
|
||
if err := db.MigrateUp(ctx, conn, gwTestLogger()); err != nil {
|
||
t.Fatalf("migrate: %v", err)
|
||
}
|
||
return conn
|
||
}
|
||
|
||
// seedUsageSamples 用**真实的写入器**落库三条样本（批量事务路径），使端到端
|
||
// 「写入 → 汇总」与生产同一条路。
|
||
func seedUsageSamples(t *testing.T, conn *sql.DB) {
|
||
t.Helper()
|
||
writer := model.NewUsageWriter(conn, gwTestLogger())
|
||
writer.Start(context.Background())
|
||
samples := []model.UsageSample{
|
||
{Adapter: "openai", Model: "deepseek-chat", Purpose: "chat",
|
||
PromptTokens: 1000, CompletionTokens: 200, TotalTokens: 1200,
|
||
CacheHitTokens: 900, CacheMissTokens: 100, TookMs: 50},
|
||
{Adapter: "openai", Model: "deepseek-chat", Purpose: "chat",
|
||
PromptTokens: 1000, CompletionTokens: 100, TotalTokens: 1100,
|
||
CacheHitTokens: 0, CacheMissTokens: 1000, TookMs: 70},
|
||
{Adapter: "openai", Model: "gpt-x", Purpose: "memory_extract",
|
||
PromptTokens: 500, CompletionTokens: 50, TotalTokens: 550,
|
||
CacheHitTokens: 500, CacheMissTokens: 0, TookMs: 30},
|
||
}
|
||
for _, sample := range samples {
|
||
writer.RecordUsage(sample)
|
||
}
|
||
writer.Stop()
|
||
}
|
||
|
||
// insertOldUsageRow 写一行**窗口之外**的样本（created_at 显式给定）。
|
||
func insertOldUsageRow(t *testing.T, conn *sql.DB, purpose, model, createdAt string) {
|
||
t.Helper()
|
||
if _, err := conn.ExecContext(context.Background(),
|
||
`INSERT INTO llm_usage
|
||
(purpose, adapter, model, prompt_tokens, completion_tokens,
|
||
cache_hit_tokens, cache_miss_tokens, total_tokens, took_ms, created_at)
|
||
VALUES (?, 'openai', ?, 7, 1, 0, 7, 8, 5, ?)`,
|
||
purpose, model, createdAt); err != nil {
|
||
t.Fatalf("insert old usage row: %v", err)
|
||
}
|
||
}
|
||
|
||
// getLlmUsage 以 owner token 取一次用量汇总。
|
||
func getLlmUsage(t *testing.T, url, query, token string) (int, []byte) {
|
||
t.Helper()
|
||
path := "/api/v1/admin/llm/usage"
|
||
if query != "" {
|
||
path += "?" + query
|
||
}
|
||
resp, body := doReq(t, http.MethodGet, url+path, nil, token)
|
||
return resp.StatusCode, body
|
||
}
|
||
|
||
// decodeLlmUsage 解出 data 本体。
|
||
func decodeLlmUsage(t *testing.T, body []byte) llmUsageResponseDTO {
|
||
t.Helper()
|
||
var out llmUsageResponseDTO
|
||
if err := json.Unmarshal(decodeBody(t, body)["data"], &out); err != nil {
|
||
t.Fatalf("decode usage data: %v (body=%s)", err, body)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// TestAdminLlmUsageSummarizesBudgetWindow 钉住默认窗口（近 24h）下的数值与
|
||
// hitRate 公式：窗口外的行**不算进来**，命中率 = cacheHit / (cacheHit + cacheMiss)。
|
||
func TestAdminLlmUsageSummarizesBudgetWindow(t *testing.T) {
|
||
t.Parallel()
|
||
conn := llmUsageTestDB(t, t.TempDir())
|
||
seedUsageSamples(t, conn)
|
||
// 窗口外（2020 年）的一行：默认近 24h 不该看到它。
|
||
insertOldUsageRow(t, conn, "chat", "ancient-model", "2020-01-01T00:00:00.000Z")
|
||
|
||
_, ts, jwt := buildEnv(t, cfgDefault(), Deps{Usage: model.NewUsageReader(conn)})
|
||
status, body := getLlmUsage(t, ts.URL, "", signAny(t, jwt, "acct", auth.RoleOwner))
|
||
if status != http.StatusOK {
|
||
t.Fatalf("status = %d, want 200 (body=%s)", status, body)
|
||
}
|
||
got := decodeLlmUsage(t, body)
|
||
|
||
want := llmUsageTotalsDTO{Calls: 3, Prompt: 2500, Completion: 350, CacheHit: 1400, CacheMiss: 1100, HitRate: 0.56}
|
||
if got.Totals != want {
|
||
t.Fatalf("totals = %+v, want %+v", got.Totals, want)
|
||
}
|
||
if len(got.ByPurpose) != 2 || got.ByPurpose[0].Purpose != "chat" || got.ByPurpose[0].Calls != 2 {
|
||
t.Fatalf("byPurpose = %+v, want chat 在前且 2 次调用", got.ByPurpose)
|
||
}
|
||
if got.ByPurpose[0].HitRate != 0.45 {
|
||
t.Fatalf("byPurpose[chat].hitRate = %v, want 0.45", got.ByPurpose[0].HitRate)
|
||
}
|
||
if len(got.ByModel) != 2 || got.ByModel[0].Model != "deepseek-chat" {
|
||
t.Fatalf("byModel = %+v, want deepseek-chat 在前", got.ByModel)
|
||
}
|
||
// 窗口回显必须是**补齐后的真值**（缺省 = 近 24h），且是 UTC RFC3339。
|
||
since, err := time.Parse(time.RFC3339, got.Window.Since)
|
||
if err != nil {
|
||
t.Fatalf("window.since %q 不是 RFC3339: %v", got.Window.Since, err)
|
||
}
|
||
until, err := time.Parse(time.RFC3339, got.Window.Until)
|
||
if err != nil {
|
||
t.Fatalf("window.until %q 不是 RFC3339: %v", got.Window.Until, err)
|
||
}
|
||
if span := until.Sub(since); span != model.UsageDefaultWindow {
|
||
t.Fatalf("缺省窗口 = %s, want %s", span, model.UsageDefaultWindow)
|
||
}
|
||
if got.Window.Purpose != "" {
|
||
t.Fatalf("window.purpose = %q, want 空（未过滤）", got.Window.Purpose)
|
||
}
|
||
}
|
||
|
||
// TestAdminLlmUsageHonoursWindowAndPurposeParams 钉住 since/until/purpose 真的
|
||
// 生效（用一个把 2020 那条括进来的显式窗口）。
|
||
func TestAdminLlmUsageHonoursWindowAndPurposeParams(t *testing.T) {
|
||
t.Parallel()
|
||
conn := llmUsageTestDB(t, t.TempDir())
|
||
seedUsageSamples(t, conn)
|
||
insertOldUsageRow(t, conn, "memory_extract", "ancient-model", "2020-01-01T00:00:00.000Z")
|
||
|
||
_, ts, jwt := buildEnv(t, cfgDefault(), Deps{Usage: model.NewUsageReader(conn)})
|
||
token := signAny(t, jwt, "acct", auth.RoleOwner)
|
||
|
||
status, body := getLlmUsage(t, ts.URL,
|
||
"since=2019-12-31T00:00:00Z&until=2020-01-02T00:00:00Z", token)
|
||
if status != http.StatusOK {
|
||
t.Fatalf("status = %d, want 200 (body=%s)", status, body)
|
||
}
|
||
old := decodeLlmUsage(t, body)
|
||
if old.Totals.Calls != 1 || old.Totals.Prompt != 7 {
|
||
t.Fatalf("历史窗口 totals = %+v, want 只剩 2020 那条（calls 1 / prompt 7）", old.Totals)
|
||
}
|
||
|
||
// 同一窗口 + purpose 不匹配 ⇒ 空（但仍 200，不是 404）。
|
||
status, body = getLlmUsage(t, ts.URL,
|
||
"since=2019-12-31T00:00:00Z&until=2020-01-02T00:00:00Z&purpose=chat", token)
|
||
if status != http.StatusOK {
|
||
t.Fatalf("status = %d, want 200 (body=%s)", status, body)
|
||
}
|
||
filtered := decodeLlmUsage(t, body)
|
||
if filtered.Totals.Calls != 0 || len(filtered.ByPurpose) != 0 || len(filtered.ByModel) != 0 {
|
||
t.Fatalf("purpose 过滤未生效：%+v", filtered)
|
||
}
|
||
// 分组列表恒为数组（不是 null）——客户端 JSON 解析不该踩空。
|
||
if raw := string(decodeBody(t, body)["data"]); strings.Contains(raw, `"byPurpose":null`) {
|
||
t.Fatalf("byPurpose 下发成了 null：%s", raw)
|
||
}
|
||
|
||
// 毫秒时间戳与 RFC3339 同义（客户端可复用消息游标的形态）。
|
||
status, body = getLlmUsage(t, ts.URL, "since=1577750400000&until=1577923200000", token)
|
||
if status != http.StatusOK {
|
||
t.Fatalf("毫秒窗口 status = %d, want 200 (body=%s)", status, body)
|
||
}
|
||
if ms := decodeLlmUsage(t, body); ms.Totals.Calls != 1 {
|
||
t.Fatalf("毫秒窗口 totals = %+v, want calls 1（与 RFC3339 同义）", ms.Totals)
|
||
}
|
||
}
|
||
|
||
// TestAdminLlmUsageRejectsBadWindows 钉住窗口校验：形状错、区间反了、跨度超上限
|
||
// 一律 400 + bad_request，**不静默裁剪**。
|
||
func TestAdminLlmUsageRejectsBadWindows(t *testing.T) {
|
||
t.Parallel()
|
||
conn := llmUsageTestDB(t, t.TempDir())
|
||
_, ts, jwt := buildEnv(t, cfgDefault(), Deps{Usage: model.NewUsageReader(conn)})
|
||
token := signAny(t, jwt, "acct", auth.RoleOwner)
|
||
|
||
cases := map[string]string{
|
||
"since 非法格式": "since=昨天",
|
||
"until 非法格式": "until=not-a-time",
|
||
"since 不早于 until": "since=2026-10-05T00:00:00Z&until=2026-10-05T00:00:00Z",
|
||
"窗口超 90 天上限": "since=2026-01-01T00:00:00Z&until=2026-10-05T00:00:00Z",
|
||
}
|
||
for name, query := range cases {
|
||
status, body := getLlmUsage(t, ts.URL, query, token)
|
||
if status != http.StatusBadRequest {
|
||
t.Fatalf("%s: status = %d, want 400 (body=%s)", name, status, body)
|
||
}
|
||
if code := errorCode(t, body); code != "bad_request" {
|
||
t.Fatalf("%s: code = %q, want bad_request", name, code)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestAdminLlmUsageAuthAndRole 钉住角色门：无凭据 401；viewer 403（同一道
|
||
// owner|admin 门，与 /admin/llm 一致）。
|
||
func TestAdminLlmUsageAuthAndRole(t *testing.T) {
|
||
t.Parallel()
|
||
conn := llmUsageTestDB(t, t.TempDir())
|
||
_, ts, jwt := buildEnv(t, cfgDefault(), Deps{Usage: model.NewUsageReader(conn)})
|
||
|
||
status, body := getLlmUsage(t, ts.URL, "", "")
|
||
if status != http.StatusUnauthorized {
|
||
t.Fatalf("无凭据 status = %d, want 401 (body=%s)", status, body)
|
||
}
|
||
wantErrorShape(t, body)
|
||
|
||
status, body = getLlmUsage(t, ts.URL, "", signAny(t, jwt, "watch", auth.RoleViewer))
|
||
if status != http.StatusForbidden {
|
||
t.Fatalf("viewer status = %d, want 403 (body=%s)", status, body)
|
||
}
|
||
wantErrorShape(t, body)
|
||
|
||
// admin 与 owner 一样放行。
|
||
status, _ = getLlmUsage(t, ts.URL, "", signAny(t, jwt, "op", auth.RoleAdmin))
|
||
if status != http.StatusOK {
|
||
t.Fatalf("admin status = %d, want 200", status)
|
||
}
|
||
}
|
||
|
||
// TestAdminLlmUsageHidesCredentials 钉住红线：响应体里既没有凭据、也没有任何
|
||
// 「顺手带出来」的字段 —— `totals` 的键集恰好是契约里的那六个。
|
||
func TestAdminLlmUsageHidesCredentials(t *testing.T) {
|
||
t.Parallel()
|
||
conn := llmUsageTestDB(t, t.TempDir())
|
||
seedUsageSamples(t, conn)
|
||
cfg := cfgDefault()
|
||
cfg.DB.PGURL = "postgres://cyrene_user:s3cret-pw@db.internal:5432/cyrene"
|
||
_, ts, jwt := buildEnv(t, cfg, Deps{Usage: model.NewUsageReader(conn)})
|
||
|
||
status, body := getLlmUsage(t, ts.URL, "", signAny(t, jwt, "acct", auth.RoleOwner))
|
||
if status != http.StatusOK {
|
||
t.Fatalf("status = %d, want 200 (body=%s)", status, body)
|
||
}
|
||
for _, secret := range []string{
|
||
"postgres://", "s3cret-pw", "cyrene_user", "db.internal",
|
||
"apiKey", "api_key", "DSN", "dsn", "Bearer", "sk-",
|
||
} {
|
||
if strings.Contains(string(body), secret) {
|
||
t.Fatalf("响应体泄露/夹带了 %q：%s", secret, body)
|
||
}
|
||
}
|
||
|
||
var envelope struct {
|
||
Data struct {
|
||
Totals map[string]json.RawMessage `json:"totals"`
|
||
} `json:"data"`
|
||
}
|
||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||
t.Fatalf("decode envelope: %v", err)
|
||
}
|
||
wantKeys := []string{"calls", "prompt", "completion", "cacheHit", "cacheMiss", "hitRate"}
|
||
if len(envelope.Data.Totals) != len(wantKeys) {
|
||
t.Fatalf("totals 键集 = %v, want 恰好 %v", envelope.Data.Totals, wantKeys)
|
||
}
|
||
for _, key := range wantKeys {
|
||
if _, ok := envelope.Data.Totals[key]; !ok {
|
||
t.Fatalf("totals 缺字段 %q（契约 v0.64）：%v", key, envelope.Data.Totals)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestAdminLlmUsageUnwired 钉住「没接线不假装成功」：未注入汇总出口时 500 +
|
||
// usage_unavailable，绝不回一份全 0 的假汇总。
|
||
func TestAdminLlmUsageUnwired(t *testing.T) {
|
||
t.Parallel()
|
||
_, ts, jwt := buildEnv(t, cfgDefault(), Deps{})
|
||
status, body := getLlmUsage(t, ts.URL, "", signAny(t, jwt, "acct", auth.RoleOwner))
|
||
if status != http.StatusInternalServerError {
|
||
t.Fatalf("status = %d, want 500 (body=%s)", status, body)
|
||
}
|
||
if code := errorCode(t, body); code != "usage_unavailable" {
|
||
t.Fatalf("code = %q, want usage_unavailable", code)
|
||
}
|
||
}
|