package main import ( "context" "encoding/json" "errors" "strings" "testing" "github.com/Tencent/WeKnora/internal/models/chat" "github.com/Tencent/WeKnora/internal/types" ) type fakeProbeChat struct { chatDelegate response *types.ChatResponse err error } func (f *fakeProbeChat) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) { return f.response, f.err } func TestRecordingChatPreservesPerCallDiagnosticsWithoutContent(t *testing.T) { fake := &fakeProbeChat{response: &types.ChatResponse{ Content: "answer", ReasoningContent: "private reasoning", FinishReason: "length", Usage: types.TokenUsage{PromptTokens: 10, CompletionTokens: 4, TotalTokens: 14}, }} recorder := &recordingChat{chatDelegate: fake} thinking := false options := &chat.ChatOptions{MaxTokens: 4096, Thinking: &thinking} first, err := recorder.Chat(context.Background(), nil, options) if err != nil || first != fake.response { t.Fatal("recorder changed the response") } thinking = true fake.response = &types.ChatResponse{Content: "\u4e2d\u6587", FinishReason: "stop"} if _, err := recorder.Chat(context.Background(), nil, options); err != nil { t.Fatal(err) } if recorder.calls != 2 || len(recorder.records) != 2 { t.Fatal("missing call diagnostics") } a, b := recorder.records[0], recorder.records[1] if a.Number != 1 || a.MaxTokens != 4096 || a.Thinking == nil || *a.Thinking || !a.ResponseReceived || a.Failed || a.FinishReason != "length" || a.ContentRunes != 6 || a.ReasoningRunes != 17 || a.Usage.TotalTokens != 14 { t.Fatalf("unexpected first call metadata: %+v", a) } if b.Number != 2 || b.Thinking == nil || !*b.Thinking || b.ContentRunes != 2 || b.ReasoningRunes != 0 { t.Fatalf("unexpected second call metadata: %+v", b) } encoded, err := json.Marshal(recorder.records) if err != nil { t.Fatal(err) } if strings.Contains(string(encoded), "private reasoning") || strings.Contains(string(encoded), "answer") { t.Fatal("diagnostics exposed response content") } } func TestRecordingChatRecordsNoResponseWithoutLeakingError(t *testing.T) { want := errors.New("private upstream error") recorder := &recordingChat{chatDelegate: &fakeProbeChat{err: want}} response, err := recorder.Chat(context.Background(), nil, nil) if response != nil || err != want || len(recorder.records) != 1 { t.Fatal("recorder changed the failure") } record := recorder.records[0] if record.ResponseReceived || !record.Failed || record.Thinking != nil || record.ContentRunes != 0 || record.ReasoningRunes != 0 { t.Fatalf("unexpected failure metadata: %+v", record) } encoded, _ := json.Marshal(record) if strings.Contains(string(encoded), want.Error()) { t.Fatal("diagnostics exposed upstream error") } }