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) } }