AskaEth
59a1ac4056
通道：新增 model/session_ctx.go 的 WithSessionID / SessionIDFrom（与 tool.WithTurnPersona
同族的 ctx 范式，避开 model 反向 import loop 的环），loop 在 Run 顶部播种；适配器
record(ctx,...) 接会话；usageSessionArg 空串映射为 NULL。
兼容：sink 未注入或会话未播种时行为逐字不变；既有 model 测试零断言改动即全绿。
文件：model/{session_ctx.go(新), usage.go, usage_store.go, openai.go, anthropic.go,
usage_session_test.go(新)} + loop/{run.go, usage_session_test.go(新)}。
变异三处：去 loop 打标 / 写侧不绑会话 / 适配器不读 ctx ⇒ 各自红行，还原复绿。
契约与日志字段集未动；读侧不按 session 过滤，面板汇总口径不变。
352 lines
14 KiB
Go
352 lines
14 KiB
Go
package model
|
||
|
||
import (
|
||
"context"
|
||
"database/sql"
|
||
"net/http"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||
"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
|
||
)
|
||
|
||
// 本文件钉住 LLM 用量**落库**（切片 U1）：
|
||
// - sink 为 nil ⇒ 行为逐字不变（只记日志）；
|
||
// - sink 非 nil ⇒ 样本抵达写入器，字段与日志同源；
|
||
// - 队列满 ⇒ 丢弃 + 限流 WARN，**绝不阻塞**模型调用路径；
|
||
// - 批量落库后读回，命中率按 cacheHit / (cacheHit + cacheMiss) 算出。
|
||
|
||
// usageSinkProbeTimeout 是「RecordUsage 不许阻塞」的判定预算：真实现是微秒级，
|
||
// 这里给 5 秒是为了让判定只对「真的在等」敏感，不对机器忙敏感。
|
||
const usageSinkProbeTimeout = 5 * time.Second
|
||
|
||
// usageTestDB 打开一个带全部迁移的内存库（`:memory:` 每个句柄一份私有库，
|
||
// 并行测试天然隔离）。
|
||
func usageTestDB(t *testing.T) *sql.DB {
|
||
t.Helper()
|
||
conn, err := db.Open(context.Background(), ":memory:", discardLogger("model/usage_test"))
|
||
if err != nil {
|
||
t.Fatalf("open in-memory db: %v", err)
|
||
}
|
||
t.Cleanup(func() { _ = conn.Close() })
|
||
if err := db.MigrateUp(context.Background(), conn, discardLogger("model/usage_test")); err != nil {
|
||
t.Fatalf("migrate: %v", err)
|
||
}
|
||
return conn
|
||
}
|
||
|
||
// fakeUsageSink 捕获被投递的样本，用于断言 sink ↔ logUsage 调用点的字段一致。
|
||
type fakeUsageSink struct {
|
||
samples []UsageSample
|
||
}
|
||
|
||
func (f *fakeUsageSink) RecordUsage(sample UsageSample) { f.samples = append(f.samples, sample) }
|
||
|
||
// TestUsageSink_NilKeepsLogOnlyBehavior 钉住「不注入 sink ⇒ 行为逐字不变」：
|
||
// 适配器照常完成调用，且**恰好**记一条结构化用量日志。
|
||
func TestUsageSink_NilKeepsLogOnlyBehavior(t *testing.T) {
|
||
t.Parallel()
|
||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil)
|
||
log, buf := captureJSONLog()
|
||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}, WithOpenAILogger(log))
|
||
|
||
resp, err := ad.Chat(context.Background(), Request{
|
||
Purpose: "chat",
|
||
Messages: []Message{TextMessage(RoleUser, "hi")},
|
||
}, Options{})
|
||
if err != nil {
|
||
t.Fatalf("Chat: %v", err)
|
||
}
|
||
if resp.Content != "Hello!" {
|
||
t.Fatalf("正文受落库改动影响: %q", resp.Content)
|
||
}
|
||
// 没有 sink 时也恰好一条用量日志（多一条或少一条都会在这里变红）。
|
||
rec := usageLogRecord(t, buf.String())
|
||
if rec["purpose"] != "chat" || rec["prompt"] == nil || rec["cacheMiss"] == nil {
|
||
t.Fatalf("用量日志字段缺失：%v", rec)
|
||
}
|
||
}
|
||
|
||
// TestUsageSink_ReceivesSameFieldsAsLog 钉住 sink ↔ 既有 logUsage 调用点的成对
|
||
// 契约：同一个样本既进日志、也进 sink，字段逐项一致（两处口径不分叉）。
|
||
func TestUsageSink_ReceivesSameFieldsAsLog(t *testing.T) {
|
||
t.Parallel()
|
||
body := `{
|
||
"id": "cmpl-1", "model": "deepseek-chat",
|
||
"choices": [{"index":0,"message":{"role":"assistant","content":"你好"},"finish_reason":"stop"}],
|
||
"usage": {"prompt_tokens": 1000, "completion_tokens": 200, "total_tokens": 1200,
|
||
"prompt_cache_hit_tokens": 900, "prompt_cache_miss_tokens": 100,
|
||
"completion_tokens_details": {"reasoning_tokens": 150}}
|
||
}`
|
||
srv := newCompletionServer(t, http.StatusOK, "application/json", body, nil)
|
||
sink := &fakeUsageSink{}
|
||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "deepseek-chat"}, WithUsageSink(sink))
|
||
|
||
if _, err := ad.Chat(context.Background(), Request{
|
||
Purpose: "chat",
|
||
Messages: []Message{TextMessage(RoleUser, "hi")},
|
||
}, Options{}); err != nil {
|
||
t.Fatalf("Chat: %v", err)
|
||
}
|
||
if len(sink.samples) != 1 {
|
||
t.Fatalf("sink 收到 %d 条样本, want 1", len(sink.samples))
|
||
}
|
||
got := sink.samples[0]
|
||
want := UsageSample{
|
||
Adapter: ad.Name(),
|
||
Model: "deepseek-chat",
|
||
Purpose: "chat",
|
||
PromptTokens: 1000,
|
||
CompletionTokens: 200,
|
||
TotalTokens: 1200,
|
||
CacheHitTokens: 900,
|
||
CacheMissTokens: 100,
|
||
ReasoningTokens: 150,
|
||
}
|
||
got.TookMs = 0 // 耗时不做等值断言（它是真实的墙钟量）
|
||
if got != want {
|
||
t.Fatalf("sink 样本 = %+v, want %+v", got, want)
|
||
}
|
||
}
|
||
|
||
// TestUsageWriter_QueueFullDropsWithoutBlocking 钉住「写入绝不阻塞模型调用」：
|
||
// 队列满时再投递只丢弃（并限流 WARN），投递方**不等待**。
|
||
//
|
||
// 写入器刻意不 Start（没有消费者），队列因此停在满的状态 —— 这正是真实部署里
|
||
// 库卡住/变慢时的形态。
|
||
func TestUsageWriter_QueueFullDropsWithoutBlocking(t *testing.T) {
|
||
t.Parallel()
|
||
conn := usageTestDB(t)
|
||
log, buf := captureJSONLog()
|
||
writer := NewUsageWriter(conn, log)
|
||
if writer == nil {
|
||
t.Fatal("NewUsageWriter 在接了库时不该返回 nil")
|
||
}
|
||
sample := UsageSample{Adapter: "openai", Model: "m", Purpose: "chat", PromptTokens: 1}
|
||
|
||
for i := 0; i < usageQueueCapacity; i++ {
|
||
writer.RecordUsage(sample)
|
||
}
|
||
if got := len(writer.queue); got != usageQueueCapacity {
|
||
t.Fatalf("队列长度 = %d, want 满 %d", got, usageQueueCapacity)
|
||
}
|
||
|
||
const overflow = 64
|
||
delivered := make(chan struct{})
|
||
go func() {
|
||
for i := 0; i < overflow; i++ {
|
||
writer.RecordUsage(sample)
|
||
}
|
||
close(delivered)
|
||
}()
|
||
select {
|
||
case <-delivered:
|
||
case <-time.After(usageSinkProbeTimeout):
|
||
t.Fatal("队列满时 RecordUsage 阻塞了 —— 它必须丢弃而不能等")
|
||
}
|
||
|
||
// 丢弃不许静默：一条限流 WARN（后续丢弃被 30s 窗口压住，不刷屏），
|
||
// 且累计丢弃数是**全部**溢出条数。
|
||
raw := buf.String()
|
||
if !strings.Contains(raw, "llm usage queue full") {
|
||
t.Fatalf("队列满没有告警（不静默的红线）：%s", raw)
|
||
}
|
||
if got := strings.Count(raw, "llm usage queue full"); got != 1 {
|
||
t.Fatalf("限流告警记了 %d 条, want 1（否则库一挂就把日志刷爆）：%s", got, raw)
|
||
}
|
||
if writer.dropped != overflow {
|
||
t.Fatalf("累计丢弃 = %d, want %d", writer.dropped, overflow)
|
||
}
|
||
// 没有 Start ⇒ 一条都没落库（丢弃是真的丢弃，不是「攒着」）。
|
||
if rows := usageRowCount(t, conn); rows != 0 {
|
||
t.Fatalf("未启动写入器却落了 %d 行", rows)
|
||
}
|
||
}
|
||
|
||
// TestUsageWriter_PersistsBatchAndReaderSummarizes 钉住「批量写入 → 读回」：
|
||
// Stop 会把队列里剩下的样本在一个事务里写完，读侧按时间窗聚合，命中率按
|
||
// cacheHit / (cacheHit + cacheMiss) 算出。
|
||
func TestUsageWriter_PersistsBatchAndReaderSummarizes(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := context.Background()
|
||
conn := usageTestDB(t)
|
||
log, buf := captureJSONLog()
|
||
writer := NewUsageWriter(conn, log)
|
||
writer.Start(ctx)
|
||
|
||
samples := []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()
|
||
|
||
if rows := usageRowCount(t, conn); rows != int64(len(samples)) {
|
||
t.Fatalf("落库行数 = %d, want %d（日志：%s）", rows, len(samples), buf.String())
|
||
}
|
||
// 这里直接喂写入器、样本里**没有**会话 ⇒ session_id 落 NULL：无归属的调用
|
||
// （管理面板试连、未接循环的后台流程）与「属于某会话」在列里必须可区分，
|
||
// 所以「没有」写 NULL 而不是空串（打标链见 usage_session_test.go）。
|
||
var nullSessions int64
|
||
if err := conn.QueryRowContext(ctx,
|
||
`SELECT COUNT(*) FROM llm_usage WHERE session_id IS NULL`).Scan(&nullSessions); err != nil {
|
||
t.Fatalf("count null session_id: %v", err)
|
||
}
|
||
if nullSessions != int64(len(samples)) {
|
||
t.Fatalf("session_id 为 NULL 的行 = %d, want %d", nullSessions, len(samples))
|
||
}
|
||
|
||
until := time.Now().Add(time.Minute)
|
||
summary, err := NewUsageReader(conn).Summary(ctx, UsageQuery{Since: until.Add(-time.Hour), Until: until})
|
||
if err != nil {
|
||
t.Fatalf("Summary: %v", err)
|
||
}
|
||
wantTotals := UsageTotals{Calls: 3, Prompt: 2500, Completion: 350, CacheHit: 1400, CacheMiss: 1100}
|
||
if summary.Totals != wantTotals {
|
||
t.Fatalf("totals = %+v, want %+v", summary.Totals, wantTotals)
|
||
}
|
||
if got := summary.Totals.HitRate(); got != 0.56 {
|
||
t.Fatalf("hitRate = %v, want 0.56（1400/2500）", got)
|
||
}
|
||
if len(summary.ByPurpose) != 2 || summary.ByPurpose[0].Key != "chat" {
|
||
t.Fatalf("byPurpose = %+v, want chat 在前（2 次调用）", summary.ByPurpose)
|
||
}
|
||
if summary.ByPurpose[0].Calls != 2 || summary.ByPurpose[0].HitRate() != 0.45 {
|
||
t.Fatalf("byPurpose[0] = %+v, want calls 2 / hitRate 0.45", summary.ByPurpose[0])
|
||
}
|
||
if summary.ByPurpose[1].Key != "memory_extract" || summary.ByPurpose[1].HitRate() != 1 {
|
||
t.Fatalf("byPurpose[1] = %+v, want memory_extract / hitRate 1", summary.ByPurpose[1])
|
||
}
|
||
if len(summary.ByModel) != 2 || summary.ByModel[0].Key != "deepseek-chat" || summary.ByModel[0].Calls != 2 {
|
||
t.Fatalf("byModel = %+v, want deepseek-chat 在前（2 次调用）", summary.ByModel)
|
||
}
|
||
// 空窗（无调用）不是错误：零合计 + 命中率 0（分母 0 的口径）。
|
||
empty, err := NewUsageReader(conn).Summary(ctx, UsageQuery{
|
||
Since: until.Add(-48 * time.Hour), Until: until.Add(-47 * time.Hour),
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Summary(empty): %v", err)
|
||
}
|
||
if empty.Totals != (UsageTotals{}) || empty.Totals.HitRate() != 0 {
|
||
t.Fatalf("空窗 = %+v, want 全零且 hitRate 0", empty.Totals)
|
||
}
|
||
if len(empty.ByPurpose) != 0 || len(empty.ByModel) != 0 {
|
||
t.Fatalf("空窗不该有分组行：%+v / %+v", empty.ByPurpose, empty.ByModel)
|
||
}
|
||
}
|
||
|
||
// TestUsageReader_WindowAndPurposeFilter 钉住时间窗与用途过滤真的生效（写入侧
|
||
// 逐行给定 created_at，因此窗口边界是确定的，不靠墙钟）。
|
||
func TestUsageReader_WindowAndPurposeFilter(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := context.Background()
|
||
conn := usageTestDB(t)
|
||
insertUsageRow(t, conn, "chat", "m1", "2026-10-01T00:00:00.000Z", 100, 10, 90, 10)
|
||
insertUsageRow(t, conn, "chat", "m1", "2026-10-03T00:00:00.000Z", 200, 20, 180, 20)
|
||
insertUsageRow(t, conn, "noise_review", "m2", "2026-10-03T00:00:00.000Z", 40, 4, 0, 40)
|
||
|
||
reader := NewUsageReader(conn)
|
||
day := func(iso string) time.Time {
|
||
parsed, err := time.Parse(time.RFC3339, iso)
|
||
if err != nil {
|
||
t.Fatalf("parse %s: %v", iso, err)
|
||
}
|
||
return parsed
|
||
}
|
||
|
||
// 只覆盖第二天 ⇒ 两条算进来（m1 的 200 与 noise_review 的 40）。
|
||
all, err := reader.Summary(ctx, UsageQuery{
|
||
Since: day("2026-10-02T00:00:00Z"),
|
||
Until: day("2026-10-04T00:00:00Z"),
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Summary(window): %v", err)
|
||
}
|
||
if all.Totals.Calls != 2 || all.Totals.Prompt != 240 {
|
||
t.Fatalf("窗口内合计 = %+v, want calls 2 / prompt 240", all.Totals)
|
||
}
|
||
// 同窗口 + purpose 过滤 ⇒ 只剩 chat 那条。
|
||
chatOnly, err := reader.Summary(ctx, UsageQuery{
|
||
Since: day("2026-10-02T00:00:00Z"),
|
||
Until: day("2026-10-04T00:00:00Z"),
|
||
Purpose: "chat",
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Summary(purpose): %v", err)
|
||
}
|
||
if chatOnly.Totals.Calls != 1 || chatOnly.Totals.Prompt != 200 || chatOnly.ByPurpose[0].Key != "chat" {
|
||
t.Fatalf("purpose 过滤结果 = %+v, want 只剩 chat 的 200", chatOnly.Totals)
|
||
}
|
||
// 全窗（含第一天）⇒ 三条。
|
||
full, err := reader.Summary(ctx, UsageQuery{
|
||
Since: day("2026-09-30T00:00:00Z"),
|
||
Until: day("2026-10-05T00:00:00Z"),
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("Summary(full): %v", err)
|
||
}
|
||
if full.Totals.Calls != 3 || full.Totals.CacheHit != 270 || full.Totals.CacheMiss != 70 {
|
||
t.Fatalf("全窗合计 = %+v, want calls 3 / hit 270 / miss 70", full.Totals)
|
||
}
|
||
}
|
||
|
||
// TestPurgeUsageRemovesExpiredOnly 钉住保留期清理的判据：早于 before 的删掉、
|
||
// 之后的原样保留（before 恰好等于某行的时刻时保留 —— 判据是**严格早于**）。
|
||
func TestPurgeUsageRemovesExpiredOnly(t *testing.T) {
|
||
t.Parallel()
|
||
ctx := context.Background()
|
||
conn := usageTestDB(t)
|
||
insertUsageRow(t, conn, "chat", "m1", "2026-09-01T00:00:00.000Z", 10, 1, 5, 5)
|
||
insertUsageRow(t, conn, "chat", "m1", "2026-10-04T00:00:00.000Z", 20, 2, 10, 10)
|
||
|
||
cutoff, err := time.Parse(time.RFC3339, "2026-10-04T00:00:00Z")
|
||
if err != nil {
|
||
t.Fatalf("parse cutoff: %v", err)
|
||
}
|
||
removed, err := PurgeUsage(ctx, conn, cutoff)
|
||
if err != nil {
|
||
t.Fatalf("PurgeUsage: %v", err)
|
||
}
|
||
if removed != 1 {
|
||
t.Fatalf("删除 %d 行, want 1（只删严格早于 cutoff 的那条）", removed)
|
||
}
|
||
if rows := usageRowCount(t, conn); rows != 1 {
|
||
t.Fatalf("剩余行数 = %d, want 1", rows)
|
||
}
|
||
}
|
||
|
||
// insertUsageRow 直接写一行样本，created_at 由调用方给定（窗口/清理用例靠它把
|
||
// 时间钉死，不依赖墙钟）。
|
||
func insertUsageRow(t *testing.T, conn *sql.DB, purpose, model, createdAt string, prompt, completion, hit, miss int) {
|
||
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', ?, ?, ?, ?, ?, ?, 0, ?)`,
|
||
purpose, model, prompt, completion, hit, miss, prompt+completion, createdAt); err != nil {
|
||
t.Fatalf("insert usage row: %v", err)
|
||
}
|
||
}
|
||
|
||
// usageRowCount 返回 llm_usage 的当前行数。
|
||
func usageRowCount(t *testing.T, conn *sql.DB) int64 {
|
||
t.Helper()
|
||
var rows int64
|
||
if err := conn.QueryRowContext(context.Background(), `SELECT COUNT(*) FROM llm_usage`).Scan(&rows); err != nil {
|
||
t.Fatalf("count llm_usage: %v", err)
|
||
}
|
||
return rows
|
||
}
|