MailUI4Agents/gateway/internal/repo/platform_sessions_test.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)
	}
}