WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

scripts/graph_extraction_probe/main_test.go

@@ -0,0 +1,77 @@
+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")
+	}
+}