Cyrene_v2/internal/runtime/loop/usage_session_test.go

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
}