Agent-System/internal/server/chat_restore_test.go

package server_test

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"chat-agent-system/internal/bizdev"
	"chat-agent-system/internal/config"
	"chat-agent-system/internal/server"
)

func TestServer_ChatRestoreAfterRestart(t *testing.T) {
	cfg := &config.Config{
		HTTPHost:      "127.0.0.1",
		HTTPPort:      8889,
		DataDir:       t.TempDir(),
		AdminUsername: "admin",
		AdminPassword: "admin-test-pass",
		AgentBackend:  "stub",
	}
	svc := bizdev.MustNewService(t, cfg)
	srv := server.NewServer(cfg, svc)
	handler := srv.Routes()
	adminToken := loginTestAdmin(t, handler, cfg)

	sessID := "sess_restore_test_001"
	userID := int64(88)

	// Turn 1: 用户发送第一条消息
	chatReq1 := map[string]any{
		"session_id":   sessID,
		"user_id":      userID,
		"project_name": "高频超声手术刀系统",
		"device_type":  "超声手术设备",
		"message":      "你好，我们正在规划立项一款高频超声手术刀系统，预算200万",
	}
	bodyBytes1, _ := json.Marshal(chatReq1)
	req1 := httptest.NewRequest("POST", "/api/v1/chat", bytes.NewReader(bodyBytes1))
	req1.Header.Set("Content-Type", "application/json")
	rr1 := httptest.NewRecorder()
	handler.ServeHTTP(rr1, req1)

	if rr1.Code != http.StatusOK {
		t.Fatalf("turn 1 failed with status %d: %s", rr1.Code, rr1.Body.String())
	}

	var resp1 map[string]any
	if err := json.Unmarshal(rr1.Body.Bytes(), &resp1); err != nil {
		t.Fatalf("unmarshal resp1 failed: %v", err)
	}
	reply1, _ := resp1["reply_text"].(string)
	if reply1 == "" {
		t.Fatalf("turn 1 reply text is empty")
	}
	t.Logf("✓ Turn 1 完成: %s", reply1)

	// 验证消息已被持久化至 ContextManager
	msgs, err := bizdev.Snapshot(svc).Ctx.GetChatHistory(context.Background(), sessID, 10, 0)
	if err != nil || len(msgs) != 2 {
		t.Fatalf("expected 2 messages in ContextManager, got %d (err: %v)", len(msgs), err)
	}

	// 模拟服务重启 / 内存淘汰 (G9 CleanStaleSessions 触发清空内存 map)
	time.Sleep(5 * time.Millisecond)
	cleaned := srv.CleanStaleSessions(1 * time.Millisecond)
	if cleaned < 1 {
		t.Fatalf("expected at least 1 session to be cleaned from memory, got %d", cleaned)
	}
	t.Logf("✓ 成功模拟会话从内存完全淘汰 (释放会话数=%d)", cleaned)

	// Turn 2: 服务重启后，用户携相同 session_id 再次发起追问 (G3 自动恢复)
	chatReq2 := map[string]any{
		"session_id": sessID,
		"user_id":    userID,
		"message":    "关于刚才说的这个项目，国内型检和审评审批有哪些核心强标？",
	}
	bodyBytes2, _ := json.Marshal(chatReq2)
	req2 := httptest.NewRequest("POST", "/api/v1/chat", bytes.NewReader(bodyBytes2))
	req2.Header.Set("Content-Type", "application/json")
	rr2 := httptest.NewRecorder()
	handler.ServeHTTP(rr2, req2)

	if rr2.Code != http.StatusOK {
		t.Fatalf("turn 2 failed with status %d: %s", rr2.Code, rr2.Body.String())
	}

	var resp2 map[string]any
	if err := json.Unmarshal(rr2.Body.Bytes(), &resp2); err != nil {
		t.Fatalf("unmarshal resp2 failed: %v", err)
	}
	reply2, _ := resp2["reply_text"].(string)
	if reply2 == "" {
		t.Fatalf("turn 2 reply text is empty")
	}
	t.Logf("✓ Turn 2 完成: %s", reply2)

	// 验证完整历史记录包含全部 4 条消息 (Turn 1 user+asst, Turn 2 user+asst)
	allMsgs, err := bizdev.Snapshot(svc).Ctx.GetChatHistory(context.Background(), sessID, 10, 0)
	if err != nil {
		t.Fatalf("failed to query chat history: %v", err)
	}
	if len(allMsgs) != 4 {
		t.Fatalf("expected 4 total messages after restore, got %d", len(allMsgs))
	}
	if allMsgs[0].Role != "user" || !strings.Contains(allMsgs[0].Content, "高频超声手术刀系统") {
		t.Errorf("turn 1 user message mismatch: %+v", allMsgs[0])
	}
	if allMsgs[2].Role != "user" || !strings.Contains(allMsgs[2].Content, "核心强标") {
		t.Errorf("turn 2 user message mismatch: %+v", allMsgs[2])
	}

	// 通过 /api/v1/projects 检查恢复后的上下文事实（须带认证）
	reqCtx := httptest.NewRequest("GET", "/api/v1/projects?session_id="+sessID, nil)
	withBearerAuth(reqCtx, adminToken)
	rrCtx := httptest.NewRecorder()
	handler.ServeHTTP(rrCtx, reqCtx)
	if rrCtx.Code != http.StatusOK {
		t.Fatalf("GET /api/v1/projects failed: %d %s", rrCtx.Code, rrCtx.Body.String())
	}
	var ctxMap map[string]any
	_ = json.Unmarshal(rrCtx.Body.Bytes(), &ctxMap)
	facts, _ := ctxMap["blackboard_facts"].(map[string]any)
	if facts != nil {
		t.Logf("✓ 恢复后的项目黑板事实: %v", facts)
	}
}