package ctxmgr_test import ( "chat-agent-system/internal/db" "context" "fmt" "os" "testing" "time" "chat-agent-system/internal/ctxmgr" "github.com/jackc/pgx/v5/pgxpool" ) func TestDeriveMessageCategory(t *testing.T) { tests := []struct { name string role string intent string branch string meta map[string]any expected string }{ { name: "milestone metadata", role: "user", meta: map[string]any{"milestone": true}, expected: ctxmgr.CategoryMilestone, }, { name: "attachments file op", role: "user", meta: map[string]any{"attachments": []any{"file.pdf"}}, expected: ctxmgr.CategoryFile, }, { name: "file_op export", role: "assistant", meta: map[string]any{"file_op": "export"}, expected: ctxmgr.CategoryFile, }, { name: "interactive_card review", role: "assistant", meta: map[string]any{"interactive_card": map[string]any{"type": "confirm_reject"}}, expected: ctxmgr.CategoryReview, }, { name: "task_request task", role: "assistant", meta: map[string]any{"task_request": map[string]any{"taskKey": "intake"}}, expected: ctxmgr.CategoryTask, }, { name: "DOCGOV intent review", role: "user", intent: "DOCGOV_REVIEW", expected: ctxmgr.CategoryReview, }, { name: "TRIGGER_TASK intent task", role: "user", intent: "TRIGGER_TASK", expected: ctxmgr.CategoryTask, }, { name: "QUERY_PROGRESS consult", role: "user", intent: "QUERY_PROGRESS", expected: ctxmgr.CategoryConsult, }, { name: "CASUAL_CHAT casual", role: "user", intent: "CASUAL_CHAT", expected: ctxmgr.CategoryCasual, }, { name: "system role system", role: "system", expected: ctxmgr.CategorySystem, }, { name: "fallback uncategorized", role: "user", intent: "", expected: ctxmgr.CategoryUncategorized, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { actual := ctxmgr.DeriveMessageCategory(tt.role, tt.intent, tt.branch, tt.meta) if actual != tt.expected { t.Errorf("DeriveMessageCategory() = %s, want %s", actual, tt.expected) } }) } } func TestChatWindow_MemoryEquivalenceAndPagination(t *testing.T) { mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig()) ctx := context.Background() sessID := "sess_test_window_mem" // 1. 参数校验 _, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{SessionID: ""}) if err == nil { t.Errorf("expected error for empty sessionID") } // 写入不同自然日的消息 // Day 1: 2026-09-01 // Day 2: 2026-09-02 // ... // Day 10: 2026-09-10 (包含 305 条消息触发截断) shanghaiTZ := time.FixedZone("CST", 8*3600) baseDate, _ := time.ParseInLocation("2006-01-02", "2026-09-01", shanghaiTZ) var targetAnchorID int64 for d := 0; d < 10; d++ { curDate := baseDate.AddDate(0, 0, d) curDayKey := curDate.Format("2006-01-02") msgCount := 5 if d == 9 { msgCount = 310 // 触发 MaxMsgPerDay(300) 截断 } for i := 0; i < msgCount; i++ { msgTime := curDate.Add(time.Duration(i) * time.Minute) meta := map[string]any{ "day_key": curDayKey, "created_at": msgTime, } id, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", fmt.Sprintf("Day %s Msg %d", curDayKey, i), "QUERY_PROGRESS", "", meta) if err != nil { t.Fatalf("RecordUserMessageWithID failed: %v", err) } if d == 5 && i == 2 { targetAnchorID = id } } } // 2. 默认拉取 7 天窗口 res, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{ SessionID: sessID, Days: 7, }) if err != nil { t.Fatalf("GetChatWindow failed: %v", err) } if res.SessionID != sessID { t.Errorf("expected SessionID %s, got %s", sessID, res.SessionID) } // 总共有 10 天,默认拉 7 天,应返回最近的 7 天 if len(res.Days) != 7 { t.Fatalf("expected 7 days loaded, got %d", len(res.Days)) } // 正序检查:最早日在前,最新日在后 for i := 1; i < len(res.Days); i++ { if res.Days[i-1].DayKey >= res.Days[i].DayKey { t.Errorf("expected ascending day order: day[%d]=%s should be before day[%d]=%s", i-1, res.Days[i-1].DayKey, i, res.Days[i].DayKey) } } // 验证还有更早日 if !res.HasMoreBefore { t.Errorf("expected has_more_before to be true") } oldestInWindow := res.OldestLoadedDay // 验证最新的一天(Day 10)发生截断 latestDay := res.Days[len(res.Days)-1] if latestDay.Count != 310 { t.Errorf("expected Day 10 count 310, got %d", latestDay.Count) } if !latestDay.Truncated { t.Errorf("expected Day 10 truncated to be true") } if latestDay.LoadedCount != 300 { t.Errorf("expected loaded_count 300, got %d", latestDay.LoadedCount) } // 3. 触顶加载更早日 (before_day) olderRes, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{ SessionID: sessID, Days: 7, BeforeDay: oldestInWindow, }) if err != nil { t.Fatalf("GetChatWindow before_day failed: %v", err) } // 剩余 3 天 if len(olderRes.Days) != 3 { t.Errorf("expected 3 older days, got %d", len(olderRes.Days)) } // 此段不应再有更早日 if olderRes.HasMoreBefore { t.Errorf("expected olderRes.HasMoreBefore to be false") } // 4. ListChatDates 统计测试 dates, err := mgr.ListChatDates(ctx, sessID) if err != nil { t.Fatalf("ListChatDates failed: %v", err) } if len(dates) != 10 { t.Fatalf("expected 10 dates, got %d", len(dates)) } // 验证分类统计 foundDay10 := false for _, dt := range dates { if dt.Count == 310 { foundDay10 = true if dt.Categories[ctxmgr.CategoryConsult] != 310 { t.Errorf("expected 310 consult messages in day 10 stats, got %d", dt.Categories[ctxmgr.CategoryConsult]) } } } if !foundDay10 { t.Errorf("day 10 with 310 messages not found in date stats") } // 5. GetChatDayMessages 锚点上下文居中测试 day5Key := baseDate.AddDate(0, 0, 5).Format("2006-01-02") day5Res, err := mgr.GetChatDayMessages(ctx, sessID, day5Key, targetAnchorID, 0, 0, 4) if err != nil { t.Fatalf("GetChatDayMessages failed: %v", err) } if day5Res.TotalCount != 5 { t.Errorf("expected total 5 messages in Day 5, got %d", day5Res.TotalCount) } // limit 4 -> returned 4 if len(day5Res.Messages) != 4 { t.Errorf("expected 4 messages returned around anchor, got %d", len(day5Res.Messages)) } } func TestChatWindow_DatabaseIntegration(t *testing.T) { dsn := os.Getenv("CONTEXT_DB_DSN") if dsn == "" { dsn = db.SkipIfNoTestDatabaseURL(t) } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() pool, err := pgxpool.New(ctx, dsn) if err != nil { t.Skipf("cannot connect to context db (%v), skipping DB integration test", err) return } defer pool.Close() if err := pool.Ping(ctx); err != nil { t.Skipf("context db ping failed (%v), skipping DB integration test", err) return } if err := db.RunMigrationsOnceForTest(ctx, pool, dsn); err != nil { t.Fatalf("migration failed: %v", err) } mgr := ctxmgr.NewContextManager(pool, nil, ctxmgr.DefaultTimeDecayConfig()) sessID := fmt.Sprintf("sess_db_window_%d", time.Now().UnixNano()) // 插入跨越两天的测试消息 shanghaiTZ := time.FixedZone("CST", 8*3600) d1 := time.Now().In(shanghaiTZ).AddDate(0, 0, -1) d2 := time.Now().In(shanghaiTZ) d1Key := d1.Format("2006-01-02") d2Key := d2.Format("2006-01-02") id1, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "昨天立项需求沟通", "SUPPLEMENT_INFO", "", map[string]any{ "day_key": d1Key, "created_at": d1, }) if err != nil { t.Fatalf("RecordUserMessageWithID d1 failed: %v", err) } id2, err := mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "今天补充型检报告", "QUERY_REGULATIONS", "", map[string]any{ "day_key": d2Key, "created_at": d2, }) if err != nil { t.Fatalf("RecordUserMessageWithID d2 failed: %v", err) } _, err = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "启动注册申报任务", "TRIGGER_TASK", "", map[string]any{ "day_key": d2Key, "created_at": d2.Add(time.Minute), }) if err != nil { t.Fatalf("RecordUserMessageWithID d2 task failed: %v", err) } // 1. 测试真实 DB 的 GetChatWindow winRes, err := mgr.GetChatWindow(ctx, ctxmgr.ChatWindowQuery{ SessionID: sessID, Days: 7, }) if err != nil { t.Fatalf("GetChatWindow failed: %v", err) } if len(winRes.Days) != 2 { t.Fatalf("expected 2 days in window, got %d", len(winRes.Days)) } if winRes.Days[0].DayKey != d1Key || winRes.Days[1].DayKey != d2Key { t.Errorf("expected sorted days [%s, %s], got [%s, %s]", d1Key, d2Key, winRes.Days[0].DayKey, winRes.Days[1].DayKey) } if winRes.Days[0].Messages[0].ID != id1 || winRes.Days[1].Messages[0].ID != id2 { t.Errorf("expected message IDs [%d, %d], got [%d, %d]", id1, id2, winRes.Days[0].Messages[0].ID, winRes.Days[1].Messages[0].ID) } // 2. 测试真实 DB 的 ListChatDates dateStats, err := mgr.ListChatDates(ctx, sessID) if err != nil { t.Fatalf("ListChatDates failed: %v", err) } if len(dateStats) != 2 { t.Fatalf("expected 2 date stats, got %d", len(dateStats)) } // 3. 测试真实 DB 的 GetChatDayMessages dayMsgs, err := mgr.GetChatDayMessages(ctx, sessID, d1Key, id1, 0, 0, 10) if err != nil { t.Fatalf("GetChatDayMessages failed: %v", err) } if dayMsgs.TotalCount != 1 || len(dayMsgs.Messages) != 1 || dayMsgs.Messages[0].ID != id1 { t.Errorf("expected day 1 message ID %d, got %v", id1, dayMsgs) } // 4. 测试真实 DB 的 ListChatCategories catStats, err := mgr.ListChatCategories(ctx, sessID) if err != nil { t.Fatalf("ListChatCategories failed: %v", err) } if len(catStats) < 2 { t.Fatalf("expected at least 2 categories, got %d", len(catStats)) } } func TestListChatCategories_MemoryMode(t *testing.T) { mgr := ctxmgr.NewContextManager(nil, nil, ctxmgr.DefaultTimeDecayConfig()) ctx := context.Background() sessID := "sess_mem_cat_test" // 注入不同分类的消息 _, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "阶段推进", "ADVANCE_STAGE", "", nil) // task _, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "启动任务", "TRIGGER_TASK", "", nil) // task _, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "合规咨询", "QUERY_REGULATIONS", "", nil) // consult _, _ = mgr.RecordUserMessageWithID(ctx, sessID, 1, "user", "立项里程碑", "", "", map[string]any{"milestone": true}) // milestone stats, err := mgr.ListChatCategories(ctx, sessID) if err != nil { t.Fatalf("ListChatCategories failed: %v", err) } if len(stats) != 3 { t.Fatalf("expected 3 distinct categories, got %d", len(stats)) } // 按 count DESC 排序,task 应排在首位 (count=2) if stats[0].Category != "task" || stats[0].Count != 2 || stats[0].Label != "任务执行" { t.Errorf("expected first category to be task (count=2), got %+v", stats[0]) } }