Cyrene_v2/internal/runtime/model/usage_session_test.go

package model

import (
	"context"
	"net/http"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// 本文件钉住「会话标识端到端进 llm_usage.session_id」的前半程（模型层）：
// loop 经 ctx 播种会话（session_ctx.go）⇒ 适配器在**两条**调用路径上都把
// 它读进 UsageSample ⇒ 没播种时是空串（落库据此写 NULL，而不是空串）。
//
// 后半程（真的写进列）由 usage_store_test.go 的
// TestUsageWriter_PersistsSessionID 钉住，loop 侧的播种由
// internal/runtime/loop/usage_session_test.go 端到端钉住。

// sessionUsageRequest 是最小可用请求（只关心用量样本里的会话字段）。
func sessionUsageRequest() Request {
	return Request{Purpose: "chat", Messages: []Message{TextMessage(RoleUser, "hi")}}
}

// TestUsageSink_CarriesSessionFromContext 钉住打标通道：会话挂在 ctx 上时，
// 流式与非流式两条路径交给 sink 的样本都带着它。
//
// 两条路径都要覆盖：它们进入适配器的方式不同（ChatStream 由循环直接驱动，
// Chat 由 agent.RunStep 驱动），漏任何一条都会让那半边的用量归不了档。
func TestUsageSink_CarriesSessionFromContext(t *testing.T) {
	t.Parallel()
	const session = "web:sess-usage"

	jsonSrv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil)
	streamSrv := newCompletionServer(t, http.StatusOK, "text/event-stream", 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", nil)
	anthropicSrv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, nil)

	ctx := WithSessionID(context.Background(), session)
	probes := []struct {
		name string
		call func(sink UsageSink) error
	}{
		{"openai 非流式 Chat", func(sink UsageSink) error {
			ad := NewOpenAIChat(config.LLMConfig{BaseURL: jsonSrv.URL, Model: "gpt-test"}, WithUsageSink(sink))
			_, err := ad.Chat(ctx, sessionUsageRequest(), Options{})
			return err
		}},
		{"openai 流式 ChatStream", func(sink UsageSink) error {
			ad := NewOpenAIChat(config.LLMConfig{BaseURL: streamSrv.URL, Model: "gpt-test"}, WithUsageSink(sink))
			_, err := ad.ChatStream(ctx, sessionUsageRequest(), nil)
			return err
		}},
		{"anthropic 非流式 Chat", func(sink UsageSink) error {
			ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: anthropicSrv.URL, AnthropicModel: "claude-test"},
				WithAnthropicUsageSink(sink))
			_, err := ad.Chat(ctx, sessionUsageRequest(), Options{})
			return err
		}},
	}

	for _, p := range probes {
		sink := &fakeUsageSink{}
		if err := p.call(sink); err != nil {
			t.Fatalf("%s: 调用失败: %v", p.name, err)
		}
		if len(sink.samples) != 1 {
			t.Fatalf("%s: sink 收到 %d 条样本, want 1", p.name, len(sink.samples))
		}
		if got := sink.samples[0].SessionID; got != session {
			t.Fatalf("%s: 样本 SessionID = %q, want %q（会话打标在这一路上丢了）", p.name, got, session)
		}
	}
}

// TestUsageSink_UnseededSessionStaysEmpty 钉住「没播种 ⇒ 空串」：管理面板试连、
// 未接循环的后台流程本就不属于任何会话，样本里必须是空串 —— 写入侧据此写 NULL；
// 一旦这里变成别的值（比如 "unknown"），列里就分不出「无归属」与「真会话」。
func TestUsageSink_UnseededSessionStaysEmpty(t *testing.T) {
	t.Parallel()
	srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil)
	sink := &fakeUsageSink{}
	ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}, WithUsageSink(sink))

	if _, err := ad.Chat(context.Background(), sessionUsageRequest(), Options{}); err != nil {
		t.Fatalf("Chat: %v", err)
	}
	if len(sink.samples) != 1 {
		t.Fatalf("sink 收到 %d 条样本, want 1", len(sink.samples))
	}
	if got := sink.samples[0].SessionID; got != "" {
		t.Fatalf("未播种会话时样本 SessionID = %q, want 空串", got)
	}
}

// TestUsageWriter_PersistsSessionID 钉住落库侧：样本带会话 ⇒ 该列写会话；
// 样本不带会话 ⇒ 写 NULL（**不是空串**），且读回来的会话就是写入的那个。
func TestUsageWriter_PersistsSessionID(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	conn := usageTestDB(t)
	writer := NewUsageWriter(conn, discardLogger("model/usage_session_test"))
	writer.Start(ctx)

	const session = "web:sess-persist"
	writer.RecordUsage(UsageSample{Adapter: "openai", Model: "gpt-test", Purpose: "chat", SessionID: session})
	writer.RecordUsage(UsageSample{Adapter: "openai", Model: "gpt-test", Purpose: "chat"})
	writer.Stop()

	var stored string
	if err := conn.QueryRowContext(ctx,
		`SELECT session_id FROM llm_usage WHERE session_id IS NOT NULL`).Scan(&stored); err != nil {
		t.Fatalf("读回 session_id: %v", err)
	}
	if stored != session {
		t.Fatalf("落库 session_id = %q, want %q", stored, session)
	}

	var nullRows, emptyRows int64
	if err := conn.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM llm_usage WHERE session_id IS NULL`).Scan(&nullRows); err != nil {
		t.Fatalf("count null session_id: %v", err)
	}
	if err := conn.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM llm_usage WHERE session_id = ''`).Scan(&emptyRows); err != nil {
		t.Fatalf("count empty session_id: %v", err)
	}
	if nullRows != 1 {
		t.Fatalf("session_id IS NULL 的行 = %d, want 1（无归属的调用必须写 NULL）", nullRows)
	}
	if emptyRows != 0 {
		t.Fatalf("session_id = '' 的行 = %d, want 0（空串会把「无归属」写成一个真实值）", emptyRows)
	}
}