355 lines
13 KiB
Go
355 lines
13 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
// seedPlatformAgent 注册一个 Agent，平台会话镜像与它绑定。
|
||
// 与 quota_test.go 的 seedAgent 区分开：那个要指定 default_rounds，这里不关心。
|
||
func seedPlatformAgent(t *testing.T, name string) {
|
||
t.Helper()
|
||
if _, err := db.DB.ExecContext(context.Background(),
|
||
`INSERT INTO agents (agent_name, secret, platform, status) VALUES ($1, 'x', $1, 'online')`,
|
||
name); err != nil {
|
||
t.Fatalf("seed agent %s: %v", name, err)
|
||
}
|
||
}
|
||
|
||
// seedSessionWS 建一个带 workspace 与别名的会话。
|
||
func seedSessionWS(t *testing.T, alias, workspace, subject string) uuid.UUID {
|
||
t.Helper()
|
||
var id uuid.UUID
|
||
err := db.DB.QueryRowContext(context.Background(), `
|
||
INSERT INTO sessions (session_alias, workspace, from_agent, subject, alias_source)
|
||
VALUES ($1, $2, 'admin', $3, 'platform')
|
||
RETURNING session_id
|
||
`, alias, workspace, subject).Scan(&id)
|
||
if err != nil {
|
||
t.Fatalf("seed session %s: %v", alias, err)
|
||
}
|
||
return id
|
||
}
|
||
|
||
// seedMailWS 插一封带明确 to_workspace 的邮件。
|
||
func seedMailWS(t *testing.T, sessionID uuid.UUID, from, to, toWS, subject string) uuid.UUID {
|
||
t.Helper()
|
||
var id uuid.UUID
|
||
err := db.DB.QueryRowContext(context.Background(), `
|
||
INSERT INTO mails (session_id, from_name, from_workspace, to_name, to_workspace,
|
||
subject, body, cc_list, created_at)
|
||
VALUES ($1, $2, '', $3, $4, $5, 'body', '[]', $6)
|
||
RETURNING mail_id
|
||
`, sessionID, from, to, toWS, subject, nextSeedTime()).Scan(&id)
|
||
if err != nil {
|
||
t.Fatalf("seed mail: %v", err)
|
||
}
|
||
return id
|
||
}
|
||
|
||
// 这个测试是「会话别名没有正确显示曾经发生在工作区下的会话」那次故障的回归。
|
||
//
|
||
// 旧实现按 mails 反推工作区，条件是
|
||
// `to_workspace = $path OR from_workspace = $path`。
|
||
// 而 Agent 回信时 from_workspace 存的是 **Agent 名**（如 "dsh"）而不是路径，
|
||
// 于是一旦会话里只剩 Agent 的回信可匹配，反推就落空、别名列不出来。
|
||
// 现在 workspace 存在会话自己身上，与邮件里那些脏数据无关。
|
||
func TestSuggestSessionCandidatesUsesSessionWorkspace(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
|
||
sid := seedSessionWS(t, "brisk-harbor", "/home/program/agentmail", "缓存选型")
|
||
// 只有 Agent 的回信：from_workspace 是脏的（Agent 名），to_workspace 是人类（空）
|
||
if _, err := db.DB.ExecContext(context.Background(), `
|
||
INSERT INTO mails (session_id, from_name, from_workspace, to_name, to_workspace,
|
||
subject, body, cc_list, created_at)
|
||
VALUES ($1, 'dsh', 'dsh', 'admin', '', 'Re: 缓存选型', 'body', '[]', $2)
|
||
`, sid, nextSeedTime()); err != nil {
|
||
t.Fatalf("seed agent reply: %v", err)
|
||
}
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "dsh", "/home/program/agentmail")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 1 {
|
||
t.Fatalf("应有 1 个候选，实际 %d —— 会话的 workspace 列没被用上", len(got))
|
||
}
|
||
if got[0].Alias != "brisk-harbor" {
|
||
t.Errorf("别名错误：%q", got[0].Alias)
|
||
}
|
||
if got[0].Source != "mail" {
|
||
t.Errorf("来源应为 mail，实际 %q", got[0].Source)
|
||
}
|
||
if got[0].Title != "缓存选型" {
|
||
t.Errorf("标题应带出来：%q", got[0].Title)
|
||
}
|
||
}
|
||
|
||
// 历史会话的 workspace 列是空的（新加的列），必须回退到 mails.to_workspace 反推，
|
||
// 否则升级后所有老会话一夜之间从候选列表里消失。
|
||
func TestSuggestSessionCandidatesFallsBackToMails(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "opencode")
|
||
|
||
// workspace 留空，模拟升级前建立的会话
|
||
sid := seedSessionWS(t, "legacy-thread", "", "老线索")
|
||
seedMailWS(t, sid, "admin", "opencode", "/home/legacy", "老线索")
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "opencode", "/home/legacy")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 1 || got[0].Alias != "legacy-thread" {
|
||
t.Fatalf("老会话应能靠 mails 反推出来，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// 工作区不匹配的会话不能出现：候选项点下去就会被填进 session 位，
|
||
// 而 session 位是三态语义 —— 指向别处的会话会直接 404「无法送达」。
|
||
func TestSuggestSessionCandidatesFiltersByWorkspace(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
|
||
mine := seedSessionWS(t, "here-thread", "/home/a", "本区")
|
||
seedMailWS(t, mine, "admin", "dsh", "/home/a", "本区")
|
||
other := seedSessionWS(t, "there-thread", "/home/b", "别区")
|
||
seedMailWS(t, other, "admin", "dsh", "/home/b", "别区")
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "dsh", "/home/a")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 1 || got[0].Alias != "here-thread" {
|
||
t.Fatalf("只应给出本工作区的会话，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// path 为空（地址写成 `dsh` 而不带 @/path）时不按工作区过滤：
|
||
// 用户还没写到 path 段就该看到全部可续的会话。
|
||
func TestSuggestSessionCandidatesEmptyPathReturnsAll(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
|
||
a := seedSessionWS(t, "ws-a", "/home/a", "A")
|
||
seedMailWS(t, a, "admin", "dsh", "/home/a", "A")
|
||
b := seedSessionWS(t, "ws-b", "/home/b", "B")
|
||
seedMailWS(t, b, "admin", "dsh", "/home/b", "B")
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "dsh", "")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 2 {
|
||
t.Fatalf("path 为空应给出全部 2 条，实际 %d", len(got))
|
||
}
|
||
}
|
||
|
||
// 平台侧会话（人直接在 opencode/DSH 界面上开的）经心跳上报后也要能被选中 ——
|
||
// 这正是「定期从 agent 平台同步会话」要解决的问题。
|
||
func TestSuggestSessionCandidatesIncludesPlatformMirror(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "opencode")
|
||
|
||
now := time.Now()
|
||
err := ReplacePlatformSessions(context.Background(), "opencode", []PlatformSession{
|
||
{PlatformID: "ses_1", Workspace: "/home/program/agentmail", Slug: "witty-planet",
|
||
Title: "重构导入路径", MailDriven: false, UpdatedAt: &now},
|
||
{PlatformID: "ses_2", Workspace: "/home/other", Slug: "brave-comet",
|
||
Title: "别的工作区", MailDriven: false, UpdatedAt: &now},
|
||
})
|
||
if err != nil {
|
||
t.Fatalf("ReplacePlatformSessions: %v", err)
|
||
}
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "opencode", "/home/program/agentmail")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 1 {
|
||
t.Fatalf("应有 1 个平台候选，实际 %d：%+v", len(got), got)
|
||
}
|
||
if got[0].Alias != "witty-planet" || got[0].Source != "platform" {
|
||
t.Errorf("平台候选错误：%+v", got[0])
|
||
}
|
||
if got[0].Title != "重构导入路径" {
|
||
t.Errorf("标题应带出来：%q", got[0].Title)
|
||
}
|
||
}
|
||
|
||
// 同一别名两边都有时保留 mail 来源：它是「一定送得到」的保证，
|
||
// 镜像只是平台的说法。但镜像的标题该补上去 —— 平台标题通常比会话主题更贴切。
|
||
func TestSuggestSessionCandidatesMailWinsOverMirror(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "opencode")
|
||
|
||
// 本侧线索：有别名但主题为空
|
||
sid := seedSessionWS(t, "witty-planet", "/home/x", "")
|
||
seedMailWS(t, sid, "admin", "opencode", "/home/x", "某事")
|
||
|
||
now := time.Now()
|
||
if err := ReplacePlatformSessions(context.Background(), "opencode", []PlatformSession{
|
||
{PlatformID: "ses_1", Workspace: "/home/x", Slug: "witty-planet",
|
||
Title: "平台侧的标题", UpdatedAt: &now},
|
||
}); err != nil {
|
||
t.Fatalf("ReplacePlatformSessions: %v", err)
|
||
}
|
||
|
||
got, err := SuggestSessionCandidates(context.Background(), "admin", "opencode", "/home/x")
|
||
if err != nil {
|
||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||
}
|
||
if len(got) != 1 {
|
||
t.Fatalf("同名应合并成 1 条，实际 %d：%+v", len(got), got)
|
||
}
|
||
if got[0].Source != "mail" {
|
||
t.Errorf("应保留 mail 来源（它保证送得到），实际 %q", got[0].Source)
|
||
}
|
||
if got[0].Title != "平台侧的标题" {
|
||
t.Errorf("本侧标题为空时应补上镜像的：%q", got[0].Title)
|
||
}
|
||
}
|
||
|
||
// 上报是整表替换：平台侧删掉的会话必须从候选列表里消失。
|
||
// 增量合并会让它永远留着，而 session 位指向一条不存在的会话会直接 404。
|
||
func TestReplacePlatformSessionsIsFullReplace(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
ctx := context.Background()
|
||
|
||
if err := ReplacePlatformSessions(ctx, "dsh", []PlatformSession{
|
||
{PlatformID: "s1", Workspace: "/w", Slug: "one"},
|
||
{PlatformID: "s2", Workspace: "/w", Slug: "two"},
|
||
}); err != nil {
|
||
t.Fatalf("首次上报: %v", err)
|
||
}
|
||
if got, _ := SuggestSessionCandidates(ctx, "admin", "dsh", "/w"); len(got) != 2 {
|
||
t.Fatalf("首次上报应有 2 条，实际 %d", len(got))
|
||
}
|
||
|
||
// 第二次只报一条：另一条在平台侧已被删除
|
||
if err := ReplacePlatformSessions(ctx, "dsh", []PlatformSession{
|
||
{PlatformID: "s1", Workspace: "/w", Slug: "one"},
|
||
}); err != nil {
|
||
t.Fatalf("二次上报: %v", err)
|
||
}
|
||
got, _ := SuggestSessionCandidates(ctx, "admin", "dsh", "/w")
|
||
if len(got) != 1 || got[0].Alias != "one" {
|
||
t.Fatalf("整表替换失效，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// 无 slug 的平台会话不进候选：slug 是填进 session 位的值，
|
||
// 没有它这一项点下去只能得到一个空的 session 段。
|
||
func TestPlatformSessionsWithoutSlugAreSkipped(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
ctx := context.Background()
|
||
|
||
if err := ReplacePlatformSessions(ctx, "dsh", []PlatformSession{
|
||
{PlatformID: "s1", Workspace: "/w", Slug: ""},
|
||
{PlatformID: "s2", Workspace: "/w", Slug: "named"},
|
||
}); err != nil {
|
||
t.Fatalf("上报: %v", err)
|
||
}
|
||
got, _ := SuggestSessionCandidates(ctx, "admin", "dsh", "/w")
|
||
if len(got) != 1 || got[0].Alias != "named" {
|
||
t.Fatalf("无 slug 的应被跳过，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// 上报里的重复 platform_id 不该让整次事务失败（主键冲突）。
|
||
func TestReplacePlatformSessionsDedupes(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
ctx := context.Background()
|
||
|
||
if err := ReplacePlatformSessions(ctx, "dsh", []PlatformSession{
|
||
{PlatformID: "dup", Workspace: "/w", Slug: "first"},
|
||
{PlatformID: "dup", Workspace: "/w", Slug: "second"},
|
||
}); err != nil {
|
||
t.Fatalf("重复 id 不该报错: %v", err)
|
||
}
|
||
got, _ := SuggestSessionCandidates(ctx, "admin", "dsh", "/w")
|
||
if len(got) != 1 || got[0].Alias != "first" {
|
||
t.Fatalf("应保留第一条，实际 %+v", got)
|
||
}
|
||
}
|
||
|
||
// SetSessionWorkspace 只在为空时写入：会话的工作区在建立时就定下了，
|
||
// 之后不该被一封发往别处的邮件改掉 —— 那会让它在候选列表里凭空换工作区。
|
||
func TestSetSessionWorkspaceDoesNotOverwrite(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
sid := seedSessionWS(t, "fixed-ws", "/home/original", "某事")
|
||
if err := SetSessionWorkspace(ctx, sid, "/home/hijacked"); err != nil {
|
||
t.Fatalf("SetSessionWorkspace: %v", err)
|
||
}
|
||
var ws string
|
||
if err := db.DB.QueryRowContext(ctx,
|
||
`SELECT workspace FROM sessions WHERE session_id = $1`, sid).Scan(&ws); err != nil {
|
||
t.Fatalf("read back: %v", err)
|
||
}
|
||
if ws != "/home/original" {
|
||
t.Errorf("已有 workspace 被覆盖成 %q", ws)
|
||
}
|
||
|
||
// 空的那种要能补上（历史会话回填）
|
||
empty := seedSessionWS(t, "empty-ws", "", "某事")
|
||
if err := SetSessionWorkspace(ctx, empty, "/home/filled"); err != nil {
|
||
t.Fatalf("SetSessionWorkspace(empty): %v", err)
|
||
}
|
||
if err := db.DB.QueryRowContext(ctx,
|
||
`SELECT workspace FROM sessions WHERE session_id = $1`, empty).Scan(&ws); err != nil {
|
||
t.Fatalf("read back: %v", err)
|
||
}
|
||
if ws != "/home/filled" {
|
||
t.Errorf("空 workspace 未被补上，实际 %q", ws)
|
||
}
|
||
}
|
||
|
||
// CreateSession 要把 workspace 存下来 —— 这是整条链的起点，
|
||
// 漏在这里的话后面所有查询都只能靠 mails 反推。
|
||
func TestCreateSessionStoresWorkspace(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
id, err := CreateSession(ctx, nil, "admin", "带工作区", "/home/program/agentmail")
|
||
if err != nil {
|
||
t.Fatalf("CreateSession: %v", err)
|
||
}
|
||
var ws string
|
||
if err := db.DB.QueryRowContext(ctx,
|
||
`SELECT workspace FROM sessions WHERE session_id = $1`, id).Scan(&ws); err != nil {
|
||
t.Fatalf("read back: %v", err)
|
||
}
|
||
if ws != "/home/program/agentmail" {
|
||
t.Errorf("workspace 未落库：%q", ws)
|
||
}
|
||
}
|
||
|
||
// 归档的会话不进候选：归档就是「这条线索结束了」，
|
||
// 还出现在补全里等于邀请用户往一条已关闭的线索里发信。
|
||
func TestSuggestSessionCandidatesExcludesArchived(t *testing.T) {
|
||
setupTestDB(t)
|
||
seedPlatformAgent(t, "dsh")
|
||
ctx := context.Background()
|
||
|
||
sid := seedSessionWS(t, "done-thread", "/home/a", "已完成")
|
||
seedMailWS(t, sid, "admin", "dsh", "/home/a", "已完成")
|
||
if _, err := db.DB.ExecContext(ctx,
|
||
`UPDATE sessions SET status = 'archived' WHERE session_id = $1`, sid); err != nil {
|
||
t.Fatalf("archive: %v", err)
|
||
}
|
||
|
||
got, _ := SuggestSessionCandidates(ctx, "admin", "dsh", "/home/a")
|
||
if len(got) != 0 {
|
||
t.Fatalf("归档会话不该出现，实际 %+v", got)
|
||
}
|
||
}
|