package loop import ( "context" "database/sql" "io" "net/http" "net/http/httptest" "strings" "testing" "git.yeij.top/AskaEth/Cyrene_v2/internal/config" "git.yeij.top/AskaEth/Cyrene_v2/internal/db" "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" ) // 本文件钉住切片 U1 遗留缺口的**整条链**:loop 在 Run 顶部播种会话 //(model/session_ctx.go)⇒ 适配器把它读进用量样本 ⇒ UsageWriter 写进 // `llm_usage.session_id`。 // // 为什么必须真落库、而不是只断言中间样本:缺口本身就是「这一列恒为 NULL」, // 只看中间对象会漏掉写入侧(把空串当值写、或干脆没绑参都照样「绿」)。这里因此 // 用真适配器(真 HTTP 请求)+真 UsageWriter+真 SQLite 走一遍,最后直接查列。 // // 流式与非流式**两条**路径都覆盖:它们是模型调用的两个不同入口 // (streamStep→ChatStream / agentStep→RunStep→Chat),漏一条就等于那半边没接线。 // TestRun_TagsModelUsageWithSessionID 端到端断言:一次回合里产生的每一行用量, // session_id 都是该回合自己的会话 id(不是 NULL,也不是别的会话)。 func TestRun_TagsModelUsageWithSessionID(t *testing.T) { t.Parallel() const session = "web:sess-e2e-usage" ctx := context.Background() for _, stream := range []bool{false, true} { name := "非流式路径" if stream { name = "流式路径" } t.Run(name, func(t *testing.T) { conn := openUsageTestDB(t) writer := model.NewUsageWriter(conn, testLogger()) if writer == nil { t.Fatal("NewUsageWriter 在接了库时不该返回 nil") } writer.Start(ctx) srv := openAIUsageServer(t, stream) ad := model.NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}, model.WithUsageSink(writer)) opts := Options{Agent: newTestAgent(t, ad, nil), Checkpoint: newTestCheckpointStore(t), MaxRounds: 4} if stream { // 只有挂了 OnStream 才会走流式路径(见 loop.step),因此它是 // 「覆盖 ChatStream」的开关,而不是可选的观察钩子。 opts.OnStream = func(model.StreamEvent) {} } l := New(opts) if _, err := l.Run(ctx, session, []model.Message{model.TextMessage(model.RoleUser, "hi")}); err != nil { t.Fatalf("Run: %v", err) } // Stop 会把这回合攒下的样本写完再返回,因此此刻查库是稳定的。 writer.Stop() var total, wrong int64 if err := conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM llm_usage`).Scan(&total); err != nil { t.Fatalf("count llm_usage: %v", err) } if total == 0 { t.Fatal("本回合一次模型调用都没落库(适配器边界没接上 sink?)") } // 既不是 NULL、也不是别的值:一条都不许错。 if err := conn.QueryRowContext(ctx, `SELECT COUNT(*) FROM llm_usage WHERE session_id IS NULL OR session_id <> ?`, session).Scan(&wrong); err != nil { t.Fatalf("count mis-tagged rows: %v", err) } if wrong != 0 { var sample sql.NullString _ = conn.QueryRowContext(ctx, `SELECT session_id FROM llm_usage LIMIT 1`).Scan(&sample) t.Fatalf("有 %d/%d 行的 session_id 不是本回合的 %q(首行 = %v)", wrong, total, session, sample) } }) } } // openUsageTestDB 打开一个带全部迁移的内存库(`llm_usage` 由此建出)。每个用例 // 一份私有库,并行用例天然隔离。 func openUsageTestDB(t *testing.T) *sql.DB { t.Helper() ctx := context.Background() conn, err := db.Open(ctx, ":memory:", testLogger()) if err != nil { t.Fatalf("db.Open: %v", err) } t.Cleanup(func() { _ = conn.Close() }) if err := db.MigrateUp(ctx, conn, testLogger()); err != nil { t.Fatalf("db.MigrateUp: %v", err) } return conn } // openAIUsageServer 返回一个 OpenAI 兼容端点:stream 为真时回 SSE 帧(末块带 // 用量),否则回一份普通 JSON 补全 —— 两种线格式都带上用量块,用量样本才真的 // 会产生(没有用量就没有样本,用例会误判成「没接线」)。 func openAIUsageServer(t *testing.T, stream bool) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { if stream { w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, strings.Join([]string{ `data: {"model":"gpt-test","choices":[{"index":0,"delta":{"content":"你好"}}]}`, `data: {"model":"gpt-test","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, `data: [DONE]`, }, "\n")+"\n") return } w.Header().Set("Content-Type", "application/json") _, _ = io.WriteString(w, `{ "id":"cmpl-1","model":"gpt-test", "choices":[{"index":0,"message":{"role":"assistant","content":"你好"},"finish_reason":"stop"}], "usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7} }`) })) t.Cleanup(srv.Close) return srv }