319 lines
8.9 KiB
Go
319 lines
8.9 KiB
Go
package ctxmgr_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"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 := "postgres://app:dev_secret_db_pass@127.0.0.1:5433/agent_system_dev?sslmode=disable"
|
|
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
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|