chore: directory migration - gateway→server, web→client/electron
This commit is contained in:
274
server/internal/repo/adopt_alias_test.go
Normal file
274
server/internal/repo/adopt_alias_test.go
Normal file
@ -0,0 +1,274 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// ─── 接管会话的别名不可被平台命名同步覆盖 ───
|
||||
//
|
||||
// 锁的是一次生产事故的**第二环**（第一环是调度器漏 platform_session_id）：
|
||||
//
|
||||
// 12:12 人选中补全里的「项目定位」→ 接管平台会话 01a05a5e，本侧建 26e26477
|
||||
// 12:20 日历提醒省略 session 位 → 落进 26e26477（第三环，见 TestDefaultSession...）
|
||||
// 12:20 插件收不到 platform_session_id → 另开一条 pi 会话
|
||||
// 12:20 那条新会话的名字经 /sessions/{id}/sync 回写
|
||||
// → SyncSessionAlias 把 26e26477 的别名冲成「日程提醒：…」
|
||||
//
|
||||
// 结果：人在补全里选的名字凭空消失，同一条会话在候选列表里出现两次
|
||||
// （一次用被冲掉的别名、一次用镜像里的原始 slug），而另一条真实会话被
|
||||
// 按别名字符串去重吃掉了。
|
||||
//
|
||||
// 接管会话的别名是**人从补全里选中的平台 slug**，任何平台命名同步都不该动它。
|
||||
func TestSyncSessionAliasNeverOverwritesAdoptedAlias(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "01a05a5e", "项目定位", "/home/program/agentmail", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
if got := SessionAliasOf(ctx, id); got != "项目定位" {
|
||||
t.Fatalf("接管后别名 = %q，期望 项目定位", got)
|
||||
}
|
||||
|
||||
// 平台侧同步一个完全不同的名字（生产上就是日历提醒的主题）
|
||||
final, err := SyncSessionAlias(ctx, id, "日程提醒：小宅自测")
|
||||
if err != nil {
|
||||
t.Fatalf("SyncSessionAlias: %v", err)
|
||||
}
|
||||
if final != "项目定位" {
|
||||
t.Errorf("同步返回 %q —— 接管会话的别名不该被改", final)
|
||||
}
|
||||
if got := SessionAliasOf(ctx, id); got != "项目定位" {
|
||||
t.Errorf("库里别名变成了 %q —— 人在补全里选的名字被冲掉了", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 普通（非接管）会话仍然接受平台命名同步 —— 别名复用平台命名是既定决策，
|
||||
// 上面那道门不能把它一起关掉。
|
||||
func TestSyncSessionAliasStillWorksForNormalSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := CreateSession(ctx, nil, "pi", "邮件驱动的会话", "/tmp/ws")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
if _, err := EnsureSessionAlias(ctx, id, "pi-初始别名"); err != nil {
|
||||
t.Fatalf("EnsureSessionAlias: %v", err)
|
||||
}
|
||||
|
||||
final, err := SyncSessionAlias(ctx, id, "平台生成的名字")
|
||||
if err != nil {
|
||||
t.Fatalf("SyncSessionAlias: %v", err)
|
||||
}
|
||||
if final != "平台生成的名字" {
|
||||
t.Errorf("普通会话应当接受同步，得到 %q", final)
|
||||
}
|
||||
}
|
||||
|
||||
// 接管会话**没有**别名时（理论上不会发生，AdoptPlatformSession 一定给一个）
|
||||
// 仍然允许写入 —— 否则那条会话永远无法寻址。
|
||||
func TestSyncSessionAliasFillsEmptyAdoptedAlias(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := CreateSession(ctx, nil, "pi", "标题", "/tmp/ws")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
// 手工造出「有 platform_id 但无别名」的状态
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET platform_id = 'pid-x' WHERE session_id = $1`, id); err != nil {
|
||||
t.Fatalf("置 platform_id: %v", err)
|
||||
}
|
||||
|
||||
final, err := SyncSessionAlias(ctx, id, "补上一个名字")
|
||||
if err != nil {
|
||||
t.Fatalf("SyncSessionAlias: %v", err)
|
||||
}
|
||||
if final != "补上一个名字" {
|
||||
t.Errorf("无别名的接管会话应当允许写入，得到 %q", final)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 接管会话不是「默认会话」 ───
|
||||
//
|
||||
// 事故的**第三环**：日历提醒的收件地址省略 session 位（`homeagent` 而不是
|
||||
// `homeagent@/x.某会话`），走 FindOrCreateDefaultSession。它原来只按
|
||||
// 「参与过 + workspace 匹配 + 未归档」挑最近活跃的一条 —— 于是挑中了人
|
||||
// 刚刚显式指定的那条接管会话。
|
||||
//
|
||||
// 接管会话是人**点名**要谈的一条线索，不该被省略 session 位的邮件当默认会话。
|
||||
func TestFindOrCreateDefaultSessionSkipsAdoptedSessions(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
// 一条接管会话，且有邮件（满足 EXISTS 条件）
|
||||
adopted, err := AdoptPlatformSession(ctx, "pi", "pid-adopted", "人选的线索", "/home/program/agentmail", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, adopted, nil, "jianf", "", "pi", "/home/program/agentmail",
|
||||
"人发的第一封", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
// 省略 session 位投递 → 不该落进那条接管会话
|
||||
got, err := FindOrCreateDefaultSession(ctx, "pi", "/home/program/agentmail", "calendar", "日程提醒")
|
||||
if err != nil {
|
||||
t.Fatalf("FindOrCreateDefaultSession: %v", err)
|
||||
}
|
||||
if got == adopted {
|
||||
t.Error("省略 session 位的邮件落进了接管会话 —— 那是人显式指定的线索")
|
||||
}
|
||||
|
||||
// 该新建一条，且它不带 platform_id
|
||||
if pid := PlatformIDOf(ctx, got); pid != "" {
|
||||
t.Errorf("新建的默认会话不该有 platform_id，得到 %q", pid)
|
||||
}
|
||||
}
|
||||
|
||||
// 普通会话仍然可以作为默认会话被复用 —— 上面那道门不能把它一起关掉，
|
||||
// 否则每封省略 session 位的邮件都会新开一条会话。
|
||||
func TestFindOrCreateDefaultSessionStillReusesNormalSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
first, err := FindOrCreateDefaultSession(ctx, "pi", "/tmp/ws", "jianf", "第一封")
|
||||
if err != nil {
|
||||
t.Fatalf("第一次: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, first, nil, "jianf", "", "pi", "/tmp/ws", "第一封", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
second, err := FindOrCreateDefaultSession(ctx, "pi", "/tmp/ws", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatalf("第二次: %v", err)
|
||||
}
|
||||
if second != first {
|
||||
t.Error("普通默认会话应当被复用，否则每封省略 session 位的邮件都开新会话")
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 候选列表按 platform_id 去重 ───
|
||||
//
|
||||
// 事故的**第四环**：`SuggestSessionCandidates` 的去重只比别名字符串。
|
||||
// 别名一被冲掉，同一条会话就在列表里出现两次：
|
||||
//
|
||||
// 候选 1 日程提醒：…（被冲掉的别名） source=mail ← 26e26477
|
||||
// 候选 2 项目定位（镜像里的原始 slug） source=platform ← 也是 26e26477
|
||||
//
|
||||
// 更糟的是**另一条真实会话被吃掉了**：它的 slug 恰好等于候选 1 那个
|
||||
// 被冲掉的别名，于是 `seen[slug]` 命中、被 continue 跳过。
|
||||
// 人在界面上看到两条，实际只有一条能选，而第三条不存在于列表里。
|
||||
func TestSuggestSessionCandidatesDedupesByPlatformID(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedPlatformAgent(t, "pi")
|
||||
|
||||
// 接管一条平台会话
|
||||
adopted, err := AdoptPlatformSession(ctx, "pi", "01a05a5e", "项目定位", "/home/program/agentmail", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, adopted, nil, "jianf", "", "pi", "/home/program/agentmail",
|
||||
"主题", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
// 镜像里同时有它与另一条真实会话
|
||||
now := time.Now()
|
||||
if err := ReplacePlatformSessions(ctx, "pi", []PlatformSession{
|
||||
{PlatformID: "01a05a5e", Workspace: "/home/program/agentmail", Slug: "项目定位",
|
||||
Title: "项目定位", MailDriven: true, UpdatedAt: &now},
|
||||
{PlatformID: "01a06aa5", Workspace: "/home/program/agentmail", Slug: "另一条真实会话",
|
||||
Title: "另一条", MailDriven: true, UpdatedAt: &now},
|
||||
}); err != nil {
|
||||
t.Fatalf("上报镜像: %v", err)
|
||||
}
|
||||
|
||||
got, err := SuggestSessionCandidates(ctx, "jianf", "pi", "/home/program/agentmail")
|
||||
if err != nil {
|
||||
t.Fatalf("SuggestSessionCandidates: %v", err)
|
||||
}
|
||||
|
||||
// 应当恰好两条：接管那条（mail 来源）+ 另一条真实会话（platform 来源）
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("应有 2 个候选，实际 %d：%+v", len(got), got)
|
||||
}
|
||||
|
||||
byAlias := map[string]SessionCandidate{}
|
||||
for _, c := range got {
|
||||
byAlias[c.Alias] = c
|
||||
}
|
||||
if c, ok := byAlias["项目定位"]; !ok {
|
||||
t.Error("接管会话应当在候选里")
|
||||
} else if c.Source != "mail" {
|
||||
t.Errorf("接管会话的来源应是 mail（保证送得到），得到 %q", c.Source)
|
||||
}
|
||||
if c, ok := byAlias["另一条真实会话"]; !ok {
|
||||
t.Error("另一条真实会话被吃掉了 —— 那正是 bug 的表现")
|
||||
} else if c.Source != "platform" {
|
||||
t.Errorf("未接管的平台会话来源应是 platform，得到 %q", c.Source)
|
||||
}
|
||||
}
|
||||
|
||||
// 别名被冲掉之后也不该出现重复项。
|
||||
//
|
||||
// 这是事故现场的**精确复现**：本侧别名与镜像 slug 不一致（别名被另一条会话的
|
||||
// 命名同步冲掉了），此时按别名字符串去重必然漏，只有按 platform_id 才对。
|
||||
func TestSuggestSessionCandidatesNoDupWhenAliasDiverged(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedPlatformAgent(t, "pi")
|
||||
|
||||
adopted, err := AdoptPlatformSession(ctx, "pi", "01a05a5e", "项目定位", "/home/program/agentmail", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, adopted, nil, "jianf", "", "pi", "/home/program/agentmail",
|
||||
"主题", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
// 模拟别名被冲掉（绕过 SyncSessionAlias 的守卫直接改库 ——
|
||||
// 存量数据里可能已经有这种状态）
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET session_alias = $1 WHERE session_id = $2`,
|
||||
"日程提醒：小宅自测", adopted); err != nil {
|
||||
t.Fatalf("改别名: %v", err)
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if err := ReplacePlatformSessions(ctx, "pi", []PlatformSession{
|
||||
{PlatformID: "01a05a5e", Workspace: "/home/program/agentmail", Slug: "项目定位",
|
||||
Title: "项目定位", MailDriven: true, UpdatedAt: &now},
|
||||
}); err != nil {
|
||||
t.Fatalf("上报镜像: %v", err)
|
||||
}
|
||||
|
||||
got, err := SuggestSessionCandidates(ctx, "jianf", "pi", "/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].Source != "mail" {
|
||||
t.Errorf("应保留 mail 来源（它保证送得到），得到 %q", got[0].Source)
|
||||
}
|
||||
}
|
||||
|
||||
// 存量数据里可能已经有这种状态（守卫是后加的）。
|
||||
228
server/internal/repo/adopt_test.go
Normal file
228
server/internal/repo/adopt_test.go
Normal file
@ -0,0 +1,228 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
func seedPlatformMirror(t *testing.T, agentName string, list []PlatformSession) {
|
||||
t.Helper()
|
||||
if err := ReplacePlatformSessions(context.Background(), agentName, list); err != nil {
|
||||
t.Fatalf("上报镜像: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 补全把平台会话列为候选，投递侧必须能命中同一条。
|
||||
// 此前 FindNamedSessionFor 只查 sessions 表 —— 候选列表在承诺一件做不到的事。
|
||||
func TestFindPlatformSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
seedPlatformMirror(t, "pi", []PlatformSession{
|
||||
{PlatformID: "pi-sess-1", Workspace: "/home/program/agentmail",
|
||||
Slug: "设计文档-项目定位", Title: "邮件驱动·多智能体协作平台", UpdatedAt: &now},
|
||||
{PlatformID: "pi-sess-2", Workspace: "/tmp/other",
|
||||
Slug: "别处的会话", Title: "无关", UpdatedAt: &now},
|
||||
})
|
||||
|
||||
t.Run("按 slug + workspace 命中", func(t *testing.T) {
|
||||
pid, ws, title, err := FindPlatformSession(ctx, "pi", "设计文档-项目定位", "/home/program/agentmail")
|
||||
if err != nil {
|
||||
t.Fatalf("查找: %v", err)
|
||||
}
|
||||
if pid != "pi-sess-1" {
|
||||
t.Errorf("platform_id = %q", pid)
|
||||
}
|
||||
if ws != "/home/program/agentmail" {
|
||||
t.Errorf("workspace = %q", ws)
|
||||
}
|
||||
if title != "邮件驱动·多智能体协作平台" {
|
||||
t.Errorf("title = %q", title)
|
||||
}
|
||||
})
|
||||
|
||||
// 地址省略 path 位时不限工作区
|
||||
t.Run("workspace 为空时不限", func(t *testing.T) {
|
||||
if pid, _, _, err := FindPlatformSession(ctx, "pi", "别处的会话", ""); err != nil || pid != "pi-sess-2" {
|
||||
t.Errorf("得到 %q err=%v", pid, err)
|
||||
}
|
||||
})
|
||||
|
||||
// workspace 不匹配时不该命中 —— 那会让邮件投进另一个项目的会话
|
||||
t.Run("workspace 不匹配不命中", func(t *testing.T) {
|
||||
if _, _, _, err := FindPlatformSession(ctx, "pi", "别处的会话", "/home/program/agentmail"); err == nil {
|
||||
t.Error("workspace 不同却命中了")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("别的 Agent 的镜像不串", func(t *testing.T) {
|
||||
if _, _, _, err := FindPlatformSession(ctx, "dsh", "设计文档-项目定位", ""); err == nil {
|
||||
t.Error("dsh 命中了 pi 的会话")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("空参数返回 not found 而不是 panic", func(t *testing.T) {
|
||||
if _, _, _, err := FindPlatformSession(ctx, "", "x", ""); err != ErrSessionNotFound {
|
||||
t.Errorf("空 agent 应给 ErrSessionNotFound，得到 %v", err)
|
||||
}
|
||||
if _, _, _, err := FindPlatformSession(ctx, "pi", "", ""); err != ErrSessionNotFound {
|
||||
t.Errorf("空 slug 应给 ErrSessionNotFound，得到 %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// 接管后本侧有正式身份：可寻址（别名）、绑定 platform_id、workspace 用会话真实的。
|
||||
func TestAdoptPlatformSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pi-sess-1", "设计文档-项目定位",
|
||||
"/home/program/agentmail", "邮件驱动·多智能体协作平台")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
|
||||
// 别名复用平台 slug：人在补全里看到的就是那个名字，换掉会让他找不到
|
||||
if alias := SessionAliasOf(ctx, id); alias != "设计文档-项目定位" {
|
||||
t.Errorf("别名 = %q，期望复用平台 slug", alias)
|
||||
}
|
||||
if pid := PlatformIDOf(ctx, id); pid != "pi-sess-1" {
|
||||
t.Errorf("platform_id = %q", pid)
|
||||
}
|
||||
|
||||
// workspace 取平台会话的真实 cwd
|
||||
var ws string
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT workspace FROM sessions WHERE session_id = $1`, id).Scan(&ws); err != nil {
|
||||
t.Fatalf("读 workspace: %v", err)
|
||||
}
|
||||
if ws != "/home/program/agentmail" {
|
||||
t.Errorf("workspace = %q", ws)
|
||||
}
|
||||
}
|
||||
|
||||
// 普通会话的 platform_id 必须是空串（不是接管来的）。
|
||||
func TestPlatformIDOfEmptyForNormalSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := CreateSession(ctx, nil, "pi", "普通邮件会话", "/tmp/x")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
if pid := PlatformIDOf(ctx, id); pid != "" {
|
||||
t.Errorf("普通会话的 platform_id 应为空，得到 %q", pid)
|
||||
}
|
||||
}
|
||||
|
||||
// 一条平台会话只能被接管一次。
|
||||
//
|
||||
// 第二次投递必须复用第一次建的本侧会话 —— 否则同一条 TUI 对话会在邮箱里
|
||||
// 裂成多条互不相干的线索：人看到三个同名会话，而回信只落在其中一条上。
|
||||
func TestFindSessionByPlatformIDPreventsDoubleAdopt(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pi-sess-1", "某会话", "/tmp/ws", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
// 接管后还没有邮件 —— 此时反查不到（EXISTS 子句要求有本 Agent 参与的邮件）
|
||||
if _, err := FindSessionByPlatformID(ctx, "pi", "pi-sess-1"); err == nil {
|
||||
t.Log("注意：无邮件时也能反查到")
|
||||
}
|
||||
|
||||
// 投一封进去，让参与关系成立
|
||||
if _, err := CreateMail(ctx, id, nil, "jianf", "", "pi", "/tmp/ws", "主题", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
got, err := FindSessionByPlatformID(ctx, "pi", "pi-sess-1")
|
||||
if err != nil {
|
||||
t.Fatalf("反查: %v", err)
|
||||
}
|
||||
if got != id {
|
||||
t.Errorf("反查到 %v，期望 %v", got, id)
|
||||
}
|
||||
|
||||
// 别的 platform_id 查不到
|
||||
if _, err := FindSessionByPlatformID(ctx, "pi", "pi-sess-999"); err != ErrSessionNotFound {
|
||||
t.Errorf("不存在的 platform_id 应给 ErrSessionNotFound，得到 %v", err)
|
||||
}
|
||||
// 别的 Agent 查不到（参与关系不成立）
|
||||
if _, err := FindSessionByPlatformID(ctx, "dsh", "pi-sess-1"); err != ErrSessionNotFound {
|
||||
t.Errorf("dsh 不该查到 pi 的接管会话，得到 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 整表替换镜像后，已接管的本侧会话不受影响。
|
||||
//
|
||||
// 镜像是平台当前状态的快照、会被整表替换；而 sessions.platform_id 是本侧的
|
||||
// 持久绑定。平台侧那条会话被删掉之后，本侧线索与历史邮件仍然要在。
|
||||
func TestAdoptedSessionSurvivesMirrorReplace(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
now := time.Now()
|
||||
seedPlatformMirror(t, "pi", []PlatformSession{
|
||||
{PlatformID: "pi-sess-1", Workspace: "/tmp/ws", Slug: "会话甲", UpdatedAt: &now},
|
||||
})
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pi-sess-1", "会话甲", "/tmp/ws", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
|
||||
// 平台侧删了那条会话（新快照里没有它）
|
||||
seedPlatformMirror(t, "pi", []PlatformSession{
|
||||
{PlatformID: "pi-sess-2", Workspace: "/tmp/ws", Slug: "会话乙", UpdatedAt: &now},
|
||||
})
|
||||
|
||||
// 本侧绑定与别名都还在
|
||||
if pid := PlatformIDOf(ctx, id); pid != "pi-sess-1" {
|
||||
t.Errorf("镜像替换后 platform_id 丢了：%q", pid)
|
||||
}
|
||||
if alias := SessionAliasOf(ctx, id); alias != "会话甲" {
|
||||
t.Errorf("别名丢了：%q", alias)
|
||||
}
|
||||
// 但镜像里查不到了（补全不再列它，符合预期）
|
||||
if _, _, _, err := FindPlatformSession(ctx, "pi", "会话甲", ""); err != ErrSessionNotFound {
|
||||
t.Errorf("镜像里应已消失，得到 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 接管用的 slug 与本侧某条无关会话撞名时要自动加后缀（别名全局唯一）。
|
||||
func TestAdoptHandlesAliasCollision(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
// 先占掉这个别名
|
||||
taken := "撞名的别名"
|
||||
if _, err := CreateSession(ctx, &taken, "pi", "已存在", "/tmp/a"); err != nil {
|
||||
t.Fatalf("建占位会话: %v", err)
|
||||
}
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pi-sess-x", taken, "/tmp/b", "标题")
|
||||
if err != nil {
|
||||
t.Fatalf("接管: %v", err)
|
||||
}
|
||||
alias := SessionAliasOf(ctx, id)
|
||||
if alias == "" {
|
||||
t.Fatal("接管后没有别名 —— 这条会话将无法寻址")
|
||||
}
|
||||
if alias == taken {
|
||||
t.Errorf("别名与已存在的重复了：%q", alias)
|
||||
}
|
||||
// 绑定仍然正确
|
||||
if pid := PlatformIDOf(ctx, id); pid != "pi-sess-x" {
|
||||
t.Errorf("platform_id = %q", pid)
|
||||
}
|
||||
}
|
||||
215
server/internal/repo/agent_calendar_test.go
Normal file
215
server/internal/repo/agent_calendar_test.go
Normal file
@ -0,0 +1,215 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
)
|
||||
|
||||
// Agent 只能看到自己建的日程。别人的日程里可能有它无权知道的会议与地址。
|
||||
func TestListCalendarEventsCreatedBy(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
mine := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "pi 自己建的", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi",
|
||||
})
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "dsh 建的", EventTime: time.Now().Add(time.Hour), CreatedBy: "dsh",
|
||||
})
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "人建的", EventTime: time.Now().Add(time.Hour), CreatedBy: "jianf",
|
||||
})
|
||||
|
||||
from := time.Now().Add(-time.Hour)
|
||||
to := time.Now().AddDate(0, 1, 0)
|
||||
|
||||
got, err := ListCalendarEventsCreatedBy(ctx, "pi", from, to, "active")
|
||||
if err != nil {
|
||||
t.Fatalf("列出: %v", err)
|
||||
}
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("pi 应只看到 1 条，得到 %d", len(got))
|
||||
}
|
||||
if got[0].EventID != mine.EventID {
|
||||
t.Errorf("看到了别人的事件：%s", got[0].Title)
|
||||
}
|
||||
|
||||
// 没建过任何事件的 Agent 得到空数组而不是 nil（nil 序列化成 null 前端会崩）
|
||||
empty, err := ListCalendarEventsCreatedBy(ctx, "opencode", from, to, "active")
|
||||
if err != nil {
|
||||
t.Fatalf("列出: %v", err)
|
||||
}
|
||||
if empty == nil {
|
||||
t.Error("应返回空数组而不是 nil")
|
||||
}
|
||||
if len(empty) != 0 {
|
||||
t.Errorf("应为空，得到 %d 条", len(empty))
|
||||
}
|
||||
}
|
||||
|
||||
// 「发给我但不是我建的」同样不返回：那些事件的编辑权不属于我，
|
||||
// 列出来只会让模型试图改它然后拿到 404。
|
||||
func TestListCalendarEventsCreatedByIgnoresRecipient(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "dsh 建的、发给 pi",
|
||||
EventTime: time.Now().Add(time.Hour),
|
||||
CreatedBy: "dsh",
|
||||
Recipients: []string{"pi"},
|
||||
})
|
||||
|
||||
got, err := ListCalendarEventsCreatedBy(ctx, "pi",
|
||||
time.Now().Add(-time.Hour), time.Now().AddDate(0, 1, 0), "active")
|
||||
if err != nil {
|
||||
t.Fatalf("列出: %v", err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Errorf("收件人不等于创建者，不该出现在列表里（得到 %d 条）", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestListCalendarEventsCreatedByStatusFilter(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "生效中", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi", Status: "active",
|
||||
})
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "已暂停", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi", Status: "paused",
|
||||
})
|
||||
|
||||
from := time.Now().Add(-time.Hour)
|
||||
to := time.Now().AddDate(0, 1, 0)
|
||||
|
||||
if got, _ := ListCalendarEventsCreatedBy(ctx, "pi", from, to, "active"); len(got) != 1 {
|
||||
t.Errorf("active 过滤应给 1 条，得到 %d", len(got))
|
||||
}
|
||||
// 空串 = 不过滤（handler 里 status=all 映射成空串）
|
||||
if got, _ := ListCalendarEventsCreatedBy(ctx, "pi", from, to, ""); len(got) != 2 {
|
||||
t.Errorf("不过滤应给 2 条，得到 %d", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
// 总量上限的依据。速率限制压不住「每小时建 19 条连建一周」，
|
||||
// 而日历事件是长效的 —— 攒下来的每条都持续产生提醒。
|
||||
func TestCountActiveEventsBy(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "生效", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi", Status: "active",
|
||||
})
|
||||
}
|
||||
// cancelled 与 paused 不该计入「生效中」
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "取消了", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi", Status: "cancelled",
|
||||
})
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "暂停了", EventTime: time.Now().Add(time.Hour), CreatedBy: "pi", Status: "paused",
|
||||
})
|
||||
seedEvent(t, &models.CalendarEvent{
|
||||
Title: "别人的", EventTime: time.Now().Add(time.Hour), CreatedBy: "dsh", Status: "active",
|
||||
})
|
||||
|
||||
n, err := CountActiveEventsBy(ctx, "pi")
|
||||
if err != nil {
|
||||
t.Fatalf("计数: %v", err)
|
||||
}
|
||||
if n != 3 {
|
||||
t.Errorf("pi 的生效事件应为 3，得到 %d", n)
|
||||
}
|
||||
|
||||
if n, _ := CountActiveEventsBy(ctx, "从来没建过"); n != 0 {
|
||||
t.Errorf("没建过应为 0，得到 %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 速率限制 ───
|
||||
|
||||
func TestAllowAgentCalendarEvent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
limit := CalendarRateLimit()
|
||||
for i := 0; i < limit; i++ {
|
||||
ok, _ := AllowAgentCalendarEvent(ctx, "pi")
|
||||
if !ok {
|
||||
t.Fatalf("第 %d 次（上限 %d）就被拒了", i+1, limit)
|
||||
}
|
||||
}
|
||||
ok, retry := AllowAgentCalendarEvent(ctx, "pi")
|
||||
if ok {
|
||||
t.Error("超过上限应被拒")
|
||||
}
|
||||
if retry <= 0 {
|
||||
t.Errorf("被拒时应给出 retryAfter，得到 %d", retry)
|
||||
}
|
||||
}
|
||||
|
||||
// 日历桶与新建会话桶必须独立：建满 20 条日程不该连带堵住新建会话。
|
||||
func TestCalendarRateBucketIsSeparateFromSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < CalendarRateLimit(); i++ {
|
||||
AllowAgentCalendarEvent(ctx, "pi")
|
||||
}
|
||||
if ok, _ := AllowAgentCalendarEvent(ctx, "pi"); ok {
|
||||
t.Fatal("准备阶段：日历桶应已满")
|
||||
}
|
||||
// 新建会话桶应完全不受影响
|
||||
if ok, _ := AllowNewSession(ctx, "pi"); !ok {
|
||||
t.Error("日历桶满不该堵住新建会话 —— 两个桶必须独立")
|
||||
}
|
||||
}
|
||||
|
||||
// 不同 Agent 的桶互不干扰。
|
||||
func TestCalendarRateBucketPerAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < CalendarRateLimit(); i++ {
|
||||
AllowAgentCalendarEvent(ctx, "pi")
|
||||
}
|
||||
if ok, _ := AllowAgentCalendarEvent(ctx, "dsh"); !ok {
|
||||
t.Error("pi 建满不该影响 dsh")
|
||||
}
|
||||
}
|
||||
|
||||
// 创建失败要归还名额：那次创建实际没有发生。
|
||||
func TestReleaseAgentCalendarEvent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < CalendarRateLimit(); i++ {
|
||||
AllowAgentCalendarEvent(ctx, "pi")
|
||||
}
|
||||
if ok, _ := AllowAgentCalendarEvent(ctx, "pi"); ok {
|
||||
t.Fatal("准备阶段：应已满")
|
||||
}
|
||||
// 归还一个（模拟刚才那次被拒之前的失败创建）
|
||||
ReleaseAgentCalendarEvent(ctx, "pi")
|
||||
if ok, _ := AllowAgentCalendarEvent(ctx, "pi"); !ok {
|
||||
t.Error("归还名额后应能再建一条")
|
||||
}
|
||||
}
|
||||
|
||||
// 空 Agent 名放行且不记账：这条路径只在鉴权已经失败时才可能走到，
|
||||
// 记账会污染桶（bucket 名变成 "calendar:"）。
|
||||
func TestCalendarRateEmptyAgentPassesThrough(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < CalendarRateLimit()+5; i++ {
|
||||
if ok, _ := AllowAgentCalendarEvent(ctx, ""); !ok {
|
||||
t.Fatal("空 Agent 名应一律放行")
|
||||
}
|
||||
}
|
||||
}
|
||||
253
server/internal/repo/agent_disable_test.go
Normal file
253
server/internal/repo/agent_disable_test.go
Normal file
@ -0,0 +1,253 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 停用是可逆的「归档」，不是删除。这组测试钉住三件事：
|
||||
// 停用后从候选里消失、密钥被撤销、重新注册不能复活它。
|
||||
|
||||
func TestSetAgentDisabledHidesFromCandidates(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for _, n := range []string{"keeper", "goner"} {
|
||||
if err := CreateOrUpdateAgent(ctx, n, "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册 %s: %v", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "goner", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 默认列表（地址补全、GET /agents、可授权范围都走这条）不含已停用的
|
||||
got, err := ListAgents(ctx, "")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgents: %v", err)
|
||||
}
|
||||
names := map[string]bool{}
|
||||
for _, a := range got {
|
||||
names[a.Name] = true
|
||||
}
|
||||
if names["goner"] {
|
||||
t.Error("已停用的 Agent 仍出现在默认列表里 —— 人会把任务派给一个不会响应的地址")
|
||||
}
|
||||
if !names["keeper"] {
|
||||
t.Error("停用一个把别的也弄没了")
|
||||
}
|
||||
|
||||
// statusFilter="all" 时要能看到 —— 那是管理页恢复它的唯一入口
|
||||
all, err := ListAgents(ctx, "all")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgents(all): %v", err)
|
||||
}
|
||||
found := false
|
||||
for _, a := range all {
|
||||
if a.Name == "goner" {
|
||||
found = true
|
||||
if a.Status != "disabled" {
|
||||
t.Errorf("状态应为 disabled，实际 %q", a.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("statusFilter=all 也看不到已停用的，就再也无法恢复它了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAgentDisabledRevokesKeys(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
admin := seedAdminForTest(t, ctx)
|
||||
for i := 0; i < 2; i++ {
|
||||
if _, err := CreateAgentKey(ctx, "bot", "permanent", "k", 0, admin, ""); err != nil {
|
||||
t.Fatalf("建密钥: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
revoked, err := SetAgentDisabled(ctx, "bot", true)
|
||||
if err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
if revoked != 2 {
|
||||
t.Errorf("应撤销 2 把密钥，实际 %d", revoked)
|
||||
}
|
||||
|
||||
keys, err := ListAgentKeys(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAgentKeys: %v", err)
|
||||
}
|
||||
if len(keys) != 0 {
|
||||
t.Errorf("停用后仍留着 %d 把密钥 —— 插件还能用它调 /mail/send，"+
|
||||
"停用的语义是「不再参与工作」而不只是「不出现在补全里」", len(keys))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDisabledAgentCannotReRegister(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("首次注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 插件启动时会重新注册。不拒的话 status 被写回 online，停用等于没做。
|
||||
err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil)
|
||||
if !errors.Is(err, ErrAgentDisabled) {
|
||||
t.Fatalf("已停用的 Agent 重新注册应当被拒，实际 err=%v", err)
|
||||
}
|
||||
|
||||
disabled, err := AgentDisabled(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("AgentDisabled: %v", err)
|
||||
}
|
||||
if !disabled {
|
||||
t.Error("注册尝试把停用状态冲掉了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeartbeatDoesNotReviveDisabledAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
// 心跳是 30 秒一次的。不排除 disabled 的话停用最多维持半分钟。
|
||||
if _, err := HeartbeatAgent(ctx, "bot"); err != nil {
|
||||
t.Fatalf("心跳本身不该报错: %v", err)
|
||||
}
|
||||
|
||||
disabled, _ := AgentDisabled(ctx, "bot")
|
||||
if !disabled {
|
||||
t.Error("心跳把已停用的 Agent 改回在线了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRestoreAgentGoesOfflineNotOnline(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "bot", false); err != nil {
|
||||
t.Fatalf("恢复: %v", err)
|
||||
}
|
||||
|
||||
all, _ := ListAgents(ctx, "all")
|
||||
for _, a := range all {
|
||||
if a.Name != "bot" {
|
||||
continue
|
||||
}
|
||||
// 恢复成 online 会让界面显示一个其实没在跑的 Agent 为在线；
|
||||
// 它是否真的活着由下一次心跳决定。
|
||||
if a.Status != "offline" {
|
||||
t.Errorf("恢复后应为 offline，实际 %q", a.Status)
|
||||
}
|
||||
}
|
||||
|
||||
// 恢复后能重新注册
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Errorf("恢复后应当能重新注册: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAgentDisabledUnknownAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := SetAgentDisabled(ctx, "nope", true)
|
||||
if !errors.Is(err, sql.ErrNoRows) {
|
||||
t.Errorf("停用不存在的 Agent 应回 ErrNoRows，实际 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 停用不得动邮件与会话 —— 往来里有一半是人自己写的。
|
||||
func TestSetAgentDisabledKeepsMailAndSessions(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
sid, err := CreateSession(ctx, nil, "bot", "一件事", "/tmp/ws")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, sid, nil,
|
||||
"human", "", "bot", "/tmp/ws", "主题", "正文", nil); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
mails, err := ListInbox(ctx, "bot", "all", 10)
|
||||
if err != nil {
|
||||
t.Fatalf("ListInbox: %v", err)
|
||||
}
|
||||
if len(mails) != 1 {
|
||||
t.Errorf("停用把邮件删了：剩 %d 封。那些往来里有一半是人自己写的", len(mails))
|
||||
}
|
||||
}
|
||||
|
||||
// 模型范围与平台会话镜像也保留：恢复后不必重配。
|
||||
func TestSetAgentDisabledKeepsModelScope(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := CreateOrUpdateAgent(ctx, "bot", "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册: %v", err)
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "bot", []ModelRef{{Provider: "p", Model: "m"}}); err != nil {
|
||||
t.Fatalf("设范围: %v", err)
|
||||
}
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "bot", true); err != nil {
|
||||
t.Fatalf("停用: %v", err)
|
||||
}
|
||||
|
||||
allowed, err := ListAllowedModels(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAllowedModels: %v", err)
|
||||
}
|
||||
if len(allowed) != 1 {
|
||||
t.Errorf("停用把模型范围清了，恢复后管理员得重配一遍：%+v", allowed)
|
||||
}
|
||||
}
|
||||
|
||||
// seedAdminForTest 插一个管理员并返回它的 user_id（CreateAgentKey 要 created_by）。
|
||||
func seedAdminForTest(t *testing.T, ctx context.Context) uuid.UUID {
|
||||
t.Helper()
|
||||
var id uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ('key-admin', 'Admin', 'x', 'admin') RETURNING user_id`).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatalf("seed admin: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
381
server/internal/repo/attachments.go
Normal file
381
server/internal/repo/attachments.go
Normal file
@ -0,0 +1,381 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 附件 ----------
|
||||
//
|
||||
// 元数据在库、内容在磁盘（internal/blob）。两者的一致性由调用顺序保证：
|
||||
// 先落盘再入库 —— 反过来会出现「库里有记录但文件不存在」的下载 500。
|
||||
// 落盘成功但入库失败时最多留下一个无引用的文件，由 GC 回收，不影响正确性。
|
||||
|
||||
var (
|
||||
// ErrAttachmentNotFound 附件不存在
|
||||
ErrAttachmentNotFound = errors.New("attachment not found")
|
||||
// ErrAttachmentNotOwned 附件不属于该上传者
|
||||
ErrAttachmentNotOwned = errors.New("attachment not owned by uploader")
|
||||
// ErrAttachmentAlreadyAttached 附件已挂到别的邮件上
|
||||
ErrAttachmentAlreadyAttached = errors.New("attachment already attached")
|
||||
)
|
||||
|
||||
const attachmentCols = `attachment_id, mail_id, uploader, filename, content_type, size_bytes, sha256, created_at`
|
||||
|
||||
func scanAttachment(sc interface{ Scan(...any) error }) (*models.Attachment, error) {
|
||||
var a models.Attachment
|
||||
if err := sc.Scan(&a.ID, &a.MailID, &a.Uploader, &a.Filename,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
// CreateAttachment 登记一条待挂载的附件（mail_id 为空）。
|
||||
func CreateAttachment(ctx context.Context, uploader, filename, contentType string, size int64, sum string) (*models.Attachment, error) {
|
||||
a := &models.Attachment{
|
||||
Uploader: uploader,
|
||||
Filename: filename,
|
||||
ContentType: contentType,
|
||||
SizeBytes: size,
|
||||
SHA256: sum,
|
||||
}
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO attachments (uploader, filename, content_type, size_bytes, sha256)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
RETURNING attachment_id, created_at
|
||||
`, uploader, filename, contentType, size, sum).Scan(&a.ID, &a.CreatedAt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// GetAttachment 读取一条附件元数据。
|
||||
func GetAttachment(ctx context.Context, id uuid.UUID) (*models.Attachment, error) {
|
||||
a, err := scanAttachment(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+attachmentCols+` FROM attachments WHERE attachment_id = $1`, id))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrAttachmentNotFound
|
||||
}
|
||||
return a, err
|
||||
}
|
||||
|
||||
// ListAttachmentsFor 列出某封邮件的附件。
|
||||
func ListAttachmentsFor(ctx context.Context, mailID uuid.UUID) ([]models.Attachment, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT `+attachmentCols+` FROM attachments WHERE mail_id = $1 ORDER BY created_at`, mailID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []models.Attachment{}
|
||||
for rows.Next() {
|
||||
a, err := scanAttachment(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, *a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// EnsureAttachable 只做**读取校验**：这批附件是否存在、属于该上传者、且尚未挂载。
|
||||
//
|
||||
// # 为什么要有一个「只查不改」的版本
|
||||
//
|
||||
// 原先只有 AttachToMail，而它在 CreateMail **之后**调用。于是附件不合法时
|
||||
// （不属于我 / 已随别的邮件发出）请求返回 403/409，但那封邮件**已经入库、已经
|
||||
// 通知了收件人、已经扣掉了会话预算**。实测两封探针邮件（403 与 409）都躺在库里，
|
||||
// used_rounds 也涨了。发件方看到 4xx 会重试，收件方于是收到两封。
|
||||
//
|
||||
// 纯输入校验必须在产生任何副作用之前做完 —— 与「400 之后会话已建好」是同一个教训。
|
||||
//
|
||||
// 它不能取代 AttachToMail 里的原子判断：两次调用之间仍有竞态窗口
|
||||
// （另一个请求把同一个附件挂走了）。那条路径靠调用方回滚，见 handler.attachAll。
|
||||
func EnsureAttachable(ctx context.Context, ids []uuid.UUID, uploader string) error {
|
||||
for _, id := range ids {
|
||||
a, err := GetAttachment(ctx, id)
|
||||
if err != nil {
|
||||
return err // ErrAttachmentNotFound 或库错误
|
||||
}
|
||||
if a.Uploader != uploader {
|
||||
return ErrAttachmentNotOwned
|
||||
}
|
||||
if a.MailID != nil {
|
||||
return ErrAttachmentAlreadyAttached
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AttachToMail 把一批待挂载附件绑到某封邮件上。
|
||||
//
|
||||
// 每条都要求：存在、属于该上传者、且尚未挂载。
|
||||
// 用 WHERE mail_id IS NULL AND uploader = ? 一条 UPDATE 完成判断与写入，
|
||||
// 避免「先查后改」在并发下把同一个附件挂到两封邮件上。
|
||||
func AttachToMail(ctx context.Context, mailID uuid.UUID, ids []uuid.UUID, uploader string) error {
|
||||
for _, id := range ids {
|
||||
tag, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE attachments SET mail_id = $1
|
||||
WHERE attachment_id = $2 AND uploader = $3 AND mail_id IS NULL
|
||||
`, mailID, id, uploader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// 没改到：查明原因，给调用方一个能照着修的错误
|
||||
a, gErr := GetAttachment(ctx, id)
|
||||
if gErr != nil {
|
||||
return gErr
|
||||
}
|
||||
if a.Uploader != uploader {
|
||||
return ErrAttachmentNotOwned
|
||||
}
|
||||
return ErrAttachmentAlreadyAttached
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CopyAttachmentsTo 把源邮件的附件复制到目标邮件（转发时用）。
|
||||
//
|
||||
// 内容寻址下「复制」只是新增一条指向同一 sha256 的元数据，不拷磁盘文件。
|
||||
// uploader 记为转发人：附件随新邮件重新分发，其可见范围由新邮件的参与方决定，
|
||||
// 而不是沿用原上传者。返回复制的数量。
|
||||
func CopyAttachmentsTo(ctx context.Context, srcMailID, dstMailID uuid.UUID, forwarder string) (int, error) {
|
||||
src, err := ListAttachmentsFor(ctx, srcMailID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, a := range src {
|
||||
_, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO attachments (mail_id, uploader, filename, content_type, size_bytes, sha256)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
`, dstMailID, forwarder, a.Filename, a.ContentType, a.SizeBytes, a.SHA256)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
return len(src), nil
|
||||
}
|
||||
|
||||
// DeleteAttachment 删除一条附件元数据，返回它的 sha256 以及该内容是否已无人引用。
|
||||
// 内容寻址下多条记录可能共享同一个文件，只有最后一条引用消失才能删磁盘文件。
|
||||
func DeleteAttachment(ctx context.Context, id uuid.UUID) (sum string, orphaned bool, err error) {
|
||||
a, err := GetAttachment(ctx, id)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
if _, err = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM attachments WHERE attachment_id = $1`, id); err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
|
||||
var refs int
|
||||
if err = db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, a.SHA256).Scan(&refs); err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
return a.SHA256, refs == 0, nil
|
||||
}
|
||||
|
||||
// SweepOrphanAttachments 清理超过 age 仍未挂载到邮件的附件记录，
|
||||
// 返回可以从磁盘删除的 sha256 列表（已确认无任何记录引用）。
|
||||
//
|
||||
// 上传后没走完发信流程（用户取消、Agent 崩溃）会留下这类记录，
|
||||
// 不清理的话磁盘只会单调增长。
|
||||
func SweepOrphanAttachments(ctx context.Context, age time.Duration) ([]string, error) {
|
||||
cutoff := time.Now().Add(-age)
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT attachment_id, sha256 FROM attachments
|
||||
WHERE mail_id IS NULL AND created_at < $1`, cutoff)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
type orphan struct {
|
||||
id uuid.UUID
|
||||
sum string
|
||||
}
|
||||
var found []orphan
|
||||
for rows.Next() {
|
||||
var o orphan
|
||||
if err := rows.Scan(&o.id, &o.sum); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
found = append(found, o)
|
||||
}
|
||||
rows.Close()
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var removable []string
|
||||
for _, o := range found {
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM attachments WHERE attachment_id = $1`, o.id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var refs int
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, o.sum).Scan(&refs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if refs == 0 {
|
||||
removable = append(removable, o.sum)
|
||||
}
|
||||
}
|
||||
return removable, nil
|
||||
}
|
||||
|
||||
// SweepUnreferencedBlobs 删掉磁盘上没有任何库记录指向的内容文件。
|
||||
//
|
||||
// # 为什么 SweepOrphanAttachments 不够
|
||||
//
|
||||
// 那个函数走的是 `SELECT … FROM attachments WHERE mail_id IS NULL` —— 它只能看见
|
||||
// **库里还有记录**的孤儿。一旦记录本身消失（清库、手工 DELETE、迁移），
|
||||
// 对应的文件就永远脱离了 GC 的视野：本机实测磁盘 8 个 blob 里 7 个没有任何库记录，
|
||||
// 全部来自 09-03 那次清库，之后一直躺在那里。
|
||||
//
|
||||
// 这个反向清理从**磁盘**出发：枚举全部内容文件，凡是 attachments 与
|
||||
// calendar_attachments 都不引用的就删。返回删掉的数量。
|
||||
//
|
||||
// # 为什么要 minAge
|
||||
//
|
||||
// 上传是「先落盘、再入库」（顺序不能反，否则会出现「库里有记录、磁盘没文件」的
|
||||
// 下载 500）。那两步之间有一个窗口，此刻文件确实没有任何库记录 —— 不设年龄下限
|
||||
// 会把正在上传的文件删掉。取一个远大于单次上传耗时的值。
|
||||
func SweepUnreferencedBlobs(ctx context.Context, blobs BlobLister, minAge time.Duration) (int, error) {
|
||||
if blobs == nil {
|
||||
return 0, nil
|
||||
}
|
||||
sums, err := blobs.List()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(sums) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// 一次查回全部被引用的 sha256。逐个文件查一次库是 N 次往返，
|
||||
// 而这两张表加起来通常只有几百行。
|
||||
referenced := map[string]struct{}{}
|
||||
for _, q := range []string{
|
||||
`SELECT sha256 FROM attachments`,
|
||||
`SELECT sha256 FROM calendar_attachments`,
|
||||
} {
|
||||
rows, qErr := db.DB.QueryContext(ctx, q)
|
||||
if qErr != nil {
|
||||
return 0, qErr
|
||||
}
|
||||
for rows.Next() {
|
||||
var s string
|
||||
if sErr := rows.Scan(&s); sErr != nil {
|
||||
rows.Close()
|
||||
return 0, sErr
|
||||
}
|
||||
referenced[s] = struct{}{}
|
||||
}
|
||||
rows.Close()
|
||||
if rErr := rows.Err(); rErr != nil {
|
||||
return 0, rErr
|
||||
}
|
||||
}
|
||||
|
||||
cutoff := time.Now().Add(-minAge)
|
||||
removed := 0
|
||||
for sum, mod := range sums {
|
||||
if _, ok := referenced[sum]; ok {
|
||||
continue
|
||||
}
|
||||
if mod.After(cutoff) {
|
||||
continue // 可能正在上传（落盘与入库之间的窗口）
|
||||
}
|
||||
if rErr := blobs.Remove(sum); rErr != nil {
|
||||
continue // 删不掉就下一轮再试，不该让整次清理中断
|
||||
}
|
||||
removed++
|
||||
}
|
||||
return removed, nil
|
||||
}
|
||||
|
||||
// BlobLister 是 SweepUnreferencedBlobs 需要的存储能力。
|
||||
//
|
||||
// 用 map[string]time.Time 而不是自定义结构体：那样 blob 包就不必 import repo
|
||||
// （底层存储依赖上层仓储会很怪），而 Go 的接口是结构化匹配的，签名一致即可。
|
||||
type BlobLister interface {
|
||||
// List 返回 sha256 → 该内容文件的修改时间。
|
||||
List() (map[string]time.Time, error)
|
||||
Remove(sum string) error
|
||||
}
|
||||
|
||||
// AttachmentAccessible 判断某人是否有权读取某附件：
|
||||
// 已挂载的看邮件所属会话的参与关系，未挂载的只有上传者本人能看。
|
||||
func AttachmentAccessible(ctx context.Context, a *models.Attachment, name string) (bool, error) {
|
||||
if a.MailID == nil {
|
||||
return a.Uploader == name, nil
|
||||
}
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*) FROM mails m
|
||||
WHERE m.mail_id = $1
|
||||
AND (m.from_name = $2 OR m.to_name = $2 OR `+db.CCHas("m.cc_list", 2)+`)
|
||||
`, *a.MailID, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// ListAttachmentsForMails 批量取多封邮件的附件，返回 mail_id → 附件列表。
|
||||
//
|
||||
// 为什么要批量：会话线程与收发件箱都是「一批邮件」，逐封调 ListAttachmentsFor
|
||||
// 就是 N+1 —— 一个 200 封的会话打开一次要打 200 次库。
|
||||
// 用 IN (...) 一次取回后在内存里分组。
|
||||
//
|
||||
// 占位符手工拼而非用数组参数：SQLite 驱动不支持 PG 的 = ANY($1)，
|
||||
// 而这里的元素是已解析的 uuid.UUID，不存在注入面。
|
||||
func ListAttachmentsForMails(ctx context.Context, mailIDs []uuid.UUID) (map[uuid.UUID][]models.Attachment, error) {
|
||||
out := map[uuid.UUID][]models.Attachment{}
|
||||
if len(mailIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
ph := make([]string, len(mailIDs))
|
||||
args := make([]any, len(mailIDs))
|
||||
for i, id := range mailIDs {
|
||||
ph[i] = fmt.Sprintf("$%d", i+1)
|
||||
args[i] = id
|
||||
}
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT `+attachmentCols+` FROM attachments
|
||||
WHERE mail_id IN (`+strings.Join(ph, ",")+`)
|
||||
ORDER BY created_at`, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
a, err := scanAttachment(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if a.MailID == nil {
|
||||
continue // WHERE 已排除，只是防御
|
||||
}
|
||||
out[*a.MailID] = append(out[*a.MailID], *a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
142
server/internal/repo/attachments_test.go
Normal file
142
server/internal/repo/attachments_test.go
Normal file
@ -0,0 +1,142 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// seedClock 给测试数据发严格递增的时间戳。
|
||||
//
|
||||
// 不靠挂钟：测试在一个循环里连插几封，很可能落在同一毫秒里，
|
||||
// 于是「会话里最早/最后那封」的排序由 mail_id（随机 UUID）决定 —— 结果随机。
|
||||
// 生产里两封邮件至少隔着一次模型推理，同毫秒撞车不现实；
|
||||
// 但测试必须确定，所以显式发号。
|
||||
var seedClock = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
func nextSeedTime() string {
|
||||
seedClock = seedClock.Add(time.Second)
|
||||
return seedClock.Format("2006-01-02 15:04:05.000")
|
||||
}
|
||||
|
||||
// seedMailIn 在指定会话里插一封邮件，返回其 id。
|
||||
// 时间戳严格递增，因此调用顺序就是邮件的先后顺序。
|
||||
func seedMailIn(t *testing.T, sessionID uuid.UUID, from, to, 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, 'body', '[]', $5)
|
||||
RETURNING mail_id
|
||||
`, sessionID, from, to, subject, nextSeedTime()).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatalf("seed mail: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func seedSessionRow(t *testing.T, alias string) uuid.UUID {
|
||||
t.Helper()
|
||||
var id uuid.UUID
|
||||
err := db.DB.QueryRowContext(context.Background(), `
|
||||
INSERT INTO sessions (from_agent, subject, session_alias)
|
||||
VALUES ('opencode', 'attach test', $1)
|
||||
RETURNING session_id
|
||||
`, alias).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatalf("seed session: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func attach(t *testing.T, mailID uuid.UUID, name, sum string) {
|
||||
t.Helper()
|
||||
a, err := CreateAttachment(context.Background(), "admin", name, "text/plain", 3, sum)
|
||||
if err != nil {
|
||||
t.Fatalf("create attachment: %v", err)
|
||||
}
|
||||
if err := AttachToMail(context.Background(), mailID, []uuid.UUID{a.ID}, "admin"); err != nil {
|
||||
t.Fatalf("attach: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ListAttachmentsForMails 存在的理由是消掉 N+1：
|
||||
// 原先每封邮件单独查一次，一个 200 封的会话打开要打 200 次库。
|
||||
func TestListAttachmentsForMails(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "batch-attach")
|
||||
|
||||
m1 := seedMailIn(t, sid, "admin", "opencode", "两个附件")
|
||||
m2 := seedMailIn(t, sid, "opencode", "admin", "一个附件")
|
||||
m3 := seedMailIn(t, sid, "admin", "opencode", "没有附件")
|
||||
|
||||
attach(t, m1, "a.txt", "sum-a")
|
||||
attach(t, m1, "b.txt", "sum-b")
|
||||
attach(t, m2, "c.txt", "sum-c")
|
||||
|
||||
got, err := ListAttachmentsForMails(context.Background(),
|
||||
[]uuid.UUID{m1, m2, m3})
|
||||
if err != nil {
|
||||
t.Fatalf("批量查询失败: %v", err)
|
||||
}
|
||||
|
||||
if n := len(got[m1]); n != 2 {
|
||||
t.Errorf("m1 应有 2 个附件，实际 %d", n)
|
||||
}
|
||||
if n := len(got[m2]); n != 1 {
|
||||
t.Errorf("m2 应有 1 个附件，实际 %d", n)
|
||||
}
|
||||
// 无附件的邮件不该出现在 map 里：调用方据此保持 Attachments 为 nil，
|
||||
// 这样带 omitempty 的字段不会给每封邮件的 JSON 白加一个 "attachments":[]
|
||||
if _, ok := got[m3]; ok {
|
||||
t.Errorf("m3 无附件却出现在结果里：%#v", got[m3])
|
||||
}
|
||||
|
||||
// 同一封内按 created_at 排序，顺序不能乱
|
||||
if len(got[m1]) == 2 && got[m1][0].Filename != "a.txt" {
|
||||
t.Errorf("同一封内应按上传顺序，首个是 %s", got[m1][0].Filename)
|
||||
}
|
||||
|
||||
// 每条都要带回 mail_id，否则调用方分不清是谁的
|
||||
for _, a := range got[m1] {
|
||||
if a.MailID == nil || *a.MailID != m1 {
|
||||
t.Errorf("附件 %s 的 mail_id 不对：%v", a.Filename, a.MailID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 空输入必须返回空 map 而非 nil：调用方直接索引不该 panic。
|
||||
func TestListAttachmentsForMailsEmpty(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
|
||||
for _, ids := range [][]uuid.UUID{nil, {}} {
|
||||
got, err := ListAttachmentsForMails(context.Background(), ids)
|
||||
if err != nil {
|
||||
t.Fatalf("空输入不该报错: %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatal("空输入应返回空 map 而非 nil")
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Errorf("空输入应返回空结果，实际 %d 项", len(got))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 传入不存在的 mail_id 不该报错，只是查不到 —— 调用方可能拿着已删邮件的 id。
|
||||
func TestListAttachmentsForMailsUnknownID(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
|
||||
got, err := ListAttachmentsForMails(context.Background(),
|
||||
[]uuid.UUID{uuid.New(), uuid.New()})
|
||||
if err != nil {
|
||||
t.Fatalf("未知 id 不该报错: %v", err)
|
||||
}
|
||||
if len(got) != 0 {
|
||||
t.Errorf("未知 id 应查不到，实际 %d 项", len(got))
|
||||
}
|
||||
}
|
||||
177
server/internal/repo/autoalias.go
Normal file
177
server/internal/repo/autoalias.go
Normal file
@ -0,0 +1,177 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 自动别名 —— 让 `.new` 建出来的会话立刻可被寻址。
|
||||
//
|
||||
// # 为什么必须自动命名
|
||||
//
|
||||
// `session` 位三态里 `new` 是**一次性动作**：它建出会话就用完了。之后要再投进
|
||||
// 同一条会话，只有两条路 —— `reply_to` 某封具体邮件，或者 `name@path.<别名>`。
|
||||
// 而 `CreateSession(alias=nil)` 建出来的会话别名是 NULL，于是：
|
||||
//
|
||||
// - `FindNamedSessionFor` 查不到它（`WHERE session_alias = $1` 对 NULL 不成立）
|
||||
// - `SuggestSessionCandidates` 跳过它（`session_alias IS NOT NULL AND <> ''`）
|
||||
// - 参与方拿到的 `new_mail` 里 `session_alias` 是空串
|
||||
//
|
||||
// 结果是：被抄送方收到一封 `x@/p.new` 的邮件，**除了回复那一封之外无法再投进这条
|
||||
// 会话**。再发一次 `x@/p.new` 只会建第三条会话。这不是能力缺失，是寻址断链。
|
||||
//
|
||||
// 原先的设计假定平台插件会通过 `POST /sessions/{id}/sync` 把模型生成的标题回写成
|
||||
// 别名，于是「未命名」只是短暂状态。但两件事让这个假定不成立：
|
||||
//
|
||||
// 1. 人类发的邮件根本没有平台侧，永远等不到回写；
|
||||
// 2. 回写发生在模型跑完第一轮之后，而抄送方**在那之前**就要决定回信地址。
|
||||
//
|
||||
// 因此本侧先给一个可用的别名，平台随后仍可用 `SyncSessionAlias` 改写它 ——
|
||||
// `alias_source` 保持 `platform` 正是为此：自动名不是人定的名，不该挡住平台命名。
|
||||
//
|
||||
// # 为什么不复用 SyncSessionAlias
|
||||
//
|
||||
// 那个函数假定「会话已存在、现在要改名」，并且会跳过 `manual`。这里的场景是
|
||||
// 「刚建完、还没有名字」，且必须在**建会话的同一个请求里**完成，否则中间那一瞬
|
||||
// 发出的 SSE 仍然带空别名。
|
||||
|
||||
// aliasMaxBytes 与 normalizeAlias 的截断上限一致（sessions.session_alias 为 VARCHAR(128)）。
|
||||
const aliasMaxBytes = 128
|
||||
|
||||
// autoAliasAttempts 是撞名后追加 -2、-3… 的尝试次数上限。
|
||||
// 与 SyncSessionAlias 取同一个数量级：同一主题在同一天内开几十条会话已属异常，
|
||||
// 真到了上限说明调用方在刷会话，此时报错比继续找空位更有价值。
|
||||
const autoAliasAttempts = 50
|
||||
|
||||
// AutoAliasFor 依据收件人与主题拼一个候选别名（未做唯一性检查）。
|
||||
//
|
||||
// 形如 `dsh-重构导入路径`：前缀用收件方名字，后缀用主题。**两者都要**——
|
||||
// 只用主题时「服务恢复验证」这类通用主题会在不同 Agent 之间反复撞名，
|
||||
// 只用名字则同一个 Agent 的所有会话都叫 `dsh-2`、`dsh-3`，看不出在聊什么。
|
||||
//
|
||||
// 主题为空（少见但合法）时退回单独的名字，由调用方靠后缀去重。
|
||||
func AutoAliasFor(toName, subject string) string {
|
||||
base := sanitizeAliasPart(toName)
|
||||
topic := sanitizeAliasPart(subject)
|
||||
|
||||
switch {
|
||||
case base == "" && topic == "":
|
||||
// 两边都拿不出可用字符（例如主题全是标点、名字为空）。
|
||||
// 返回空串让调用方走随机兜底，不要在这里编造。
|
||||
return ""
|
||||
case base == "":
|
||||
return truncateAlias(topic)
|
||||
case topic == "":
|
||||
return truncateAlias(base)
|
||||
default:
|
||||
return truncateAlias(base + "-" + topic)
|
||||
}
|
||||
}
|
||||
|
||||
// EnsureSessionAlias 保证会话拥有一个可寻址的别名，返回最终别名。
|
||||
//
|
||||
// 已有别名时原样返回，不做任何写入 —— 这让它可以被无条件调用，
|
||||
// 包括「默认会话」路径上那条可能是刚建的、也可能是复用的会话。
|
||||
//
|
||||
// 撞名时追加 -2、-3… 后缀；`want` 为空或全部被占用时退回
|
||||
// `session-<uuid 前 8 位>`：一个能寻址的丑名字，远胜于没有名字。
|
||||
func EnsureSessionAlias(ctx context.Context, id uuid.UUID, want string) (string, error) {
|
||||
if cur := SessionAliasOf(ctx, id); cur != "" {
|
||||
return cur, nil
|
||||
}
|
||||
|
||||
cands := make([]string, 0, autoAliasAttempts+1)
|
||||
if want != "" {
|
||||
for i := 0; i < autoAliasAttempts; i++ {
|
||||
if i == 0 {
|
||||
cands = append(cands, want)
|
||||
continue
|
||||
}
|
||||
cands = append(cands, truncateAlias(fmt.Sprintf("%s-%d", want, i+1)))
|
||||
}
|
||||
}
|
||||
// 兜底：uuid 前 8 位。碰撞概率可忽略，且与 want 无关，
|
||||
// 因此即便主题里一个可用字符都没有也总能拿到别名。
|
||||
cands = append(cands, "session-"+id.String()[:8])
|
||||
|
||||
for _, c := range cands {
|
||||
// 条件写入：`session_alias IS NULL OR = ''` 保证并发下只有一方写成功，
|
||||
// 另一方 RowsAffected=0，随后重读拿到对方写的名字 ——
|
||||
// 两个请求都返回同一个别名，而不是各自以为自己命名成功。
|
||||
res, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET session_alias = $1, updated_at = NOW()
|
||||
WHERE session_id = $2 AND (session_alias IS NULL OR session_alias = '')`,
|
||||
c, id)
|
||||
if err != nil {
|
||||
if db.IsUniqueViolation(err) {
|
||||
continue // 别名被别的会话占了，试下一个后缀
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
// 期间别人（并发请求或平台同步）已经命名过，尊重那个名字
|
||||
if cur := SessionAliasOf(ctx, id); cur != "" {
|
||||
return cur, nil
|
||||
}
|
||||
// 写不进去且读不到名字，只可能是会话刚被删
|
||||
return "", fmt.Errorf("会话 %s 已不存在，无法分配别名", id)
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("别名 %q 连同 -2..-%d 后缀与 uuid 兜底均被占用", want, autoAliasAttempts)
|
||||
}
|
||||
|
||||
// sanitizeAliasPart 把任意文本压成别名可用的片段。
|
||||
//
|
||||
// 规则与 normalizeAlias 一致（非法字符换 -、压缩连续 -、去首尾 -），
|
||||
// 另外多做两件事：
|
||||
//
|
||||
// - **去掉 Markdown / 标点噪声**：主题里的 `[联调]`、`—`、`：` 变成一串
|
||||
// 破折号毫无信息量。只保留字母、数字与非标点的 Unicode 字符（中文、日文等）。
|
||||
// - **压缩空白**：`Re: 服务恢复验证` → `Re-服务恢复验证`，而不是 `Re--服务恢复验证`。
|
||||
//
|
||||
// 保留中文是刻意的：本项目的会话主题多为中文，转拼音需要额外依赖，
|
||||
// 而 `dsh-重构导入路径` 作为地址完全可用（三维寻址只忌 `. / @` 与空白）。
|
||||
func sanitizeAliasPart(s string) string {
|
||||
var b strings.Builder
|
||||
lastDash := false
|
||||
for _, r := range s {
|
||||
keep := unicode.IsLetter(r) || unicode.IsDigit(r)
|
||||
if keep {
|
||||
b.WriteRune(r)
|
||||
lastDash = false
|
||||
continue
|
||||
}
|
||||
// 其余一切（空白、标点、符号、寻址保留字符）都折成单个 -
|
||||
if !lastDash && b.Len() > 0 {
|
||||
b.WriteByte('-')
|
||||
lastDash = true
|
||||
}
|
||||
}
|
||||
out := strings.Trim(b.String(), "-")
|
||||
// "new" 是寻址保留字，作为整体别名时必须避开。
|
||||
// 加前缀而不是拒绝：调用方给的素材没有错，是这个词恰好被占用。
|
||||
if out == "new" {
|
||||
return "session-new"
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// truncateAlias 按字节截断且不切坏多字节字符（中文主题很容易超 128 字节）。
|
||||
func truncateAlias(s string) string {
|
||||
if len(s) <= aliasMaxBytes {
|
||||
return strings.Trim(s, "-")
|
||||
}
|
||||
cut := s[:aliasMaxBytes]
|
||||
for len(cut) > 0 && !utf8.ValidString(cut) {
|
||||
cut = cut[:len(cut)-1]
|
||||
}
|
||||
return strings.Trim(cut, "-")
|
||||
}
|
||||
283
server/internal/repo/autoalias_test.go
Normal file
283
server/internal/repo/autoalias_test.go
Normal file
@ -0,0 +1,283 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 这一组测试守的是一条不变量：**`.new` 建出来的会话必须立刻可寻址**。
|
||||
//
|
||||
// 破坏方式很隐蔽 —— 邮件照样送达、收件人照样能回复那一封，只有「指名投进同一条
|
||||
// 会话」这个动作静默失败（`FindNamedSessionFor` 查不到未命名会话），再发一次
|
||||
// `.new` 就多一条平行会话。所以这里的断言都落在「事后能不能按别名找回来」上，
|
||||
// 而不是「有没有报错」。
|
||||
|
||||
func TestAutoAliasForCombinesNameAndSubject(t *testing.T) {
|
||||
// 名字与主题都要在：只用主题时「服务恢复验证」这类通用主题会跨 Agent 撞名，
|
||||
// 只用名字则同一个 Agent 的会话全叫 dsh-2、dsh-3，看不出在聊什么。
|
||||
got := AutoAliasFor("dsh", "重构导入路径")
|
||||
if got != "dsh-重构导入路径" {
|
||||
t.Fatalf("want dsh-重构导入路径, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoAliasForStripsAddressingChars(t *testing.T) {
|
||||
// 别名要参与 name@path.session 的切分，含 . / @ 或空白会让地址解析歧义。
|
||||
// 主题里的 Markdown 与标点噪声（[联调]、—、：）也不该变成一串破折号。
|
||||
cases := []struct{ in, want string }{
|
||||
{"[联调] llmsproxy / ModelRouter — 请提供部署现状", "x-联调-llmsproxy-ModelRouter-请提供部署现状"},
|
||||
{"a.b.c", "x-a-b-c"},
|
||||
{"has spaces", "x-has-spaces"},
|
||||
{"user@host", "x-user-host"},
|
||||
{"Re: 服务恢复验证", "x-Re-服务恢复验证"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := AutoAliasFor("x", c.in)
|
||||
if got != c.want {
|
||||
t.Errorf("AutoAliasFor(x, %q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
if strings.ContainsAny(got, ". \t/@") {
|
||||
t.Errorf("别名 %q 含寻址保留字符，会破坏地址解析", got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoAliasForAvoidsReservedNew(t *testing.T) {
|
||||
// "new" 是 session 位的保留字。别名若正好是它，`x@/p.new` 就同时是
|
||||
// 「投进这条会话」与「再建一条」两种意思。
|
||||
if got := AutoAliasFor("", "new"); got == "new" {
|
||||
t.Fatal("别名不得为保留字 new")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoAliasForEmptyMaterial(t *testing.T) {
|
||||
// 素材里一个可用字符都没有时返回空串，交由 EnsureSessionAlias 走 uuid 兜底，
|
||||
// 而不是在这里编造一个名字。
|
||||
if got := AutoAliasFor("", "···"); got != "" {
|
||||
t.Fatalf("want empty, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoAliasForTruncatesAtByteLimit(t *testing.T) {
|
||||
// session_alias 是 VARCHAR(128)，而中文主题很容易超；按字节截断时
|
||||
// 不能把多字节字符切坏（切坏后写库会得到非法 UTF-8）。
|
||||
got := AutoAliasFor("bot", strings.Repeat("中", 200))
|
||||
if len(got) > aliasMaxBytes {
|
||||
t.Fatalf("别名 %d 字节，超过上限 %d", len(got), aliasMaxBytes)
|
||||
}
|
||||
if !utf8Valid(got) {
|
||||
t.Fatal("截断切坏了多字节字符")
|
||||
}
|
||||
}
|
||||
|
||||
func utf8Valid(s string) bool {
|
||||
for _, r := range s {
|
||||
if r == '\uFFFD' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasMakesNewSessionAddressable(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 复刻 `.new` 的真实路径：CreateSession(alias=nil) —— 别名是 NULL。
|
||||
sid, err := CreateSession(ctx, nil, "admin", "重构导入路径", "/home/program/agentmail")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
seedMailForSession(t, sid, "admin", "dsh", "/home/program/agentmail")
|
||||
|
||||
// 命名前：按别名找不回来（这正是线上那条断链）
|
||||
if _, err := FindNamedSessionFor(ctx, "dsh", "/home/program/agentmail", "dsh-重构导入路径"); err == nil {
|
||||
t.Fatal("未命名会话竟然能按别名找到，测试前提不成立")
|
||||
}
|
||||
|
||||
alias, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("dsh", "重构导入路径"))
|
||||
if err != nil {
|
||||
t.Fatalf("命名: %v", err)
|
||||
}
|
||||
if alias == "" {
|
||||
t.Fatal("别名为空")
|
||||
}
|
||||
|
||||
// 命名后：收件方能指名投回这条会话，而不是又开一条
|
||||
got, err := FindNamedSessionFor(ctx, "dsh", "/home/program/agentmail", alias)
|
||||
if err != nil {
|
||||
t.Fatalf("按别名寻址: %v", err)
|
||||
}
|
||||
if got != sid {
|
||||
t.Fatalf("别名 %q 指向 %s，应指向 %s", alias, got, sid)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasKeepsExistingName(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 调用方显式命名过（发信时传了 session_alias），标记为 manual。
|
||||
// 自动命名绝不能覆盖它 —— 人记住的地址不该下一秒失效。
|
||||
want := "llmsproxy-joint"
|
||||
sid, err := CreateSession(ctx, &want, "dsh", "联调", "/home/program/llmsproxy")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
|
||||
got, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("opencode", "别的主题"))
|
||||
if err != nil {
|
||||
t.Fatalf("命名: %v", err)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("已有别名被改写成 %q，应保持 %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasIsIdempotent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 默认会话路径上 EnsureSessionAlias 会被每封邮件调用一次
|
||||
// （会话可能是刚建的也可能是复用的），因此重复调用必须返回同一个名字。
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "服务恢复验证", "/tmp/ws")
|
||||
first, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("dsh", "服务恢复验证"))
|
||||
if err != nil {
|
||||
t.Fatalf("首次命名: %v", err)
|
||||
}
|
||||
second, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("dsh", "服务恢复验证"))
|
||||
if err != nil {
|
||||
t.Fatalf("二次命名: %v", err)
|
||||
}
|
||||
if first != second {
|
||||
t.Fatalf("重复调用给出两个别名: %q vs %q", first, second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasSuffixesOnCollision(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 同一个 Agent + 同一主题会反复出现（「服务恢复验证」发两次）。
|
||||
// 别名全局唯一（负责寻址），撞名必须让位而不是报错 —— 发信不该因为
|
||||
// 主题重复而失败。
|
||||
want := AutoAliasFor("dsh", "服务恢复验证")
|
||||
|
||||
a, _ := CreateSession(ctx, nil, "admin", "服务恢复验证", "/tmp/ws")
|
||||
aliasA, err := EnsureSessionAlias(ctx, a, want)
|
||||
if err != nil {
|
||||
t.Fatalf("首个会话命名: %v", err)
|
||||
}
|
||||
|
||||
b, _ := CreateSession(ctx, nil, "admin", "服务恢复验证", "/tmp/ws")
|
||||
aliasB, err := EnsureSessionAlias(ctx, b, want)
|
||||
if err != nil {
|
||||
t.Fatalf("第二个会话命名: %v", err)
|
||||
}
|
||||
|
||||
if aliasA == aliasB {
|
||||
t.Fatalf("两条会话拿到同一个别名 %q", aliasA)
|
||||
}
|
||||
if aliasB != want+"-2" {
|
||||
t.Fatalf("撞名后缀应为 %s-2，实际 %q", want, aliasB)
|
||||
}
|
||||
|
||||
// 两个别名各自指向自己那条会话，没有相互覆盖
|
||||
for alias, expect := range map[string]uuid.UUID{aliasA: a, aliasB: b} {
|
||||
var got uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT session_id FROM sessions WHERE session_alias = $1`, alias).Scan(&got)
|
||||
if err != nil {
|
||||
t.Fatalf("查别名 %q: %v", alias, err)
|
||||
}
|
||||
if got != expect {
|
||||
t.Errorf("别名 %q 指向 %s，应指向 %s", alias, got, expect)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasFallsBackToUUID(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 主题与名字都拿不出可用字符时（AutoAliasFor 返回空串），
|
||||
// 仍必须得到一个能寻址的别名 —— 丑名字远胜于没有名字。
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "···", "/tmp/ws")
|
||||
alias, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("", "···"))
|
||||
if err != nil {
|
||||
t.Fatalf("兜底命名: %v", err)
|
||||
}
|
||||
if alias == "" {
|
||||
t.Fatal("兜底后别名仍为空")
|
||||
}
|
||||
if !strings.HasPrefix(alias, "session-") {
|
||||
t.Fatalf("兜底别名应形如 session-xxxxxxxx，实际 %q", alias)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnsureSessionAliasSurfacesInSuggestions(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 补全候选只收「有别名的非归档会话」（`session_alias IS NOT NULL AND <> ''`）。
|
||||
// 自动命名的另一半价值就在这里：命名前这条会话在人类的三段式补全里
|
||||
// 也是不可见的，人同样只能靠回复某封邮件才能续谈。
|
||||
seedUser(t, ctx, "admin")
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "重构导入路径", "/home/program/agentmail")
|
||||
seedMailForSession(t, sid, "admin", "dsh", "/home/program/agentmail")
|
||||
|
||||
before, err := SuggestSessionCandidates(ctx, "admin", "dsh", "/home/program/agentmail")
|
||||
if err != nil {
|
||||
t.Fatalf("补全（命名前）: %v", err)
|
||||
}
|
||||
for _, c := range before {
|
||||
if c.Source == "mail" {
|
||||
t.Fatalf("未命名会话不该出现在补全里，却拿到 %+v", c)
|
||||
}
|
||||
}
|
||||
|
||||
alias, err := EnsureSessionAlias(ctx, sid, AutoAliasFor("dsh", "重构导入路径"))
|
||||
if err != nil {
|
||||
t.Fatalf("命名: %v", err)
|
||||
}
|
||||
|
||||
after, err := SuggestSessionCandidates(ctx, "admin", "dsh", "/home/program/agentmail")
|
||||
if err != nil {
|
||||
t.Fatalf("补全（命名后）: %v", err)
|
||||
}
|
||||
found := false
|
||||
for _, c := range after {
|
||||
if c.Alias == alias {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("命名后 %q 仍未出现在补全候选里: %+v", alias, after)
|
||||
}
|
||||
}
|
||||
|
||||
// seedMailForSession 往会话里塞一封邮件。
|
||||
// FindNamedSessionFor 与 SuggestSessionCandidates 都要求「该收件人参与过」，
|
||||
// 只建会话不建邮件的话两者都查不到，测试会得出错误结论。
|
||||
func seedMailForSession(t *testing.T, sid uuid.UUID, from, to, workspace string) {
|
||||
t.Helper()
|
||||
if _, err := CreateMail(context.Background(), sid, nil,
|
||||
from, "", to, workspace, "主题", "正文", nil); err != nil {
|
||||
t.Fatalf("seed mail: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// seedUser 插一个用户。SuggestSessionCandidates 的可见性条件要查 users 表。
|
||||
func seedUser(t *testing.T, ctx context.Context, username string) {
|
||||
t.Helper()
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ($1, $1, 'x', 'admin')`, username)
|
||||
if err != nil {
|
||||
t.Fatalf("seed user %s: %v", username, err)
|
||||
}
|
||||
}
|
||||
165
server/internal/repo/budget_test.go
Normal file
165
server/internal/repo/budget_test.go
Normal file
@ -0,0 +1,165 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// setupBudgetDB 复用 quota_test.go 的临时库，再建一个会话。
|
||||
// 用真实 SQLite 而非 mock：预算的正确性核心是「判断与自增在同一条 UPDATE 里」，
|
||||
// 那正是只有真实数据库才能验证的部分。
|
||||
func setupBudgetDB(t *testing.T) uuid.UUID {
|
||||
t.Helper()
|
||||
setupTestDB(t)
|
||||
id, err := CreateSession(context.Background(), nil, "bot", "预算测试", "")
|
||||
if err != nil {
|
||||
t.Fatalf("create session: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func TestSessionBudgetZeroMeansUnlimited(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 默认 0 = 不限：引入预算不该把已在进行的会话卡死
|
||||
b, err := GetSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !b.Unlimited || b.Remaining != -1 {
|
||||
t.Fatalf("默认应为不限：%+v", b)
|
||||
}
|
||||
// 不限时反复占用都成功
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := ConsumeSessionBudget(ctx, id); err != nil {
|
||||
t.Fatalf("不限额下第 %d 次占用失败: %v", i+1, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBudgetExhausts(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := SetSessionBudget(ctx, id, 2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 1; i <= 2; i++ {
|
||||
b, err := ConsumeSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次应成功: %v", i, err)
|
||||
}
|
||||
if b.Used != i {
|
||||
t.Fatalf("第 %d 次后 used = %d", i, b.Used)
|
||||
}
|
||||
}
|
||||
b, err := ConsumeSessionBudget(ctx, id)
|
||||
if !errors.Is(err, ErrSessionBudgetExhausted) {
|
||||
t.Fatalf("第 3 次应耗尽，得到 err=%v b=%+v", err, b)
|
||||
}
|
||||
if b.Remaining != 0 {
|
||||
t.Fatalf("耗尽后剩余应为 0：%+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
// 判断与自增必须在同一条 UPDATE 里，否则并发下会把预算刷穿
|
||||
func TestSessionBudgetConcurrentDoesNotOverdraw(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
const limit = 10
|
||||
if _, err := SetSessionBudget(ctx, id, limit); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
ok := 0
|
||||
for i := 0; i < 40; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := ConsumeSessionBudget(ctx, id); err == nil {
|
||||
mu.Lock()
|
||||
ok++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if ok != limit {
|
||||
t.Fatalf("40 并发下成功 %d 次，期望恰好 %d 次（预算被刷穿或误拒）", ok, limit)
|
||||
}
|
||||
b, _ := GetSessionBudget(ctx, id)
|
||||
if b.Used != limit {
|
||||
t.Fatalf("used = %d，期望 %d", b.Used, limit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBudgetResetAndLowerBelowUsed(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
SetSessionBudget(ctx, id, 5)
|
||||
for i := 0; i < 3; i++ {
|
||||
ConsumeSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// 调到低于已用次数 = 「就到这里为止」，是人的合法意图，不该报错
|
||||
b, err := SetSessionBudget(ctx, id, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("下调预算不该失败: %v", err)
|
||||
}
|
||||
if b.Remaining != 0 {
|
||||
t.Fatalf("已用 3 上限 1 时剩余应为 0：%+v", b)
|
||||
}
|
||||
if _, err := ConsumeSessionBudget(ctx, id); !errors.Is(err, ErrSessionBudgetExhausted) {
|
||||
t.Fatal("下调后应立即拦住")
|
||||
}
|
||||
|
||||
// 重置只清已用次数，不动上限
|
||||
b, err = ResetSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.Used != 0 || b.Max != 1 {
|
||||
t.Fatalf("重置后应为 0/1：%+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
// 会话预算先扣、全局配额后扣；全局拦下时必须把会话那次退回去，
|
||||
// 否则那格白掉了 —— 那次往返实际上没有发生
|
||||
func TestRefundSessionBudget(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
SetSessionBudget(ctx, id, 3)
|
||||
ConsumeSessionBudget(ctx, id)
|
||||
|
||||
RefundSessionBudget(ctx, id)
|
||||
b, _ := GetSessionBudget(ctx, id)
|
||||
if b.Used != 0 {
|
||||
t.Fatalf("退还后 used 应为 0：%+v", b)
|
||||
}
|
||||
|
||||
// 已经是 0 时再退不该变成负数
|
||||
RefundSessionBudget(ctx, id)
|
||||
b, _ = GetSessionBudget(ctx, id)
|
||||
if b.Used != 0 {
|
||||
t.Fatalf("重复退还把 used 变成了 %d", b.Used)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetSessionBudgetMissingSession(t *testing.T) {
|
||||
setupBudgetDB(t)
|
||||
if _, err := GetSessionBudget(context.Background(), uuid.New()); err == nil {
|
||||
t.Fatal("不存在的会话应报错")
|
||||
} else if errors.Is(err, sql.ErrNoRows) {
|
||||
t.Fatal("应包装成可读错误而不是裸 sql.ErrNoRows")
|
||||
}
|
||||
}
|
||||
541
server/internal/repo/calendar.go
Normal file
541
server/internal/repo/calendar.go
Normal file
@ -0,0 +1,541 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/lunar"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
var ErrEventNotFound = errors.New("calendar event not found")
|
||||
|
||||
// calendarCols 是所有 SELECT 共用的列清单。
|
||||
//
|
||||
// 抽出来是因为原先有**四处**手抄同一串列名（Get / List / DueEvents 各一处），
|
||||
// 而 Scan 的参数顺序必须与之逐一对应。加一列时漏改任何一处都不会编译报错 ——
|
||||
// 只会在运行时得到 "Scan: expected N destination arguments" 或者更糟：
|
||||
// 列数恰好相同而值错位（曾在 ListSessionsFor 上真的发生过，
|
||||
// 加了预算两列没加进 Scan，整个联系人栏 500）。
|
||||
const calendarCols = `event_id, title, description, reminder_text, agent_name, to_address,
|
||||
recipients, delivery_mode, event_time, remind_before, recurrence, recurrence_end,
|
||||
status, last_fired_at, fired_for, permission_mode, created_at, updated_at, created_by`
|
||||
|
||||
// rowScanner 让 QueryRow 与 Rows 共用同一个 scan 实现。
|
||||
type rowScanner interface {
|
||||
Scan(dest ...any) error
|
||||
}
|
||||
|
||||
// scanCalendarEvent 按 calendarCols 的顺序读一行。
|
||||
//
|
||||
// recipients 存的是 JSON 文本，必须先读进 []byte 再 Unmarshal ——
|
||||
// 直接 Scan 进 []string 会静默失败（driver 不知道怎么转）。
|
||||
func scanCalendarEvent(sc rowScanner) (*models.CalendarEvent, error) {
|
||||
var e models.CalendarEvent
|
||||
var recipientsJSON []byte
|
||||
if err := sc.Scan(
|
||||
&e.EventID, &e.Title, &e.Description, &e.ReminderText,
|
||||
&e.AgentName, &e.ToAddress,
|
||||
&recipientsJSON, &e.DeliveryMode,
|
||||
&e.EventTime, &e.RemindBefore, &e.Recurrence, &e.RecurrenceEnd,
|
||||
&e.Status, &e.LastFiredAt, &e.FiredFor, &e.PermissionMode,
|
||||
&e.CreatedAt, &e.UpdatedAt, &e.CreatedBy,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(recipientsJSON) > 0 {
|
||||
// 解析失败不算致命：退回 to_address/agent_name 兜底链，
|
||||
// 事件仍能投递。让一条脏 JSON 把整个列表打成 500 更糟。
|
||||
_ = json.Unmarshal(recipientsJSON, &e.Recipients)
|
||||
}
|
||||
if e.Recipients == nil {
|
||||
// Go 的 nil slice 序列化成 null，前端 .map 会崩
|
||||
e.Recipients = []string{}
|
||||
}
|
||||
return &e, nil
|
||||
}
|
||||
|
||||
// marshalRecipients 把收件人列表序列化成入库的 JSON 文本。
|
||||
func marshalRecipients(list []string) string {
|
||||
if list == nil {
|
||||
list = []string{}
|
||||
}
|
||||
b, err := json.Marshal(list)
|
||||
if err != nil {
|
||||
return "[]"
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// ─── CRUD ───
|
||||
|
||||
func CreateCalendarEvent(ctx context.Context, e *models.CalendarEvent) (*models.CalendarEvent, error) {
|
||||
e.EventID = uuid.New().String()
|
||||
e.CreatedAt = time.Now()
|
||||
e.UpdatedAt = e.CreatedAt
|
||||
if e.Status == "" {
|
||||
e.Status = "active"
|
||||
}
|
||||
if e.Recurrence == "" {
|
||||
e.Recurrence = "none"
|
||||
}
|
||||
|
||||
if e.DeliveryMode == "" {
|
||||
e.DeliveryMode = models.DeliverSeparate
|
||||
}
|
||||
// 档位合法化：脏值 fail-closed 到默认档，不透传成库里的非法值
|
||||
//（否则后续读路径会拿到一个 ModeNeedsHuman 判定不了的值）。
|
||||
e.PermissionMode = models.NormalizePermissionMode(e.PermissionMode)
|
||||
|
||||
if e.Recipients == nil {
|
||||
e.Recipients = []string{}
|
||||
}
|
||||
|
||||
_, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO calendar_events
|
||||
(event_id, title, description, reminder_text, agent_name, to_address,
|
||||
recipients, delivery_mode,
|
||||
event_time, remind_before, recurrence, recurrence_end, status, permission_mode, created_by,
|
||||
created_at, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
e.EventID, e.Title, e.Description, e.ReminderText,
|
||||
e.AgentName, e.ToAddress,
|
||||
marshalRecipients(e.Recipients), e.DeliveryMode,
|
||||
e.EventTime, e.RemindBefore, e.Recurrence, e.RecurrenceEnd,
|
||||
e.Status, e.PermissionMode, e.CreatedBy, e.CreatedAt, e.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
func GetCalendarEvent(ctx context.Context, eventID string) (*models.CalendarEvent, error) {
|
||||
e, err := scanCalendarEvent(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+calendarCols+` FROM calendar_events WHERE event_id = ?`, eventID))
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrEventNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
return e, nil
|
||||
}
|
||||
|
||||
func UpdateCalendarEvent(ctx context.Context, eventID string, e *models.CalendarEvent) error {
|
||||
e.UpdatedAt = time.Now()
|
||||
// 档位同样在 update 路径上规范化
|
||||
e.PermissionMode = models.NormalizePermissionMode(e.PermissionMode)
|
||||
result, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE calendar_events SET
|
||||
title = ?, description = ?, reminder_text = ?,
|
||||
agent_name = ?, to_address = ?,
|
||||
recipients = ?, delivery_mode = ?,
|
||||
event_time = ?, remind_before = ?, recurrence = ?, recurrence_end = ?,
|
||||
status = ?, permission_mode = ?, updated_at = ?
|
||||
WHERE event_id = ?`,
|
||||
e.Title, e.Description, e.ReminderText,
|
||||
e.AgentName, e.ToAddress,
|
||||
marshalRecipients(e.Recipients), e.EffectiveDeliveryMode(),
|
||||
e.EventTime, e.RemindBefore, e.Recurrence, e.RecurrenceEnd,
|
||||
e.Status, e.PermissionMode, e.UpdatedAt, eventID,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
if n == 0 {
|
||||
return ErrEventNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func DeleteCalendarEvent(ctx context.Context, eventID string) error {
|
||||
result, err := db.DB.ExecContext(ctx, `DELETE FROM calendar_events WHERE event_id = ?`, eventID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, _ := result.RowsAffected()
|
||||
if n == 0 {
|
||||
return ErrEventNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ─── 查询 ───
|
||||
|
||||
// ListCalendarEvents 返回指定时间范围内的事件（日历视图）。
|
||||
func ListCalendarEvents(ctx context.Context, from, to time.Time, status string) ([]models.CalendarEvent, error) {
|
||||
if status == "" {
|
||||
status = "active"
|
||||
}
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT `+calendarCols+`
|
||||
FROM calendar_events
|
||||
WHERE event_time >= ? AND event_time <= ?
|
||||
AND (status = ? OR ? = '')
|
||||
ORDER BY event_time ASC`, from, to, status, status)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var events []models.CalendarEvent
|
||||
for rows.Next() {
|
||||
e, err := scanCalendarEvent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
events = append(events, *e)
|
||||
}
|
||||
return events, rows.Err()
|
||||
}
|
||||
|
||||
// ─── 调度器 ───
|
||||
|
||||
// DueEvents 返回下一分钟内需要触发的事件。
|
||||
//
|
||||
// 调度器每分钟调用一次：event_time + remind_before <= now+60s 且尚未触发（last_fired_at 为 NULL
|
||||
// 或小于 event_time）的 active 事件。
|
||||
// DueEvents 取出该触发的事件。
|
||||
//
|
||||
// 60 秒 lookahead 让提醒宁早不晚：调度周期是 30 秒，不提前看的话
|
||||
// 一个刚好落在两个 tick 之间的提醒会迟到最多 30 秒。
|
||||
//
|
||||
// **去重判据是 fired_for（已触发的 occurrence）与 event_time 相等**，
|
||||
// 不是 last_fired_at 与 event_time 比大小 —— 后者在 lookahead 窗口内
|
||||
// 恒为真（触发时刻早于 event_time），会让同一条提醒每个 tick 重发一次。
|
||||
func DueEvents(ctx context.Context) ([]models.CalendarEvent, error) {
|
||||
now := time.Now()
|
||||
deadline := now.Add(60 * time.Second)
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT `+calendarCols+`
|
||||
FROM calendar_events
|
||||
WHERE status = 'active'
|
||||
AND datetime(event_time, '-' || remind_before || ' minutes') <= ?
|
||||
AND (fired_for IS NULL OR fired_for <> event_time)
|
||||
ORDER BY event_time ASC`, deadline)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var events []models.CalendarEvent
|
||||
for rows.Next() {
|
||||
e, err := scanCalendarEvent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
events = append(events, *e)
|
||||
}
|
||||
return events, rows.Err()
|
||||
}
|
||||
|
||||
// MarkEventFired 标记事件已触发，防止重复。
|
||||
func MarkEventFired(ctx context.Context, eventID string) error {
|
||||
// fired_for 直接从 event_time 列复制而不是在 Go 侧格式化再写回：
|
||||
// 两者必须逐字节相同（判据是字符串相等），经过一轮 time.Time 往返
|
||||
// 有可能改变表示形式。
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE calendar_events SET last_fired_at = ?, fired_for = event_time
|
||||
WHERE event_id = ?`,
|
||||
time.Now(), eventID)
|
||||
return err
|
||||
}
|
||||
|
||||
// AdvanceRecurrence 为重复事件计算下一次触发时间。
|
||||
//
|
||||
// 返回 false 表示重复已过期（recurrence_end 已过），事件应置为 cancelled。
|
||||
func AdvanceRecurrence(ctx context.Context, eventID string) (bool, error) {
|
||||
var recurrence string
|
||||
var eventTime time.Time
|
||||
var recurrenceEnd *time.Time
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT recurrence, event_time, recurrence_end
|
||||
FROM calendar_events WHERE event_id = ?`, eventID).Scan(
|
||||
&recurrence, &eventTime, &recurrenceEnd)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
if recurrence == models.RecurNone {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// **一路推到未来**，不是只推一步。
|
||||
//
|
||||
// 只推一步的后果（实测）：一条 100 天前设的每日事件，每轮扫描都判定
|
||||
// 「已过期该触发」→ 发一封 → event_time 只前进一天 → 下一轮又过期。
|
||||
// 30 轮扫描触发 30 次，而调度周期是 30 秒 —— 人会收到一串垃圾提醒，
|
||||
// 连发 100 封才追上今天。
|
||||
//
|
||||
// 跳过的那些 occurrence **不补发**：定时提醒的价值在于「按时」，
|
||||
// 三个月前那次站会提醒现在发出去毫无意义，只会淹掉真正该看的那封。
|
||||
// 本轮仍会发一封（fireEvent 已经在发了），代表「这条规则还活着」。
|
||||
next, err := advanceToFuture(recurrence, eventTime, time.Now(), recurrenceEnd)
|
||||
if err != nil {
|
||||
// 推不出下一次（例如「每年农历闰六月」而目标年无闰六月）：
|
||||
// 置为 cancelled 而不是留在 active 空转。留着会让调度器每 30 秒
|
||||
// 重试同一个算不出来的规则，日志里刷同一条错误直到有人发现。
|
||||
_, cErr := db.DB.ExecContext(ctx,
|
||||
`UPDATE calendar_events SET status = 'cancelled' WHERE event_id = ?`, eventID)
|
||||
if cErr != nil {
|
||||
return false, cErr
|
||||
}
|
||||
return false, err
|
||||
}
|
||||
// 零值 = 已越过 recurrence_end（"none" 在函数开头就返回了，到不了这里）。
|
||||
// **必须置 cancelled**：留在 active 会让 DueEvents 每轮都捞到这条
|
||||
// 早已过期的事件，而 fired_for 已经等于 event_time 所以它又不会被触发 ——
|
||||
// 表现是一条永远排在到期列表里、永远不动的僵尸事件。
|
||||
if next.IsZero() {
|
||||
_, cErr := db.DB.ExecContext(ctx,
|
||||
`UPDATE calendar_events SET status = 'cancelled', updated_at = ? WHERE event_id = ?`,
|
||||
time.Now(), eventID)
|
||||
return false, cErr
|
||||
}
|
||||
|
||||
// 走到这里说明 next 既在未来又在终止时间之内。
|
||||
_, err = db.DB.ExecContext(ctx,
|
||||
`UPDATE calendar_events SET event_time = ?, updated_at = ? WHERE event_id = ?`,
|
||||
next, time.Now(), eventID)
|
||||
return true, err
|
||||
}
|
||||
|
||||
// advanceToFuture 从 from 起反复按规则推进，直到越过 now。
|
||||
//
|
||||
// 单独成不碰数据库的函数是为了可测。三个终止条件，缺一不可：
|
||||
//
|
||||
// 1. 越过 now —— 正常出口
|
||||
// 2. 越过 recurrenceEnd —— 返回零值，调用方据此置 cancelled
|
||||
// 3. maxAdvanceSteps 上限 —— 防御性的。规则算得出但不前进（理论上
|
||||
// NextOccurrence 不会返回 <= 当前值，但农历那条路径依赖外部库，
|
||||
// 一旦它某年给出反直觉结果，没有上限就是个死循环 goroutine，
|
||||
// 而它跑在调度器里 —— 整个提醒系统会一起卡住）
|
||||
//
|
||||
// 上限取 4000：按每日重复算约 11 年，足够覆盖「很久以前设的提醒」，
|
||||
// 而 4000 次纯内存日期运算在一个 tick 里跑完毫无压力。
|
||||
const maxAdvanceSteps = 4000
|
||||
|
||||
func advanceToFuture(recurrence string, from, now time.Time, recurrenceEnd *time.Time) (time.Time, error) {
|
||||
cur := from
|
||||
for i := 0; i < maxAdvanceSteps; i++ {
|
||||
next, err := NextOccurrence(recurrence, cur)
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if next.IsZero() {
|
||||
return time.Time{}, nil // 不重复
|
||||
}
|
||||
if !next.After(cur) {
|
||||
// 规则不前进 —— 与死循环等价，当作算不出来
|
||||
return time.Time{}, fmt.Errorf("重复规则 %q 未能前进（停在 %s）", recurrence, cur.Format(time.RFC3339))
|
||||
}
|
||||
cur = next
|
||||
if recurrenceEnd != nil && cur.After(*recurrenceEnd) {
|
||||
return time.Time{}, nil // 已过终止时间
|
||||
}
|
||||
if cur.After(now) {
|
||||
return cur, nil
|
||||
}
|
||||
}
|
||||
return time.Time{}, fmt.Errorf("重复规则 %q 推进 %d 次仍未越过当前时刻", recurrence, maxAdvanceSteps)
|
||||
}
|
||||
|
||||
// NextOccurrence 按重复规则算出下一次触发时刻。
|
||||
//
|
||||
// 独立成不碰数据库的纯函数是为了可测：农历推进错了不会报错，
|
||||
// 只会让提醒发在错误的日子，而那种错误要等真的过了一个月才看得见。
|
||||
//
|
||||
// 返回零值 time 且 err == nil 表示「不重复」（规则是 none 或未知值）。
|
||||
//
|
||||
// **农历规则不能用 AddDate 近似**：农历月 29~30 天不定、农历年 353~385 天
|
||||
// （闰年多一整月）。用固定天数推进一年能偏半个月 —— 农历生日提醒会
|
||||
// 逐年漂移到完全不相干的日子上。
|
||||
func NextOccurrence(recurrence string, from time.Time) (time.Time, error) {
|
||||
switch recurrence {
|
||||
case models.RecurDaily:
|
||||
return from.AddDate(0, 0, 1), nil
|
||||
case models.RecurWeekly:
|
||||
return from.AddDate(0, 0, 7), nil
|
||||
case models.RecurMonthly:
|
||||
// 公历每月：AddDate 在月末会溢出（1 月 31 日 +1 月 = 3 月 3 日）。
|
||||
// 夹到目标月的最后一天 —— 与农历那边的 clamp 语义一致：
|
||||
// 「每月 31 日」的意思是「月末」，滚到下月初是错的。
|
||||
return addSolarMonthClamped(from, 1), nil
|
||||
case models.RecurYearly:
|
||||
// 公历每年：2 月 29 日在平年会溢出成 3 月 1 日，同样要夹。
|
||||
// 闰日生日的约定是「平年过 2 月 28」，不是 3 月 1 日。
|
||||
return addSolarMonthClamped(from, 12), nil
|
||||
case models.RecurLunarMonthly:
|
||||
d := lunar.FromSolar(from).AddMonths(1)
|
||||
t, _, err := d.ToSolar(from.Location(), from.Hour(), from.Minute(), from.Second(), from.Nanosecond())
|
||||
return t, err
|
||||
case models.RecurLunarYearly:
|
||||
d := lunar.FromSolar(from).AddYears(1)
|
||||
t, _, err := d.ToSolar(from.Location(), from.Hour(), from.Minute(), from.Second(), from.Nanosecond())
|
||||
return t, err
|
||||
default:
|
||||
return time.Time{}, nil
|
||||
}
|
||||
}
|
||||
|
||||
// addSolarMonthClamped 在公历上加月份，日期夹到目标月的实际天数内。
|
||||
//
|
||||
// time.AddDate 的溢出行为（3 月 31 日 +1 月 = 5 月 1 日）对「每月同一日」
|
||||
// 的提醒是错的：31 日的事件会在 2 月变成 3 月 3 日，然后从此每月 3 日提醒
|
||||
// —— 一次溢出永久改变了规则。
|
||||
func addSolarMonthClamped(t time.Time, n int) time.Time {
|
||||
y, m, d := t.Date()
|
||||
m += time.Month(n)
|
||||
for m > 12 {
|
||||
m -= 12
|
||||
y++
|
||||
}
|
||||
// 目标月第 0 天 = 上个月最后一天，用它拿到月长
|
||||
last := time.Date(y, m+1, 0, 0, 0, 0, 0, t.Location()).Day()
|
||||
if d > last {
|
||||
d = last
|
||||
}
|
||||
return time.Date(y, m, d, t.Hour(), t.Minute(), t.Second(), t.Nanosecond(), t.Location())
|
||||
}
|
||||
|
||||
// ─── 附件 ───
|
||||
|
||||
func AddCalendarAttachment(ctx context.Context, a *models.CalendarAttachment) error {
|
||||
a.AttachmentID = uuid.New().String()
|
||||
a.CreatedAt = time.Now()
|
||||
_, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO calendar_attachments (attachment_id, event_id, filename, sha256, size_bytes, created_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)`,
|
||||
a.AttachmentID, a.EventID, a.Filename, a.SHA256, a.SizeBytes, a.CreatedAt)
|
||||
return err
|
||||
}
|
||||
|
||||
func ListCalendarAttachments(ctx context.Context, eventID string) ([]models.CalendarAttachment, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT attachment_id, event_id, filename, sha256, size_bytes, created_at
|
||||
FROM calendar_attachments WHERE event_id = ? ORDER BY created_at`, eventID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var atts []models.CalendarAttachment
|
||||
for rows.Next() {
|
||||
var a models.CalendarAttachment
|
||||
if err := rows.Scan(&a.AttachmentID, &a.EventID, &a.Filename, &a.SHA256, &a.SizeBytes, &a.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
atts = append(atts, a)
|
||||
}
|
||||
return atts, rows.Err()
|
||||
}
|
||||
|
||||
func DeleteCalendarAttachments(ctx context.Context, eventID string) error {
|
||||
_, err := db.DB.ExecContext(ctx, `DELETE FROM calendar_attachments WHERE event_id = ?`, eventID)
|
||||
return err
|
||||
}
|
||||
|
||||
// DeleteCalendarAttachment 删单条附件。
|
||||
//
|
||||
// 返回 false 表示这条不存在（而不是报错）：调用方据此回 404 而非 500。
|
||||
// 只删元数据，磁盘 blob 留给 GC —— 内容寻址下同一个 sha256 可能被别的
|
||||
// 附件引用着，跟着删会让那些引用一起坏掉。
|
||||
func DeleteCalendarAttachment(ctx context.Context, attachmentID string) (bool, error) {
|
||||
res, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM calendar_attachments WHERE attachment_id = ?`, attachmentID)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
n, err := res.RowsAffected()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n > 0, nil
|
||||
}
|
||||
|
||||
// AttachCalendarFilesToMail 把事件的附件复制成邮件附件。
|
||||
//
|
||||
// 提醒邮件是新建的，附件必须重新挂一份指向同一 sha256 的元数据 ——
|
||||
// 内容寻址下这不拷磁盘文件，只是多一条记录。
|
||||
//
|
||||
// 缺了这一步的后果：人在事件上传了附件、UI 里看得见、提醒也按时发出，
|
||||
// 但 Agent 收到的那封信里附件清单是空的 —— 事件附件与邮件附件是两张表，
|
||||
// 不复制就永远只存在于日历侧。这是「日历附件只记元数据未接投递」的另一半。
|
||||
//
|
||||
// uploader 记为 calendarSender（"calendar"）：附件随提醒邮件重新分发，
|
||||
// 其可见范围由该邮件的参与方决定，而不是沿用事件创建者。
|
||||
func AttachCalendarFilesToMail(ctx context.Context, eventID string, mailID uuid.UUID, uploader string) (int, error) {
|
||||
atts, err := ListCalendarAttachments(ctx, eventID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
n := 0
|
||||
for _, a := range atts {
|
||||
// sha256 为空说明这条记录没有真实内容（历史脏数据），跳过而不是
|
||||
// 挂一个下载必然 404 的附件
|
||||
if a.SHA256 == "" {
|
||||
continue
|
||||
}
|
||||
if _, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO attachments (mail_id, uploader, filename, content_type, size_bytes, sha256)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
`, mailID, uploader, a.Filename, "application/octet-stream", a.SizeBytes, a.SHA256); err != nil {
|
||||
return n, err
|
||||
}
|
||||
n++
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// ListCalendarEventsCreatedBy 只返回某个创建者建的事件。
|
||||
//
|
||||
// Agent 侧列表用这个而不是 ListCalendarEvents：Agent 不该看到别人（人类或
|
||||
// 其他 Agent）的日程 —— 那里可能有它无权知道的会议、地址、附件名。
|
||||
//
|
||||
// 注意**不是**「发给我的事件」：`recipients` 里有我但我没建的，同样不返回。
|
||||
// 理由是那些事件的编辑权不属于我，列出来只会让模型试图改它然后拿到 403。
|
||||
// 想知道「谁给我设了提醒」，那条信息在提醒邮件本身里。
|
||||
func ListCalendarEventsCreatedBy(ctx context.Context, creator string, from, to time.Time, status string) ([]models.CalendarEvent, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT `+calendarCols+`
|
||||
FROM calendar_events
|
||||
WHERE created_by = ?
|
||||
AND event_time >= ? AND event_time <= ?
|
||||
AND (status = ? OR ? = '')
|
||||
ORDER BY event_time ASC`, creator, from, to, status, status)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
events := []models.CalendarEvent{}
|
||||
for rows.Next() {
|
||||
e, err := scanCalendarEvent(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
events = append(events, *e)
|
||||
}
|
||||
return events, rows.Err()
|
||||
}
|
||||
|
||||
// CountActiveEventsBy 数某个创建者当前有多少条生效中的事件。
|
||||
//
|
||||
// 给 Agent 侧的总量上限用。速率限制只压住「短时间内暴建」，
|
||||
// 压不住「每小时建 19 条、连建一周」—— 而日历事件是长效的，
|
||||
// 攒下来的每一条都会持续产生提醒邮件。
|
||||
func CountActiveEventsBy(ctx context.Context, creator string) (int, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM calendar_events WHERE created_by = ? AND status = 'active'`,
|
||||
creator).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
972
server/internal/repo/calendar_test.go
Normal file
972
server/internal/repo/calendar_test.go
Normal file
@ -0,0 +1,972 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/lunar"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func seedEvent(t *testing.T, e *models.CalendarEvent) *models.CalendarEvent {
|
||||
t.Helper()
|
||||
if e.Title == "" {
|
||||
e.Title = "测试事件"
|
||||
}
|
||||
if e.EventTime.IsZero() {
|
||||
e.EventTime = time.Now().Add(time.Hour)
|
||||
}
|
||||
out, err := CreateCalendarEvent(context.Background(), e)
|
||||
if err != nil {
|
||||
t.Fatalf("建事件: %v", err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestCalendarEventCRUD(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
at := time.Now().Add(2 * time.Hour).Truncate(time.Second)
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "每日站会",
|
||||
Description: "同步进展",
|
||||
ReminderText: "日程提醒：{title}",
|
||||
AgentName: "dsh",
|
||||
ToAddress: "dsh@/home",
|
||||
EventTime: at,
|
||||
RemindBefore: 15,
|
||||
Recurrence: "daily",
|
||||
CreatedBy: "jianf",
|
||||
})
|
||||
|
||||
if e.EventID == "" {
|
||||
t.Fatal("建完事件必须有 event_id")
|
||||
}
|
||||
if e.Status != "active" {
|
||||
t.Errorf("新事件默认应为 active，得到 %q", e.Status)
|
||||
}
|
||||
|
||||
got, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读事件: %v", err)
|
||||
}
|
||||
if got.Title != "每日站会" || got.RemindBefore != 15 || got.Recurrence != "daily" {
|
||||
t.Errorf("读回的字段不符：%+v", got)
|
||||
}
|
||||
if !got.EventTime.Equal(at) {
|
||||
t.Errorf("event_time 读回错位：写 %v 读 %v", at, got.EventTime)
|
||||
}
|
||||
|
||||
got.Title = "改名后的站会"
|
||||
got.Status = "paused"
|
||||
if err := UpdateCalendarEvent(ctx, e.EventID, got); err != nil {
|
||||
t.Fatalf("改事件: %v", err)
|
||||
}
|
||||
again, _ := GetCalendarEvent(ctx, e.EventID)
|
||||
if again.Title != "改名后的站会" || again.Status != "paused" {
|
||||
t.Errorf("改后没生效：%+v", again)
|
||||
}
|
||||
|
||||
if err := DeleteCalendarEvent(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("删事件: %v", err)
|
||||
}
|
||||
if _, err := GetCalendarEvent(ctx, e.EventID); err != ErrEventNotFound {
|
||||
t.Errorf("删掉后应报 ErrEventNotFound，得到 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalendarNotFoundIsTyped(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 不存在的 id 要给出可判定的错误，而不是 sql.ErrNoRows ——
|
||||
// handler 靠它区分 404 与 500。
|
||||
if _, err := GetCalendarEvent(ctx, "00000000-0000-0000-0000-000000000000"); err != ErrEventNotFound {
|
||||
t.Errorf("Get 应报 ErrEventNotFound，得到 %v", err)
|
||||
}
|
||||
if err := DeleteCalendarEvent(ctx, "00000000-0000-0000-0000-000000000000"); err != ErrEventNotFound {
|
||||
t.Errorf("Delete 应报 ErrEventNotFound，得到 %v", err)
|
||||
}
|
||||
if err := UpdateCalendarEvent(ctx, "00000000-0000-0000-0000-000000000000",
|
||||
&models.CalendarEvent{Title: "x", EventTime: time.Now()}); err != ErrEventNotFound {
|
||||
t.Errorf("Update 应报 ErrEventNotFound，得到 %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDueEventsOnlyReturnsRipe(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// 已经该响的（事件时间在过去）
|
||||
ripe := seedEvent(t, &models.CalendarEvent{Title: "该响了", EventTime: now.Add(-time.Minute)})
|
||||
// 提前 30 分钟提醒、事件在 20 分钟后 —— 提醒点已过
|
||||
early := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "提前提醒已到", EventTime: now.Add(20 * time.Minute), RemindBefore: 30,
|
||||
})
|
||||
// 还早（1 小时后，无提前提醒）
|
||||
future := seedEvent(t, &models.CalendarEvent{Title: "还早", EventTime: now.Add(time.Hour)})
|
||||
// 已暂停的不该响
|
||||
paused := seedEvent(t, &models.CalendarEvent{Title: "暂停的", EventTime: now.Add(-time.Minute)})
|
||||
p, _ := GetCalendarEvent(ctx, paused.EventID)
|
||||
p.Status = "paused"
|
||||
if err := UpdateCalendarEvent(ctx, paused.EventID, p); err != nil {
|
||||
t.Fatalf("暂停: %v", err)
|
||||
}
|
||||
|
||||
due, err := DueEvents(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("DueEvents: %v", err)
|
||||
}
|
||||
|
||||
got := map[string]bool{}
|
||||
for _, e := range due {
|
||||
got[e.EventID] = true
|
||||
}
|
||||
if !got[ripe.EventID] {
|
||||
t.Error("到期事件没被取出")
|
||||
}
|
||||
if !got[early.EventID] {
|
||||
t.Error("remind_before 已过的事件没被取出")
|
||||
}
|
||||
if got[future.EventID] {
|
||||
t.Error("未到期事件被取出了")
|
||||
}
|
||||
if got[paused.EventID] {
|
||||
t.Error("已暂停的事件被取出了")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkEventFiredStopsRefiring(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 幂等的关键：标记后同一条不该再出现在 DueEvents 里，
|
||||
// 否则调度器每 30 秒把同一封提醒重发一遍。
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "只该响一次", EventTime: time.Now().Add(-time.Minute)})
|
||||
|
||||
due, _ := DueEvents(ctx)
|
||||
if len(due) != 1 {
|
||||
t.Fatalf("标记前应有 1 条到期，得到 %d", len(due))
|
||||
}
|
||||
|
||||
if err := MarkEventFired(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("标记: %v", err)
|
||||
}
|
||||
|
||||
due, _ = DueEvents(ctx)
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
t.Error("已标记触发的事件仍出现在 DueEvents 里")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvanceRecurrence(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
base := time.Now().Add(-time.Minute).Truncate(time.Second)
|
||||
|
||||
t.Run("一次性事件不推进", func(t *testing.T) {
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "一次性", EventTime: base, Recurrence: "none"})
|
||||
advanced, err := AdvanceRecurrence(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if advanced {
|
||||
t.Error("recurrence=none 不该推进")
|
||||
}
|
||||
})
|
||||
|
||||
for _, tc := range []struct {
|
||||
rule string
|
||||
want time.Time
|
||||
}{
|
||||
{"daily", base.AddDate(0, 0, 1)},
|
||||
{"weekly", base.AddDate(0, 0, 7)},
|
||||
{"monthly", base.AddDate(0, 1, 0)},
|
||||
} {
|
||||
t.Run(tc.rule+" 推进一个周期", func(t *testing.T) {
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: tc.rule, EventTime: base, Recurrence: tc.rule,
|
||||
})
|
||||
advanced, err := AdvanceRecurrence(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if !advanced {
|
||||
t.Fatal("应该推进")
|
||||
}
|
||||
got, _ := GetCalendarEvent(ctx, e.EventID)
|
||||
if !got.EventTime.Equal(tc.want) {
|
||||
t.Errorf("下次时间应为 %v，得到 %v", tc.want, got.EventTime)
|
||||
}
|
||||
// 推进后 event_time 已在未来，且 last_fired_at 仍为旧值 →
|
||||
// 必须重新出现在 DueEvents 里等待下一轮（否则重复事件只响一次）。
|
||||
if got.Status != "active" {
|
||||
t.Errorf("推进后应仍为 active，得到 %q", got.Status)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("超过 recurrence_end 则取消", func(t *testing.T) {
|
||||
end := base.Add(12 * time.Hour) // 下一次（+1 天）会越过它
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "快结束了", EventTime: base, Recurrence: "daily", RecurrenceEnd: &end,
|
||||
})
|
||||
advanced, err := AdvanceRecurrence(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if advanced {
|
||||
t.Error("越过 recurrence_end 时不该报告推进成功")
|
||||
}
|
||||
got, _ := GetCalendarEvent(ctx, e.EventID)
|
||||
if got.Status != "cancelled" {
|
||||
t.Errorf("越过结束时间应置为 cancelled，得到 %q", got.Status)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestListCalendarEventsRange(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
now := time.Now()
|
||||
inRange := seedEvent(t, &models.CalendarEvent{Title: "范围内", EventTime: now.Add(time.Hour)})
|
||||
seedEvent(t, &models.CalendarEvent{Title: "太远", EventTime: now.AddDate(0, 3, 0)})
|
||||
|
||||
events, err := ListCalendarEvents(ctx, now, now.Add(24*time.Hour), "active")
|
||||
if err != nil {
|
||||
t.Fatalf("列事件: %v", err)
|
||||
}
|
||||
if len(events) != 1 || events[0].EventID != inRange.EventID {
|
||||
t.Errorf("时间范围过滤不对，得到 %d 条", len(events))
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalendarAttachments(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "带附件"})
|
||||
|
||||
if err := AddCalendarAttachment(ctx, &models.CalendarAttachment{
|
||||
EventID: e.EventID, Filename: "报表.xlsx", SHA256: "abc", SizeBytes: 2048,
|
||||
}); err != nil {
|
||||
t.Fatalf("加附件: %v", err)
|
||||
}
|
||||
|
||||
atts, err := ListCalendarAttachments(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("列附件: %v", err)
|
||||
}
|
||||
if len(atts) != 1 || atts[0].Filename != "报表.xlsx" {
|
||||
t.Fatalf("附件读回不符：%+v", atts)
|
||||
}
|
||||
if atts[0].AttachmentID == "" {
|
||||
t.Error("附件必须有 attachment_id —— 没有它模型无法在 send_mail 里引用")
|
||||
}
|
||||
|
||||
// 事件没有附件时返回空而不是报错
|
||||
other := seedEvent(t, &models.CalendarEvent{Title: "没附件"})
|
||||
if atts, err := ListCalendarAttachments(ctx, other.EventID); err != nil || len(atts) != 0 {
|
||||
t.Errorf("无附件事件应返回空列表，得到 %d 条 err=%v", len(atts), err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteCalendarAttachment(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "要删附件"})
|
||||
for _, name := range []string{"甲.pdf", "乙.pdf"} {
|
||||
if err := AddCalendarAttachment(ctx, &models.CalendarAttachment{
|
||||
EventID: e.EventID, Filename: name, SHA256: "sum-" + name, SizeBytes: 10,
|
||||
}); err != nil {
|
||||
t.Fatalf("加附件 %s: %v", name, err)
|
||||
}
|
||||
}
|
||||
|
||||
atts, _ := ListCalendarAttachments(ctx, e.EventID)
|
||||
if len(atts) != 2 {
|
||||
t.Fatalf("准备阶段应有 2 个附件，得到 %d", len(atts))
|
||||
}
|
||||
|
||||
ok, err := DeleteCalendarAttachment(ctx, atts[0].AttachmentID)
|
||||
if err != nil {
|
||||
t.Fatalf("删附件: %v", err)
|
||||
}
|
||||
if !ok {
|
||||
t.Error("删掉存在的附件应返回 true")
|
||||
}
|
||||
|
||||
left, _ := ListCalendarAttachments(ctx, e.EventID)
|
||||
if len(left) != 1 {
|
||||
t.Fatalf("删一个后应剩 1 个，得到 %d", len(left))
|
||||
}
|
||||
if left[0].AttachmentID == atts[0].AttachmentID {
|
||||
t.Error("删错了对象")
|
||||
}
|
||||
|
||||
// 不存在的 id 返回 false 而不是报错 —— 调用方据此回 404 而非 500
|
||||
ok, err = DeleteCalendarAttachment(ctx, "00000000-0000-0000-0000-000000000000")
|
||||
if err != nil {
|
||||
t.Errorf("删不存在的附件不该报错，得到 %v", err)
|
||||
}
|
||||
if ok {
|
||||
t.Error("删不存在的附件应返回 false")
|
||||
}
|
||||
}
|
||||
|
||||
// 事件附件必须能复制成邮件附件。
|
||||
//
|
||||
// 少了这一步，附件只存在于日历侧：UI 里看得见、提醒按时发出、
|
||||
// 而 Agent 收到的那封信附件清单是空的 —— 两张表互不相通。
|
||||
func TestAttachCalendarFilesToMail(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "带附件的提醒"})
|
||||
if err := AddCalendarAttachment(ctx, &models.CalendarAttachment{
|
||||
EventID: e.EventID, Filename: "周报.md", SHA256: "deadbeef", SizeBytes: 512,
|
||||
}); err != nil {
|
||||
t.Fatalf("加附件: %v", err)
|
||||
}
|
||||
// sha256 为空的脏数据必须被跳过：挂上去只会得到一个下载必然 404 的附件
|
||||
if err := AddCalendarAttachment(ctx, &models.CalendarAttachment{
|
||||
EventID: e.EventID, Filename: "没内容.bin", SHA256: "", SizeBytes: 0,
|
||||
}); err != nil {
|
||||
t.Fatalf("加空附件: %v", err)
|
||||
}
|
||||
|
||||
seedAgentForAttach(t, "pi")
|
||||
sessionID := seedSessionForAttach(t, "pi")
|
||||
mailID, err := CreateMail(ctx, sessionID, nil, "calendar", "", "pi", "", "日程提醒：带附件的提醒", "正文", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
|
||||
n, err := AttachCalendarFilesToMail(ctx, e.EventID, mailID, "calendar")
|
||||
if err != nil {
|
||||
t.Fatalf("挂附件: %v", err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("应只挂 1 个（空 sha256 那条跳过），得到 %d", n)
|
||||
}
|
||||
|
||||
mailAtts, err := ListAttachmentsFor(ctx, mailID)
|
||||
if err != nil {
|
||||
t.Fatalf("列邮件附件: %v", err)
|
||||
}
|
||||
if len(mailAtts) != 1 {
|
||||
t.Fatalf("邮件上应有 1 个附件，得到 %d", len(mailAtts))
|
||||
}
|
||||
if mailAtts[0].Filename != "周报.md" || mailAtts[0].SHA256 != "deadbeef" {
|
||||
t.Errorf("附件内容不符：%+v", mailAtts[0])
|
||||
}
|
||||
// 内容寻址：复制不产生新的 sha256，指向同一份磁盘文件
|
||||
if mailAtts[0].Uploader != "calendar" {
|
||||
t.Errorf("uploader 应是 calendar，得到 %q", mailAtts[0].Uploader)
|
||||
}
|
||||
|
||||
// 没有附件的事件挂 0 个且不报错
|
||||
empty := seedEvent(t, &models.CalendarEvent{Title: "无附件"})
|
||||
if n, err := AttachCalendarFilesToMail(ctx, empty.EventID, mailID, "calendar"); err != nil || n != 0 {
|
||||
t.Errorf("无附件事件应挂 0 个，得到 %d err=%v", n, err)
|
||||
}
|
||||
}
|
||||
|
||||
func seedAgentForAttach(t *testing.T, name string) {
|
||||
t.Helper()
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO agents (agent_name, secret, platform) VALUES ($1, 'x', 'test')`,
|
||||
name); err != nil {
|
||||
t.Fatalf("seed agent: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func seedSessionForAttach(t *testing.T, agentName string) uuid.UUID {
|
||||
t.Helper()
|
||||
id := uuid.New()
|
||||
// 列名是 from_agent 而不是 agent_name（后者是 agents 表的主键名）
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO sessions (session_id, from_agent, subject, session_alias, status)
|
||||
VALUES ($1, $2, '日程提醒', 'cal-test', 'active')`,
|
||||
id, agentName); err != nil {
|
||||
t.Fatalf("seed session: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
// 落在 60 秒 lookahead 窗口内的**未来**事件，标记后不得再次到期。
|
||||
//
|
||||
// 生产实测的重发现场：一条 event_time=12:53:17 的事件在
|
||||
// 12:52:30 / 12:53:00 / 12:53:06 / 12:53:36 各发了一封相同提醒。
|
||||
// 根因是去重判据写成 `last_fired_at < event_time` —— 触发时刻（now）
|
||||
// 本来就早于 event_time，条件恒真，于是每个 tick 重发一次，
|
||||
// 直到 event_time 真正过去才自己停下。
|
||||
//
|
||||
// 改成按 occurrence 相等（fired_for = 当时的 event_time）才精确：
|
||||
// AdvanceRecurrence 改了 event_time 就该再触发，没改就永不重发。
|
||||
func TestFiredEventInLookaheadWindowDoesNotRefire(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 47 秒后 —— 在 lookahead 窗口内，所以第一次扫描就会入选
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "窗口内的未来事件", EventTime: time.Now().Add(47 * time.Second),
|
||||
})
|
||||
|
||||
due, err := DueEvents(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("首次扫描: %v", err)
|
||||
}
|
||||
if len(due) != 1 {
|
||||
t.Fatalf("lookahead 应让它提前入选，得到 %d 条", len(due))
|
||||
}
|
||||
|
||||
if err := MarkEventFired(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("标记: %v", err)
|
||||
}
|
||||
|
||||
// 模拟后续几个 tick
|
||||
for i := 0; i < 3; i++ {
|
||||
due, err = DueEvents(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次重扫: %v", i+2, err)
|
||||
}
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
t.Fatalf("第 %d 次扫描仍判定到期 —— 提醒会被重发", i+2)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 重复事件推进 event_time 之后必须重新到期：
|
||||
// 按 occurrence 去重的另一半，漏了它就变成「每个重复事件只响一次」。
|
||||
func TestRecurringEventRefiresAfterAdvance(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "每天都要响",
|
||||
EventTime: time.Now().Add(-time.Minute),
|
||||
Recurrence: "daily",
|
||||
})
|
||||
|
||||
if err := MarkEventFired(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("标记: %v", err)
|
||||
}
|
||||
due, _ := DueEvents(ctx)
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
t.Fatal("标记后不该立刻再次到期")
|
||||
}
|
||||
}
|
||||
|
||||
// 推进到下一次（+1 天）后，把时间挪到过去模拟「第二天到了」
|
||||
if _, err := AdvanceRecurrence(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("推进重复: %v", err)
|
||||
}
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE calendar_events SET event_time = ? WHERE event_id = ?`,
|
||||
time.Now().Add(-30*time.Second), e.EventID); err != nil {
|
||||
t.Fatalf("模拟次日: %v", err)
|
||||
}
|
||||
|
||||
due, _ = DueEvents(ctx)
|
||||
found := false
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("event_time 推进后应重新到期，否则重复事件只响一次")
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 重复规则推进（NextOccurrence 是纯函数，不碰数据库）───
|
||||
|
||||
func TestNextOccurrenceSolar(t *testing.T) {
|
||||
base := time.Date(2026, 9, 3, 9, 30, 0, 0, time.Local)
|
||||
|
||||
cases := []struct {
|
||||
rule string
|
||||
want string
|
||||
}{
|
||||
{models.RecurDaily, "2026-09-04"},
|
||||
{models.RecurWeekly, "2026-09-10"},
|
||||
{models.RecurMonthly, "2026-10-03"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got, err := NextOccurrence(c.rule, base)
|
||||
if err != nil {
|
||||
t.Errorf("%s: %v", c.rule, err)
|
||||
continue
|
||||
}
|
||||
if got.Format("2006-01-02") != c.want {
|
||||
t.Errorf("%s: 得到 %s，期望 %s", c.rule, got.Format("2006-01-02"), c.want)
|
||||
}
|
||||
// 时钟必须原样保留
|
||||
if got.Hour() != 9 || got.Minute() != 30 {
|
||||
t.Errorf("%s: 时钟被改动 %v", c.rule, got)
|
||||
}
|
||||
}
|
||||
|
||||
// none 与未知值都返回零值 + nil error
|
||||
for _, r := range []string{models.RecurNone, "", "每隔一个蓝月亮"} {
|
||||
got, err := NextOccurrence(r, base)
|
||||
if err != nil || !got.IsZero() {
|
||||
t.Errorf("%q 应返回零值无错，得到 %v err=%v", r, got, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// time.AddDate 的溢出对「每月同一日」是错的：3 月 31 日 +1 月 = 5 月 1 日。
|
||||
// 一次溢出会永久改变规则 —— 31 日的事件在 2 月变成 3 月 3 日，
|
||||
// 然后从此每月 3 日提醒。
|
||||
func TestNextOccurrenceMonthlyClampsMonthEnd(t *testing.T) {
|
||||
cases := []struct {
|
||||
from string
|
||||
want string
|
||||
why string
|
||||
}{
|
||||
{"2026-01-31", "2026-02-28", "1月31日 +1月 → 2月末（2026 非闰年）"},
|
||||
{"2026-03-31", "2026-04-30", "3月31日 +1月 → 4月30日"},
|
||||
{"2026-05-31", "2026-06-30", "5月31日 +1月 → 6月30日"},
|
||||
{"2028-01-31", "2028-02-29", "闰年 2 月有 29 天"},
|
||||
{"2026-01-15", "2026-02-15", "月中日期不受影响"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
from, _ := time.ParseInLocation("2006-01-02", c.from, time.Local)
|
||||
got, err := NextOccurrence(models.RecurMonthly, from)
|
||||
if err != nil {
|
||||
t.Errorf("%s: %v", c.why, err)
|
||||
continue
|
||||
}
|
||||
if got.Format("2006-01-02") != c.want {
|
||||
t.Errorf("%s: 得到 %s，期望 %s", c.why, got.Format("2006-01-02"), c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 农历月推进：公历间隔在 29~30 天之间浮动，不是固定值。
|
||||
// 这正是不能用 AddDate 的原因。
|
||||
func TestNextOccurrenceLunarMonthly(t *testing.T) {
|
||||
// 2026-09-03 = 农历七月廿二
|
||||
cur := time.Date(2026, 9, 3, 9, 0, 0, 0, time.Local)
|
||||
gaps := map[int]bool{}
|
||||
for i := 0; i < 6; i++ {
|
||||
next, err := NextOccurrence(models.RecurLunarMonthly, cur)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次推进: %v", i+1, err)
|
||||
}
|
||||
if !next.After(cur) {
|
||||
t.Fatalf("第 %d 次推进没有前进：%v → %v", i+1, cur, next)
|
||||
}
|
||||
gap := int(next.Sub(cur).Hours() / 24)
|
||||
gaps[gap] = true
|
||||
// 农历同一日：连续推进后农历「日」应保持
|
||||
if d := lunar.FromSolar(next); d.Day != 22 {
|
||||
t.Errorf("第 %d 次推进后农历日变成 %d（期望 22）：%s", i+1, d.Day, d.String())
|
||||
}
|
||||
cur = next
|
||||
}
|
||||
// 间隔必须出现过多种值，证明不是固定天数
|
||||
if len(gaps) < 2 {
|
||||
t.Errorf("六次农历月推进的公历间隔只有 %v —— 疑似退化成固定天数", gaps)
|
||||
}
|
||||
for g := range gaps {
|
||||
if g < 28 || g > 31 {
|
||||
t.Errorf("农历月间隔 %d 天不合理", g)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 农历年推进：公历日期每年漂移。用公历 yearly 会固定在同一天，
|
||||
// 与「过农历生日/祭日」的期望不符 —— 这是农历规则存在的理由。
|
||||
func TestNextOccurrenceLunarYearly(t *testing.T) {
|
||||
cur := time.Date(2026, 9, 3, 9, 0, 0, 0, time.Local)
|
||||
seen := map[string]bool{}
|
||||
for i := 0; i < 5; i++ {
|
||||
next, err := NextOccurrence(models.RecurLunarYearly, cur)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次: %v", i+1, err)
|
||||
}
|
||||
if !next.After(cur) {
|
||||
t.Fatalf("第 %d 次没有前进：%v → %v", i+1, cur, next)
|
||||
}
|
||||
// 农历月日应保持
|
||||
d := lunar.FromSolar(next)
|
||||
if d.Month != 7 || d.Day != 22 {
|
||||
t.Errorf("第 %d 次推进后农历变成 %d-%d（期望 7-22）", i+1, d.Month, d.Day)
|
||||
}
|
||||
seen[next.Format("01-02")] = true
|
||||
cur = next
|
||||
}
|
||||
if len(seen) < 3 {
|
||||
t.Errorf("五年公历月日只有 %d 种 —— 农历年重复应漂移", len(seen))
|
||||
}
|
||||
}
|
||||
|
||||
// 农历规则经过数据库这一轮也要正确（AdvanceRecurrence 里调 NextOccurrence）。
|
||||
func TestAdvanceRecurrenceLunar(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
start := time.Date(2026, 9, 3, 9, 0, 0, 0, time.Local)
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "农历每月十五（这里用廿二）",
|
||||
EventTime: start,
|
||||
Recurrence: models.RecurLunarMonthly,
|
||||
})
|
||||
|
||||
advanced, err := AdvanceRecurrence(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if !advanced {
|
||||
t.Fatal("农历重复应能推进")
|
||||
}
|
||||
|
||||
after, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读回: %v", err)
|
||||
}
|
||||
if !after.EventTime.After(start) {
|
||||
t.Errorf("event_time 未前进：%v", after.EventTime)
|
||||
}
|
||||
// 农历日保持
|
||||
if d := lunar.FromSolar(after.EventTime); d.Day != 22 {
|
||||
t.Errorf("农历日变成 %d，期望 22（%s）", d.Day, d.String())
|
||||
}
|
||||
// 公历间隔应在一个农历月内
|
||||
gap := int(after.EventTime.Sub(start).Hours() / 24)
|
||||
if gap < 28 || gap > 31 {
|
||||
t.Errorf("间隔 %d 天不像一个农历月", gap)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 多收件人 ───
|
||||
|
||||
func TestRecipientsRoundtrip(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "三个 Agent 各自汇报",
|
||||
Recipients: []string{"pi@/home/program/agentmail", "dsh", "opencode@/tmp"},
|
||||
DeliveryMode: models.DeliverSeparate,
|
||||
})
|
||||
|
||||
got, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读回: %v", err)
|
||||
}
|
||||
if len(got.Recipients) != 3 {
|
||||
t.Fatalf("收件人应有 3 个，得到 %d：%v", len(got.Recipients), got.Recipients)
|
||||
}
|
||||
if got.Recipients[0] != "pi@/home/program/agentmail" {
|
||||
t.Errorf("顺序或内容不符：%v", got.Recipients)
|
||||
}
|
||||
if got.EffectiveDeliveryMode() != models.DeliverSeparate {
|
||||
t.Errorf("投递模式 = %q", got.EffectiveDeliveryMode())
|
||||
}
|
||||
}
|
||||
|
||||
// 空收件人列表必须序列化成 []（而不是 null）：Go 的 nil slice 会变 null，
|
||||
// 前端 .map 直接崩。
|
||||
func TestRecipientsNeverNull(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{Title: "没写收件人"})
|
||||
got, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读回: %v", err)
|
||||
}
|
||||
if got.Recipients == nil {
|
||||
t.Error("Recipients 为 nil —— 会序列化成 null 让前端崩")
|
||||
}
|
||||
if len(got.Recipients) != 0 {
|
||||
t.Errorf("应是空数组，得到 %v", got.Recipients)
|
||||
}
|
||||
}
|
||||
|
||||
// 旧数据（只有 agent_name / to_address）必须继续工作 —— 历史事件不迁移。
|
||||
func TestEffectiveRecipientsFallbackChain(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
e models.CalendarEvent
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
"Recipients 优先",
|
||||
models.CalendarEvent{Recipients: []string{"a", "b"}, ToAddress: "c", AgentName: "d"},
|
||||
[]string{"a", "b"},
|
||||
},
|
||||
{
|
||||
"退回 to_address",
|
||||
models.CalendarEvent{ToAddress: "pi@/tmp.alias", AgentName: "pi"},
|
||||
[]string{"pi@/tmp.alias"},
|
||||
},
|
||||
{
|
||||
"再退回 agent_name",
|
||||
models.CalendarEvent{AgentName: "dsh"},
|
||||
[]string{"dsh"},
|
||||
},
|
||||
{
|
||||
"全空给 nil",
|
||||
models.CalendarEvent{},
|
||||
nil,
|
||||
},
|
||||
{
|
||||
"Recipients 里全是空白时继续退回",
|
||||
models.CalendarEvent{Recipients: []string{"", " "}, AgentName: "pi"},
|
||||
[]string{"pi"},
|
||||
},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := c.e.EffectiveRecipients()
|
||||
if len(got) != len(c.want) {
|
||||
t.Errorf("%s: 得到 %v，期望 %v", c.name, got, c.want)
|
||||
continue
|
||||
}
|
||||
for i := range got {
|
||||
if got[i] != c.want[i] {
|
||||
t.Errorf("%s: 第 %d 项 %q，期望 %q", c.name, i, got[i], c.want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 未知投递模式按 separate 处理：它的失败模式更轻。
|
||||
// together 用错会让本该独立判断的 Agent 互相看到回复而趋同，事后无法分离。
|
||||
func TestEffectiveDeliveryModeDefaultsToSeparate(t *testing.T) {
|
||||
for _, in := range []string{"", "separate", "垃圾值", "SEPARATE"} {
|
||||
e := models.CalendarEvent{DeliveryMode: in}
|
||||
if got := e.EffectiveDeliveryMode(); got != models.DeliverSeparate {
|
||||
t.Errorf("DeliveryMode=%q → %q，期望 separate", in, got)
|
||||
}
|
||||
}
|
||||
e := models.CalendarEvent{DeliveryMode: models.DeliverTogether}
|
||||
if e.EffectiveDeliveryMode() != models.DeliverTogether {
|
||||
t.Error("together 应被保留")
|
||||
}
|
||||
}
|
||||
|
||||
// 公历每年：2 月 29 日在平年必须夹到 2 月 28，不能溢出成 3 月 1 日。
|
||||
// 闰日生日的约定是「平年过 2 月 28」。
|
||||
func TestNextOccurrenceYearlyClampsLeapDay(t *testing.T) {
|
||||
cases := []struct {
|
||||
from string
|
||||
want string
|
||||
why string
|
||||
}{
|
||||
{"2028-02-29", "2029-02-28", "闰日 +1 年 → 平年 2 月 28"},
|
||||
{"2026-03-15", "2027-03-15", "普通日期不受影响"},
|
||||
{"2027-02-28", "2028-02-28", "平年 2/28 → 闰年仍是 2/28（不跳到 29）"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
from, _ := time.ParseInLocation("2006-01-02", c.from, time.Local)
|
||||
got, err := NextOccurrence(models.RecurYearly, from)
|
||||
if err != nil {
|
||||
t.Errorf("%s: %v", c.why, err)
|
||||
continue
|
||||
}
|
||||
if got.Format("2006-01-02") != c.want {
|
||||
t.Errorf("%s: 得到 %s，期望 %s", c.why, got.Format("2006-01-02"), c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 过期的重复事件必须一次推到未来，不能每轮补发一封。
|
||||
//
|
||||
// 实测的 bug：AdvanceRecurrence 只推进一步 —— 一条 100 天前设的每日事件，
|
||||
// 每轮扫描都判定「已过期该触发」→ 发一封 → event_time 只前进一天 →
|
||||
// 下一轮又过期。30 轮扫描触发 30 次，而调度周期是 30 秒，
|
||||
// 人会收到一串垃圾提醒，连发 100 封才追上今天。
|
||||
func TestStaleRecurringEventDoesNotFlood(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "很久以前设的每日提醒",
|
||||
EventTime: time.Now().AddDate(0, 0, -100),
|
||||
Recurrence: models.RecurDaily,
|
||||
})
|
||||
|
||||
fires := 0
|
||||
// 模拟调度器连续跑 30 轮（生产上就是 15 分钟）
|
||||
for i := 0; i < 30; i++ {
|
||||
due, err := DueEvents(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 轮扫描: %v", i+1, err)
|
||||
}
|
||||
hit := false
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
hit = true
|
||||
}
|
||||
}
|
||||
if !hit {
|
||||
break
|
||||
}
|
||||
fires++
|
||||
if err := MarkEventFired(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("标记: %v", err)
|
||||
}
|
||||
if _, err := AdvanceRecurrence(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
if fires != 1 {
|
||||
t.Errorf("过期的每日重复事件触发了 %d 次，应只触发 1 次", fires)
|
||||
}
|
||||
after, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读回: %v", err)
|
||||
}
|
||||
if !after.EventTime.After(time.Now()) {
|
||||
t.Errorf("推进后 event_time 仍在过去：%v", after.EventTime)
|
||||
}
|
||||
// 只跳到「刚过现在」的那一次，不是跳到很远的将来
|
||||
if after.EventTime.After(time.Now().AddDate(0, 0, 2)) {
|
||||
t.Errorf("推得太远了：%v", after.EventTime)
|
||||
}
|
||||
}
|
||||
|
||||
// 农历规则的过期事件同样不能刷屏。
|
||||
func TestStaleLunarRecurringDoesNotFlood(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "去年设的农历每月提醒",
|
||||
EventTime: time.Now().AddDate(-1, 0, 0),
|
||||
Recurrence: models.RecurLunarMonthly,
|
||||
})
|
||||
|
||||
if _, err := AdvanceRecurrence(ctx, e.EventID); err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
after, _ := GetCalendarEvent(ctx, e.EventID)
|
||||
if !after.EventTime.After(time.Now()) {
|
||||
t.Errorf("一年前的农历事件推进后仍在过去：%v", after.EventTime)
|
||||
}
|
||||
// 农历日必须保持
|
||||
if d := lunar.FromSolar(after.EventTime); d.Day != lunar.FromSolar(e.EventTime).Day {
|
||||
t.Errorf("农历日从 %d 变成 %d", lunar.FromSolar(e.EventTime).Day, d.Day)
|
||||
}
|
||||
}
|
||||
|
||||
// 越过 recurrence_end 时必须置 cancelled 而不是留在 active。
|
||||
//
|
||||
// 留着的表现是一条僵尸事件：DueEvents 每轮都捞到它（event_time 在过去），
|
||||
// 但 fired_for 已等于 event_time 所以又不触发 —— 永远排在到期列表里不动。
|
||||
func TestAdvanceCancelsAfterRecurrenceEnd(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
end := time.Now().AddDate(0, 0, -1) // 昨天就该停
|
||||
e := seedEvent(t, &models.CalendarEvent{
|
||||
Title: "已到期的每日重复",
|
||||
EventTime: time.Now().AddDate(0, 0, -5),
|
||||
Recurrence: models.RecurDaily,
|
||||
RecurrenceEnd: &end,
|
||||
})
|
||||
|
||||
advanced, err := AdvanceRecurrence(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if advanced {
|
||||
t.Error("已过 recurrence_end 不该报告推进成功")
|
||||
}
|
||||
after, err := GetCalendarEvent(ctx, e.EventID)
|
||||
if err != nil {
|
||||
t.Fatalf("读回: %v", err)
|
||||
}
|
||||
if after.Status != "cancelled" {
|
||||
t.Errorf("状态应是 cancelled，得到 %q —— 留在 active 会变僵尸事件", after.Status)
|
||||
}
|
||||
// 且不该再出现在到期列表里
|
||||
due, _ := DueEvents(ctx)
|
||||
for _, d := range due {
|
||||
if d.EventID == e.EventID {
|
||||
t.Error("已 cancelled 的事件仍出现在 DueEvents")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// advanceToFuture 是纯函数，单独测三个终止条件。
|
||||
func TestAdvanceToFuture(t *testing.T) {
|
||||
now := time.Date(2026, 9, 3, 12, 0, 0, 0, time.Local)
|
||||
|
||||
t.Run("跨过 now 就停", func(t *testing.T) {
|
||||
from := now.AddDate(0, 0, -100)
|
||||
got, err := advanceToFuture(models.RecurDaily, from, now, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
if !got.After(now) {
|
||||
t.Errorf("结果 %v 不在 now 之后", got)
|
||||
}
|
||||
// 恰好是越过 now 的第一次，不是更远
|
||||
if got.After(now.AddDate(0, 0, 1)) {
|
||||
t.Errorf("推过头了：%v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("越过 recurrenceEnd 给零值", func(t *testing.T) {
|
||||
end := now.AddDate(0, 0, -1)
|
||||
got, err := advanceToFuture(models.RecurDaily, now.AddDate(0, 0, -5), now, &end)
|
||||
if err != nil {
|
||||
t.Fatalf("不该报错：%v", err)
|
||||
}
|
||||
if !got.IsZero() {
|
||||
t.Errorf("应给零值，得到 %v", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("不重复给零值无错", func(t *testing.T) {
|
||||
got, err := advanceToFuture(models.RecurNone, now, now, nil)
|
||||
if err != nil || !got.IsZero() {
|
||||
t.Errorf("得到 %v err=%v", got, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("未来的事件原地推一步", func(t *testing.T) {
|
||||
from := now.AddDate(0, 0, 5)
|
||||
got, err := advanceToFuture(models.RecurDaily, from, now, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("推进: %v", err)
|
||||
}
|
||||
// from 已在未来，第一次推进就该返回
|
||||
if !got.Equal(from.AddDate(0, 0, 1)) {
|
||||
t.Errorf("得到 %v，期望 %v", got, from.AddDate(0, 0, 1))
|
||||
}
|
||||
})
|
||||
|
||||
// 上限是防御性的：农历路径依赖外部库，一旦某年给出反直觉结果，
|
||||
// 没有上限就是个死循环 goroutine，而它跑在调度器里 —— 整个提醒系统一起卡住
|
||||
t.Run("十年前的每日事件也能在上限内追上", func(t *testing.T) {
|
||||
got, err := advanceToFuture(models.RecurDaily, now.AddDate(-10, 0, 0), now, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("十年（约 3650 步）应在 %d 上限内：%v", maxAdvanceSteps, err)
|
||||
}
|
||||
if !got.After(now) {
|
||||
t.Errorf("结果 %v 不在 now 之后", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
269
server/internal/repo/defaultsession_test.go
Normal file
269
server/internal/repo/defaultsession_test.go
Normal file
@ -0,0 +1,269 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 这一组测试钉住「省略 session 位复用默认会话」与「真的新建会话」必须可区分。
|
||||
//
|
||||
// 事故背景：handler 层曾用 `parentMailID == nil` 判断「是不是新建会话」，
|
||||
// 据此决定要不要写往返预算与权限档位。但省略 session 位复用默认会话时
|
||||
// parentMailID 也是 nil —— 于是每一封续谈的信都会把这两个字段重置成默认值。
|
||||
//
|
||||
// 线上实测（修复前）：
|
||||
//
|
||||
// 第一封 to=pi@/tmp/budgetprobe max_rounds=7 → budget_max 7
|
||||
// 第二封 to=pi@/tmp/budgetprobe（省略该字段） → budget_max 20 ← 被静默改写
|
||||
//
|
||||
// 而那段代码的注释本身正在论证这不该发生（「续谈已有会话若也接受这个字段，
|
||||
// 每封新信都会悄悄改掉对方正在遵守的预算」）—— 意图是对的，守卫写错了。
|
||||
//
|
||||
// 修法：FindOrCreateDefaultSessionCreated 额外返回 created，
|
||||
// handler 改用它而不是 parentMailID。
|
||||
|
||||
func TestDefaultSessionFirstCallCreates(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !created {
|
||||
t.Fatal("从未通信过的 name@path，第一次必须报告 created=true")
|
||||
}
|
||||
if id.String() == "" {
|
||||
t.Fatal("应返回有效会话 id")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultSessionReuseReportsNotCreated(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
first, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil || !created {
|
||||
t.Fatalf("首封应新建：created=%v err=%v", created, err)
|
||||
}
|
||||
// 复用的前提是这条会话里有该收件人参与过的邮件（EXISTS 子查询）
|
||||
if _, err := CreateMail(ctx, first, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
second, created2, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second != first {
|
||||
t.Fatalf("第二封应复用同一条默认会话：first=%s second=%s", first, second)
|
||||
}
|
||||
if created2 {
|
||||
t.Fatal("复用已有默认会话时 created 必须为 false —— 这正是预算被冲掉的根因")
|
||||
}
|
||||
}
|
||||
|
||||
// 这条是上面那个线上事故的最小复现：走 created 判据时预算不被改写。
|
||||
func TestBudgetSurvivesDefaultSessionReuse(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil || !created {
|
||||
t.Fatalf("首封应新建：created=%v err=%v", created, err)
|
||||
}
|
||||
// 模拟 handler：只有 created 为真才设预算
|
||||
if _, err := SetSessionBudget(ctx, id, 7); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, id, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, created2, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created2 {
|
||||
// 若这里为真，handler 就会重设预算 —— 事故重现
|
||||
t.Fatal("复用时 created 为真会让 handler 重设预算")
|
||||
}
|
||||
|
||||
b, err := GetSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.Max != 7 {
|
||||
t.Fatalf("续谈不得改写预算：want 7, got %d", b.Max)
|
||||
}
|
||||
}
|
||||
|
||||
// 档位与预算同一个判据，一起钉住：plan 档不能因为第二封信而升成 workspace。
|
||||
func TestPermissionModeSurvivesDefaultSessionReuse(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil || !created {
|
||||
t.Fatalf("首封应新建：created=%v err=%v", created, err)
|
||||
}
|
||||
if _, err := SetSessionPermissionMode(ctx, id, "plan"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, id, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
_, created2, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created2 {
|
||||
t.Fatal("复用时 created 为真会让 handler 把档位重置成默认档")
|
||||
}
|
||||
|
||||
if got := SessionPermissionMode(ctx, id); got != "plan" {
|
||||
t.Fatalf("续谈不得改写档位：want plan, got %s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 不同工作目录是不同的默认会话，两边各自新建。
|
||||
// 这条防的是「把 created 实现成一个全局标志」之类的偷懒写法。
|
||||
func TestDefaultSessionPerWorkspace(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
a, createdA, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w1", "jianf", "甲")
|
||||
if err != nil || !createdA {
|
||||
t.Fatalf("/w1 应新建：%v %v", createdA, err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, a, nil, "jianf", "", "pi", "/w1", "甲", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
b, createdB, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w2", "jianf", "乙")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !createdB {
|
||||
t.Fatal("/w2 是另一个工作目录，应当另建一条默认会话")
|
||||
}
|
||||
if a == b {
|
||||
t.Fatal("不同工作目录不该共用同一条默认会话")
|
||||
}
|
||||
}
|
||||
|
||||
// 旧签名仍在别处被调用，保持行为不变（只是丢掉 created）。
|
||||
func TestLegacyWrapperStillWorks(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := FindOrCreateDefaultSession(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, id, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
again, err := FindOrCreateDefaultSession(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if again != id {
|
||||
t.Fatalf("包装函数应与原行为一致：%s vs %s", id, again)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 预算不被冲的回归用例（P0 第一项，P3 钉死） ───
|
||||
//
|
||||
// 守卫原本是 `parentMailID == nil`（表示新建），但省略 session 位复用默认会话时
|
||||
// parentMailID 也是 nil —— 于是「仅在新建时生效」的字段（预算、档位）在每封
|
||||
// 省略 session 位的信上都被重写了。实测：第一封 max_rounds=7 → 第二封省略该
|
||||
// 字段 → 预算被静默改成默认的 20。
|
||||
//
|
||||
// 修法：handler 改用 FindOrCreateDefaultSessionCreated 返回的 `created` 判据。
|
||||
// 本测试钉死「复用默认会话时 created=false」这一事实，让守卫不会倒退回去。
|
||||
func TestDefaultSessionReuseBudgetNotReset(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
first, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil || !created {
|
||||
t.Fatalf("首封应新建：created=%v err=%v", created, err)
|
||||
}
|
||||
// 模拟发信路径：新建会话时定预算为 7（低于默认 20）
|
||||
if _, err := SetSessionBudget(ctx, first, 7); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 复用的前提是会话里有该收件人参与过的邮件
|
||||
if _, err := CreateMail(ctx, first, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 第二封省略 session 位 → 复用默认会话，created 必须为 false
|
||||
second, created2, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second != first {
|
||||
t.Fatalf("第二封应复用同一会话：%s vs %s", first, second)
|
||||
}
|
||||
if created2 {
|
||||
t.Fatal("复用默认会话时 created 必须为 false，否则预算会被默认值冲掉")
|
||||
}
|
||||
|
||||
// 既然 created2=false，发信路径不会调 SetSessionBudget → 预算仍为 7
|
||||
b, err := GetSessionBudget(ctx, first)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.Max != 7 {
|
||||
t.Errorf("预算被冲掉：got max=%d，want 7（复用不应重设）", b.Max)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 档位同样不被冲的回归用例 ───
|
||||
//
|
||||
// 与预算同理：复用默认会话时档位也不该被重设成默认档。
|
||||
// 人指定 plan 档后，第二封信省略 session 位复用同一条会话 →
|
||||
// created=false → SetSessionPermissionMode 不被调 → 档位仍为 plan。
|
||||
func TestDefaultSessionReusePermissionModeNotReset(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
first, created, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "首封")
|
||||
if err != nil || !created {
|
||||
t.Fatalf("首封应新建：created=%v err=%v", created, err)
|
||||
}
|
||||
// 新建会话时定档位为 plan（比默认 workspace 更严）
|
||||
if _, err := SetSessionPermissionMode(ctx, first, "plan"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, first, nil, "jianf", "", "pi", "/w", "首封", "x", nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
second, created2, err := FindOrCreateDefaultSessionCreated(ctx, "pi", "/w", "jianf", "第二封")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if second != first {
|
||||
t.Fatalf("第二封应复用同一会话：%s vs %s", first, second)
|
||||
}
|
||||
if created2 {
|
||||
t.Fatal("复用默认会话时 created 必须为 false，否则档位会被默认值冲掉")
|
||||
}
|
||||
|
||||
mode := SessionPermissionMode(ctx, first)
|
||||
if mode != "plan" {
|
||||
t.Errorf("档位被冲掉：got %q，want plan（复用不应重设）", mode)
|
||||
}
|
||||
}
|
||||
135
server/internal/repo/deliverable_test.go
Normal file
135
server/internal/repo/deliverable_test.go
Normal file
@ -0,0 +1,135 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// RecipientDeliverable 是「邮件会不会进黑洞」的唯一判据。
|
||||
//
|
||||
// 事故背景：发信路径原来只校验地址语法、调用权限与会话别名，从不问
|
||||
// 「这个名字存在吗」。实测发给已彻底删除的 remotebot 返回 200，邮件入库、
|
||||
// 分配了 20 个来回预算、建好会话，而那一端永远不会有人读。发件人看到 200
|
||||
// 和一个 session_id，以为送出去了 —— 静默丢件比报错严重，报错能立刻改，
|
||||
// 静默丢件要等对方追问才发现。
|
||||
|
||||
func TestRecipientDeliverable_Human(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO users (username, password_hash) VALUES ('alice', 'x')`); err != nil {
|
||||
t.Fatalf("seed user: %v", err)
|
||||
}
|
||||
|
||||
if err := RecipientDeliverable(ctx, "alice"); err != nil {
|
||||
t.Fatalf("人类用户应当可达，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecipientDeliverable_OnlineAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
if err := RecipientDeliverable(ctx, "pi"); err != nil {
|
||||
t.Fatalf("在册且未停用的 Agent 应当可达，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecipientDeliverable_UnknownName(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
err := RecipientDeliverable(ctx, "ghost")
|
||||
if !errors.Is(err, ErrRecipientUnknown) {
|
||||
t.Fatalf("不存在的收件人应当返回 ErrRecipientUnknown，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 删除后立刻不可达 —— 这正是本次事故的场景。
|
||||
func TestRecipientDeliverable_AfterDelete(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "remotebot", 20)
|
||||
|
||||
if err := RecipientDeliverable(ctx, "remotebot"); err != nil {
|
||||
t.Fatalf("删除前应当可达，得到: %v", err)
|
||||
}
|
||||
|
||||
if _, err := DeleteAgent(ctx, "remotebot"); err != nil {
|
||||
t.Fatalf("delete agent: %v", err)
|
||||
}
|
||||
|
||||
err := RecipientDeliverable(ctx, "remotebot")
|
||||
if !errors.Is(err, ErrRecipientUnknown) {
|
||||
t.Fatalf("删除后必须不可达，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 停用选择「当场拒收」而不是「入库等恢复后补投」：停用的语义就是这个
|
||||
// Agent 现在不干活，让发件人以为信已送达更坏 —— 它会照常等回信。
|
||||
func TestRecipientDeliverable_Disabled(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "opencode", 20)
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "opencode", true); err != nil {
|
||||
t.Fatalf("disable: %v", err)
|
||||
}
|
||||
|
||||
err := RecipientDeliverable(ctx, "opencode")
|
||||
if !errors.Is(err, ErrRecipientDisabled) {
|
||||
t.Fatalf("已停用的 Agent 应当返回 ErrRecipientDisabled，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 恢复后重新可达，否则停用就成了不可逆操作。
|
||||
func TestRecipientDeliverable_ReenabledAgain(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "opencode", 20)
|
||||
|
||||
if _, err := SetAgentDisabled(ctx, "opencode", true); err != nil {
|
||||
t.Fatalf("disable: %v", err)
|
||||
}
|
||||
if _, err := SetAgentDisabled(ctx, "opencode", false); err != nil {
|
||||
t.Fatalf("re-enable: %v", err)
|
||||
}
|
||||
|
||||
if err := RecipientDeliverable(ctx, "opencode"); err != nil {
|
||||
t.Fatalf("恢复后应当重新可达，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 空名交给上层的地址解析处理，这里放行 —— 不然人类给自己发信
|
||||
// （to 位省略 name）会被这道检查误伤。
|
||||
func TestRecipientDeliverable_EmptyName(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
if err := RecipientDeliverable(context.Background(), ""); err != nil {
|
||||
t.Fatalf("空名应当放行，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 人类用户名与 Agent 名共用命名空间。一个名字同时是人类用户时按人类算 ——
|
||||
// 人的收件箱一直在，不受 Agent 停用影响。
|
||||
func TestRecipientDeliverable_HumanWinsOverDisabledAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO users (username, password_hash) VALUES ('dual', 'x')`); err != nil {
|
||||
t.Fatalf("seed user: %v", err)
|
||||
}
|
||||
seedAgent(t, "dual", 20)
|
||||
if _, err := SetAgentDisabled(ctx, "dual", true); err != nil {
|
||||
t.Fatalf("disable: %v", err)
|
||||
}
|
||||
|
||||
if err := RecipientDeliverable(ctx, "dual"); err != nil {
|
||||
t.Fatalf("同名人类用户应当优先放行，得到: %v", err)
|
||||
}
|
||||
}
|
||||
351
server/internal/repo/keys.go
Normal file
351
server/internal/repo/keys.go
Normal file
@ -0,0 +1,351 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 密钥认证 ----------
|
||||
//
|
||||
// 两类密钥共享一个全局唯一的 token 命名空间：验证时先查 agent_keys 再查 user_keys。
|
||||
// 这样一个 token 永远只有一种身份，不会出现「同一串字符既能注册 Agent 又能读人类邮箱」。
|
||||
|
||||
var (
|
||||
// ErrKeyNotFound 密钥不存在
|
||||
ErrKeyNotFound = errors.New("key not found")
|
||||
// ErrKeyUsed 一次性密钥已被使用
|
||||
ErrKeyUsed = errors.New("key already used")
|
||||
// ErrKeyExpired 定时密钥已过期
|
||||
ErrKeyExpired = errors.New("key expired")
|
||||
// ErrKeyTypeInvalid 密钥类型不受支持
|
||||
ErrKeyTypeInvalid = errors.New("invalid key type")
|
||||
// ErrKeyNeedsExpiry timed 密钥缺少有效的 expires_hours
|
||||
ErrKeyNeedsExpiry = errors.New("timed key requires positive expires_hours")
|
||||
// ErrKeyTooShort 登记的客户端密钥长度不足
|
||||
ErrKeyTooShort = errors.New("key token too short")
|
||||
)
|
||||
|
||||
// expiryFor 依据密钥类型算出过期时间。
|
||||
// 只有 timed 需要 expires_at；permanent 与 one_time 都是 NULL，
|
||||
// 各自的失效条件由 key_type 本身表达，不混用 expires_at。
|
||||
func expiryFor(keyType string, hours int) (*time.Time, error) {
|
||||
if !models.ValidKeyType(keyType) {
|
||||
return nil, ErrKeyTypeInvalid
|
||||
}
|
||||
if keyType != models.KeyTimed {
|
||||
return nil, nil
|
||||
}
|
||||
if hours <= 0 {
|
||||
return nil, ErrKeyNeedsExpiry
|
||||
}
|
||||
t := time.Now().Add(time.Duration(hours) * time.Hour)
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// checkKeyUsable 判断一条密钥记录当前是否可用。
|
||||
func checkKeyUsable(keyType string, expiresAt, usedAt *time.Time) error {
|
||||
switch keyType {
|
||||
case models.KeyOneTime:
|
||||
if usedAt != nil {
|
||||
return ErrKeyUsed
|
||||
}
|
||||
case models.KeyTimed:
|
||||
if expiresAt == nil || time.Now().After(*expiresAt) {
|
||||
return ErrKeyExpired
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------- Agent 密钥（管理员签发） ----------
|
||||
|
||||
// ErrKeyTokenTaken 登记的密钥已被占用
|
||||
var ErrKeyTokenTaken = errors.New("key token already registered")
|
||||
|
||||
// CreateAgentKey 签发一条 Agent 接入密钥。agentName 为空表示待绑定。
|
||||
//
|
||||
// presetToken 非空时登记客户端已在本地生成的密钥（插件首装场景），
|
||||
// 这样密钥全文只从客户端往服务器走一次，不需要反方向传递；留空则由服务器生成。
|
||||
func CreateAgentKey(ctx context.Context, agentName, keyType, label string, expiresHours int, createdBy uuid.UUID, presetToken string) (*models.AgentKey, error) {
|
||||
expires, err := expiryFor(keyType, expiresHours)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
token := presetToken
|
||||
if token == "" {
|
||||
if token, err = newToken(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if len(token) < 32 {
|
||||
// 太短的客户端密钥不接受，否则等于把弱口令当凭证
|
||||
return nil, ErrKeyTooShort
|
||||
}
|
||||
|
||||
// 已退役的名字不可重建 —— 登记密钥会级联建 agents 行，绕过注册检查。
|
||||
if agentName != "" {
|
||||
if retired, rErr := IsRetiredAgentName(ctx, agentName); rErr != nil {
|
||||
return nil, rErr
|
||||
} else if retired {
|
||||
return nil, fmt.Errorf("该名字已退役，不可重建")
|
||||
}
|
||||
}
|
||||
|
||||
var namePtr *string
|
||||
if agentName != "" {
|
||||
namePtr = &agentName
|
||||
}
|
||||
|
||||
k := &models.AgentKey{
|
||||
Token: token,
|
||||
TokenHint: models.TokenHint(token),
|
||||
AgentName: namePtr,
|
||||
KeyType: keyType,
|
||||
Label: label,
|
||||
ExpiresAt: expires,
|
||||
CreatedBy: &createdBy,
|
||||
}
|
||||
err = db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO agent_keys (key_token, agent_name, key_type, label, expires_at, created_by)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
RETURNING key_id, created_at
|
||||
`, token, namePtr, keyType, label, expires, createdBy).Scan(&k.ID, &k.CreatedAt)
|
||||
if db.IsUniqueViolation(err) {
|
||||
return nil, ErrKeyTokenTaken
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return k, nil
|
||||
}
|
||||
|
||||
// ListAgentKeys 列出 Agent 密钥；agentName 非空时按 Agent 过滤。
|
||||
// 返回值不含 token 全文，只有 hint。
|
||||
func ListAgentKeys(ctx context.Context, agentName string) ([]models.AgentKey, error) {
|
||||
q := `SELECT key_id, key_token, agent_name, key_type, label, expires_at, used_at, created_by, created_at
|
||||
FROM agent_keys`
|
||||
args := []any{}
|
||||
if agentName != "" {
|
||||
q += ` WHERE agent_name = $1`
|
||||
args = append(args, agentName)
|
||||
}
|
||||
q += ` ORDER BY created_at DESC`
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx, q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
keys := []models.AgentKey{}
|
||||
for rows.Next() {
|
||||
var k models.AgentKey
|
||||
var token string
|
||||
if err := rows.Scan(&k.ID, &token, &k.AgentName, &k.KeyType, &k.Label,
|
||||
&k.ExpiresAt, &k.UsedAt, &k.CreatedBy, &k.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
k.TokenHint = models.TokenHint(token) // 不回传全文
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteAgentKey 吊销一条 Agent 密钥。
|
||||
func DeleteAgentKey(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx, `DELETE FROM agent_keys WHERE key_id = $1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BindAgentKey 把一条密钥绑定到指定 Agent 名。
|
||||
func BindAgentKey(ctx context.Context, id uuid.UUID, agentName string) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET agent_name = $2 WHERE key_id = $1`, id, agentName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// VerifyAgentKey 校验 Agent 密钥并返回它绑定的 Agent 名（未绑定时返回空串）。
|
||||
//
|
||||
// 一次性密钥在校验通过时立刻写 used_at —— 用 WHERE used_at IS NULL 保证并发下
|
||||
// 只有一个请求能把它标记掉，避免两个 Agent 拿同一把一次性密钥同时注册成功。
|
||||
func VerifyAgentKey(ctx context.Context, token string) (string, error) {
|
||||
var (
|
||||
id uuid.UUID
|
||||
agentName *string
|
||||
keyType string
|
||||
expiresAt *time.Time
|
||||
usedAt *time.Time
|
||||
)
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT key_id, agent_name, key_type, expires_at, used_at
|
||||
FROM agent_keys WHERE key_token = $1
|
||||
`, token).Scan(&id, &agentName, &keyType, &expiresAt, &usedAt)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", ErrKeyNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if keyType == models.KeyOneTime {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return "", ErrKeyUsed // 并发下被别人抢先用掉了
|
||||
}
|
||||
}
|
||||
|
||||
if agentName == nil {
|
||||
return "", nil
|
||||
}
|
||||
return *agentName, nil
|
||||
}
|
||||
|
||||
// ClaimAgentKey 在待绑定密钥首次注册时把它落定到该 Agent 名。
|
||||
// 已绑定的密钥不受影响（WHERE agent_name IS NULL），因此不能借一把已绑定的密钥改注册别的 Agent。
|
||||
func ClaimAgentKey(ctx context.Context, token, agentName string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET agent_name = $2 WHERE key_token = $1 AND agent_name IS NULL`,
|
||||
token, agentName)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- 用户密钥（用户自助签发） ----------
|
||||
|
||||
// CreateUserKey 为用户签发一条客户端连接密钥。
|
||||
func CreateUserKey(ctx context.Context, userID uuid.UUID, label, keyType string, expiresHours int) (*models.UserKey, error) {
|
||||
expires, err := expiryFor(keyType, expiresHours)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
token, err := newToken()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
k := &models.UserKey{
|
||||
Token: token,
|
||||
TokenHint: models.TokenHint(token),
|
||||
UserID: userID,
|
||||
Label: label,
|
||||
KeyType: keyType,
|
||||
ExpiresAt: expires,
|
||||
}
|
||||
err = db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO user_keys (key_token, user_id, label, key_type, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
RETURNING key_id, created_at
|
||||
`, token, userID, label, keyType, expires).Scan(&k.ID, &k.CreatedAt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return k, nil
|
||||
}
|
||||
|
||||
// ListUserKeys 列出某用户的连接密钥（不含 token 全文）。
|
||||
func ListUserKeys(ctx context.Context, userID uuid.UUID) ([]models.UserKey, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT key_id, key_token, user_id, label, key_type, expires_at, used_at, created_at
|
||||
FROM user_keys WHERE user_id = $1 ORDER BY created_at DESC
|
||||
`, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
keys := []models.UserKey{}
|
||||
for rows.Next() {
|
||||
var k models.UserKey
|
||||
var token string
|
||||
if err := rows.Scan(&k.ID, &token, &k.UserID, &k.Label, &k.KeyType,
|
||||
&k.ExpiresAt, &k.UsedAt, &k.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
k.TokenHint = models.TokenHint(token)
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteUserKey 删除自己的一条密钥。带 user_id 条件，避免删掉别人的。
|
||||
func DeleteUserKey(ctx context.Context, userID, keyID uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM user_keys WHERE key_id = $1 AND user_id = $2`, keyID, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// VerifyUserKey 校验用户密钥并返回对应用户。
|
||||
// 用户必须仍处于 active 状态——禁用账号后其密钥应当立即失效。
|
||||
func VerifyUserKey(ctx context.Context, token string) (*models.User, error) {
|
||||
var (
|
||||
id uuid.UUID
|
||||
userID uuid.UUID
|
||||
keyType string
|
||||
expiresAt *time.Time
|
||||
usedAt *time.Time
|
||||
)
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT key_id, user_id, key_type, expires_at, used_at
|
||||
FROM user_keys WHERE key_token = $1
|
||||
`, token).Scan(&id, &userID, &keyType, &expiresAt, &usedAt)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrKeyNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if keyType == models.KeyOneTime {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE user_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return nil, ErrKeyUsed
|
||||
}
|
||||
}
|
||||
|
||||
u, err := GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrKeyNotFound // 账号已禁用，密钥一并失效
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
144
server/internal/repo/markread_test.go
Normal file
144
server/internal/repo/markread_test.go
Normal file
@ -0,0 +1,144 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// seedMailTo 造一封给 recipient 的未读邮件，可选带抄送。
|
||||
func seedMailTo(t *testing.T, recipient string, cc string) uuid.UUID {
|
||||
t.Helper()
|
||||
ctx := context.Background()
|
||||
sid, err := CreateSession(ctx, nil, "sender", "t", "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ccJSON := "[]"
|
||||
if cc != "" {
|
||||
ccJSON = `[{"name":"` + cc + `","path":"","session":"","raw":"` + cc + `"}]`
|
||||
}
|
||||
var id uuid.UUID
|
||||
err = db.DB.QueryRowContext(ctx,
|
||||
`INSERT INTO mails (session_id, from_name, to_name, subject, body, cc_list)
|
||||
VALUES ($1, 'sender', $2, 's', 'b', $3) RETURNING mail_id`,
|
||||
sid, recipient, ccJSON).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func statusOf(t *testing.T, id uuid.UUID) string {
|
||||
t.Helper()
|
||||
var s string
|
||||
if err := db.DB.QueryRowContext(context.Background(),
|
||||
`SELECT status FROM mails WHERE mail_id = $1`, id).Scan(&s); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func TestMarkMailsReadForOnlyOwnMail(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
mine := seedMailTo(t, "bot", "")
|
||||
others := seedMailTo(t, "other", "")
|
||||
|
||||
// 一次请求里混着别人的邮件：自己的标掉，别人的动不了。
|
||||
// 鉴权写在 UPDATE 的 WHERE 里，所以这不是「先查后拒」而是根本改不动。
|
||||
n, err := MarkMailsReadFor(ctx, "bot", []uuid.UUID{mine, others})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("影响行数 = %d，期望 1（只有自己那封）", n)
|
||||
}
|
||||
if statusOf(t, mine) != "read" {
|
||||
t.Fatal("自己的邮件没被标记")
|
||||
}
|
||||
if statusOf(t, others) != "unread" {
|
||||
t.Fatal("别人的邮件被标记了 —— 鉴权失效")
|
||||
}
|
||||
}
|
||||
|
||||
// 重复标记是幂等的：Agent 通常把上一轮列出的 id 原样传回来，
|
||||
// 其中混着已读的不该算错误
|
||||
func TestMarkMailsReadForIsIdempotent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
id := seedMailTo(t, "bot", "")
|
||||
|
||||
if n, _ := MarkMailsReadFor(ctx, "bot", []uuid.UUID{id}); n != 1 {
|
||||
t.Fatalf("首次应标掉 1 封，实际 %d", n)
|
||||
}
|
||||
n, err := MarkMailsReadFor(ctx, "bot", []uuid.UUID{id})
|
||||
if err != nil {
|
||||
t.Fatalf("重复标记不该报错: %v", err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Fatalf("重复标记影响行数 = %d，期望 0", n)
|
||||
}
|
||||
}
|
||||
|
||||
// 被抄送的邮件也在收件箱里，也该能标掉
|
||||
func TestMarkMailsReadForCoversCC(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
id := seedMailTo(t, "other", "bot") // 主收件人是 other，bot 被抄送
|
||||
|
||||
n, err := MarkMailsReadFor(ctx, "bot", []uuid.UUID{id})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("被抄送的邮件应可标记，影响行数 = %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkMailsReadForEmptyList(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
// 空列表直接返回，不该拼出 `IN ()` 这种非法 SQL
|
||||
n, err := MarkMailsReadFor(context.Background(), "bot", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("空列表不该报错: %v", err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Fatalf("空列表影响行数 = %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkAllInboxReadForSkipsArchivedAndOthers(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
a := seedMailTo(t, "bot", "")
|
||||
b := seedMailTo(t, "bot", "")
|
||||
others := seedMailTo(t, "other", "")
|
||||
|
||||
// 把 b 所在会话归档：那封在收件箱里根本看不到，
|
||||
// 标掉它只会让「标记了 N 封」与用户看到的对不上
|
||||
var sid uuid.UUID
|
||||
db.DB.QueryRowContext(ctx, `SELECT session_id FROM mails WHERE mail_id = $1`, b).Scan(&sid)
|
||||
db.DB.ExecContext(ctx, `UPDATE sessions SET status = 'archived' WHERE session_id = $1`, sid)
|
||||
|
||||
n, err := MarkAllInboxReadFor(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("影响行数 = %d，期望 1（排除归档会话）", n)
|
||||
}
|
||||
if statusOf(t, a) != "read" {
|
||||
t.Fatal("活跃会话里的未读没被标掉")
|
||||
}
|
||||
if statusOf(t, b) != "unread" {
|
||||
t.Fatal("归档会话里的邮件被标掉了")
|
||||
}
|
||||
if statusOf(t, others) != "unread" {
|
||||
t.Fatal("别人的邮件被标掉了")
|
||||
}
|
||||
}
|
||||
238
server/internal/repo/models_scope.go
Normal file
238
server/internal/repo/models_scope.go
Normal file
@ -0,0 +1,238 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// ---------- 邮件场景下的可用模型 ----------
|
||||
//
|
||||
// 两张表，两种真相：
|
||||
//
|
||||
// agent_model_catalog —— 平台**上报**它当前看得见哪些模型（注册时整表替换）
|
||||
// agent_allowed_models —— 管理员**选定**其中哪些能在邮件场景下用，rank 即优先级
|
||||
//
|
||||
// 为什么不合成一张带 allowed 标记的表：模型会从平台目录里消失（换了 provider 配置、
|
||||
// 上游临时下线），那时整行被删掉就连带把管理员的选择也删了，模型回来还得重配一遍。
|
||||
// 分开存之后，「选了什么」是持久的，目录只决定「这一项现在是否可用」。
|
||||
//
|
||||
// 为什么让平台上报而不是在 Gateway 里配一张静态表：模型清单是平台侧的事实 ——
|
||||
// opencode 的 provider 配置、DSH 的 llm 适配器注册，都可能随时变。
|
||||
// Gateway 猜不出来，猜错的后果是管理员在配置页选了一个平台其实调不到的模型。
|
||||
|
||||
// ModelRef 是一次「provider + model」路由。
|
||||
type ModelRef struct {
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
}
|
||||
|
||||
// CatalogModel 是平台上报的一个可选模型。
|
||||
type CatalogModel struct {
|
||||
Provider string `json:"provider"`
|
||||
Model string `json:"model"`
|
||||
DisplayName string `json:"display_name,omitempty"`
|
||||
// Allowed 表示它已被管理员选入邮件场景。
|
||||
// 与目录合并后一起返回，前端才能画出「已勾选」的复选框。
|
||||
Allowed bool `json:"allowed"`
|
||||
// Rank 仅在 Allowed 为真时有意义，越小越先试。
|
||||
Rank int `json:"rank,omitempty"`
|
||||
}
|
||||
|
||||
// maxCatalogModels 限制单个 Agent 上报的模型数。
|
||||
//
|
||||
// 有平台会把上游的全部模型都列出来（实测 opencode 的一个 provider 就有几十个），
|
||||
// 无上限的话一次注册能写进几千行，而配置页面上几千个复选框对人毫无用处。
|
||||
const maxCatalogModels = 300
|
||||
|
||||
// ReplaceModelCatalog 整表替换某 Agent 上报的模型目录。
|
||||
//
|
||||
// 整表替换而非增量合并：目录是平台当前状态的快照，
|
||||
// 增量合并会让已经下线的模型永远留在列表里，而那正是「选了却调不到」的来源。
|
||||
//
|
||||
// 事务包住删+插：中途失败留下一个空目录，会让配置页显示「该平台没有可用模型」
|
||||
// 而管理员根本没做任何操作。
|
||||
func ReplaceModelCatalog(ctx context.Context, agentName string, models []CatalogModel) error {
|
||||
agentName = strings.TrimSpace(agentName)
|
||||
if agentName == "" {
|
||||
return nil
|
||||
}
|
||||
if len(models) > maxCatalogModels {
|
||||
models = models[:maxCatalogModels]
|
||||
}
|
||||
|
||||
tx, err := db.DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`DELETE FROM agent_model_catalog WHERE agent_name = $1`, agentName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
for _, m := range models {
|
||||
p := strings.TrimSpace(m.Provider)
|
||||
id := strings.TrimSpace(m.Model)
|
||||
if p == "" || id == "" {
|
||||
continue // 半条记录不如不要：它在配置页上是一个点不动的空复选框
|
||||
}
|
||||
key := p + "/" + id
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO agent_model_catalog (agent_name, provider, model, display_name, reported_at)
|
||||
VALUES ($1, $2, $3, $4, NOW())`,
|
||||
agentName, p, id, strings.TrimSpace(m.DisplayName)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// ListModelCatalog 返回某 Agent 的模型目录，并标出哪些已被选入邮件场景。
|
||||
//
|
||||
// LEFT JOIN 而不是两次查询：前端要的是一份「带勾选状态的清单」，
|
||||
// 在 SQL 里合完比让前端自己对齐两个数组更难出错。
|
||||
func ListModelCatalog(ctx context.Context, agentName string) ([]CatalogModel, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT c.provider, c.model, c.display_name,
|
||||
CASE WHEN a.model IS NULL THEN 0 ELSE 1 END AS allowed,
|
||||
COALESCE(a.rank, 0)
|
||||
FROM agent_model_catalog c
|
||||
LEFT JOIN agent_allowed_models a
|
||||
ON a.agent_name = c.agent_name
|
||||
AND a.provider = c.provider
|
||||
AND a.model = c.model
|
||||
WHERE c.agent_name = $1
|
||||
ORDER BY c.provider, c.model
|
||||
`, agentName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []CatalogModel{}
|
||||
for rows.Next() {
|
||||
var m CatalogModel
|
||||
var allowed int
|
||||
if err := rows.Scan(&m.Provider, &m.Model, &m.DisplayName, &allowed, &m.Rank); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
m.Allowed = allowed == 1
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListAllowedModels 按 rank 返回该 Agent 在邮件场景下可用的模型。
|
||||
//
|
||||
// **不与目录做 JOIN**：目录是平台上次注册时的快照，插件重启前可能已经过期。
|
||||
// 真正能不能调通只有插件试过才知道 —— 这也正是插件要按顺序降级的原因。
|
||||
// 在这里用目录过滤，只会把「目录暂时没上报但其实可用」的模型挡掉。
|
||||
func ListAllowedModels(ctx context.Context, agentName string) ([]ModelRef, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT provider, model FROM agent_allowed_models
|
||||
WHERE agent_name = $1
|
||||
ORDER BY rank ASC, provider ASC, model ASC
|
||||
`, agentName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []ModelRef{}
|
||||
for rows.Next() {
|
||||
var m ModelRef
|
||||
if err := rows.Scan(&m.Provider, &m.Model); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ListStaleAllowedModels 返回已选但已不在平台目录里的模型。
|
||||
//
|
||||
// 平台可能临时下线了某个模型（换了 provider 配置、上游故障），
|
||||
// 而管理员的选择是持久的。界面上不显示这些项的话，管理员会以为自己
|
||||
// 没选过它们 —— 而它们其实还在被插件尝试（ListAllowedModels 不与目录 JOIN）。
|
||||
func ListStaleAllowedModels(ctx context.Context, agentName string) ([]ModelRef, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT a.provider, a.model
|
||||
FROM agent_allowed_models a
|
||||
WHERE a.agent_name = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM agent_model_catalog c
|
||||
WHERE c.agent_name = a.agent_name
|
||||
AND c.provider = a.provider
|
||||
AND c.model = a.model
|
||||
)
|
||||
ORDER BY a.rank ASC
|
||||
`, agentName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []ModelRef{}
|
||||
for rows.Next() {
|
||||
var m ModelRef
|
||||
if err := rows.Scan(&m.Provider, &m.Model); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// SetAllowedModels 整表替换某 Agent 的邮件场景可用模型，入参顺序即优先级。
|
||||
//
|
||||
// 允许传空列表：那表示「不限定」——插件回退到平台自己的默认模型。
|
||||
// 这与「一个都不许用」不同，后者等于让 Agent 彻底哑掉，不该是一次误删的后果。
|
||||
func SetAllowedModels(ctx context.Context, agentName string, picks []ModelRef) error {
|
||||
agentName = strings.TrimSpace(agentName)
|
||||
if agentName == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := db.DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`DELETE FROM agent_allowed_models WHERE agent_name = $1`, agentName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
rank := 0
|
||||
seen := map[string]bool{}
|
||||
for _, m := range picks {
|
||||
p := strings.TrimSpace(m.Provider)
|
||||
id := strings.TrimSpace(m.Model)
|
||||
if p == "" || id == "" {
|
||||
continue
|
||||
}
|
||||
key := p + "/" + id
|
||||
if seen[key] {
|
||||
// 重复项直接跳过而不是报错：它对最终顺序没有影响，
|
||||
// 为一次无害的重复让整次保存失败只会让人以为配置没生效。
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO agent_allowed_models (agent_name, provider, model, rank)
|
||||
VALUES ($1, $2, $3, $4)`, agentName, p, id, rank); err != nil {
|
||||
return err
|
||||
}
|
||||
rank++
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
299
server/internal/repo/models_scope_test.go
Normal file
299
server/internal/repo/models_scope_test.go
Normal file
@ -0,0 +1,299 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 目录与选择分两张表，是为了让「已选」在模型从平台目录里消失后仍然留存。
|
||||
// 这个测试钉住那个行为 —— 合并成一张带 allowed 标记的表就会失败。
|
||||
func TestAllowedModelsSurviveCatalogChurn(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "llmsproxy", Model: "AUTO", DisplayName: "AUTO"},
|
||||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||||
}); err != nil {
|
||||
t.Fatalf("首次上报目录: %v", err)
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||||
{Provider: "llmsproxy", Model: "AUTO"},
|
||||
}); err != nil {
|
||||
t.Fatalf("保存选择: %v", err)
|
||||
}
|
||||
|
||||
// 平台侧 AUTO 临时下线，只上报另一个
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||||
}); err != nil {
|
||||
t.Fatalf("二次上报目录: %v", err)
|
||||
}
|
||||
|
||||
allowed, err := ListAllowedModels(ctx, "dsh")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAllowedModels: %v", err)
|
||||
}
|
||||
if len(allowed) != 1 || allowed[0].Model != "AUTO" {
|
||||
t.Fatalf("模型从目录消失后选择也被删了：%+v —— "+
|
||||
"两张表分开存的意义就在于此", allowed)
|
||||
}
|
||||
|
||||
// 它应当被标为 stale，界面上才能提示「已选但平台没上报」
|
||||
stale, err := ListStaleAllowedModels(ctx, "dsh")
|
||||
if err != nil {
|
||||
t.Fatalf("ListStaleAllowedModels: %v", err)
|
||||
}
|
||||
if len(stale) != 1 || stale[0].Model != "AUTO" {
|
||||
t.Errorf("应有 1 个 stale，实际 %+v", stale)
|
||||
}
|
||||
|
||||
// 模型回来后不该再是 stale，也不需要重新勾选
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "llmsproxy", Model: "AUTO"},
|
||||
{Provider: "llmsproxy", Model: "claude-sonnet-4-6"},
|
||||
}); err != nil {
|
||||
t.Fatalf("三次上报: %v", err)
|
||||
}
|
||||
stale2, _ := ListStaleAllowedModels(ctx, "dsh")
|
||||
if len(stale2) != 0 {
|
||||
t.Errorf("模型回来后不该再是 stale：%+v", stale2)
|
||||
}
|
||||
}
|
||||
|
||||
// 目录整表替换：平台下线的模型必须从配置页消失，
|
||||
// 否则管理员会勾选一个平台其实调不到的模型。
|
||||
func TestReplaceModelCatalogIsFullReplace(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{
|
||||
{Provider: "p", Model: "a"},
|
||||
{Provider: "p", Model: "b"},
|
||||
}); err != nil {
|
||||
t.Fatalf("首次: %v", err)
|
||||
}
|
||||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{
|
||||
{Provider: "p", Model: "a"},
|
||||
}); err != nil {
|
||||
t.Fatalf("二次: %v", err)
|
||||
}
|
||||
got, err := ListModelCatalog(ctx, "opencode")
|
||||
if err != nil {
|
||||
t.Fatalf("ListModelCatalog: %v", err)
|
||||
}
|
||||
if len(got) != 1 || got[0].Model != "a" {
|
||||
t.Fatalf("整表替换失效：%+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ListModelCatalog 要在同一次查询里标出「已选」与 rank，
|
||||
// 前端才能画出带勾选与顺序的清单。
|
||||
func TestListModelCatalogMarksAllowedAndRank(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "p", Model: "first", DisplayName: "第一"},
|
||||
{Provider: "p", Model: "second"},
|
||||
{Provider: "p", Model: "unpicked"},
|
||||
}); err != nil {
|
||||
t.Fatalf("上报: %v", err)
|
||||
}
|
||||
// 顺序即优先级：second 排前面
|
||||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||||
{Provider: "p", Model: "second"},
|
||||
{Provider: "p", Model: "first"},
|
||||
}); err != nil {
|
||||
t.Fatalf("保存: %v", err)
|
||||
}
|
||||
|
||||
got, err := ListModelCatalog(ctx, "dsh")
|
||||
if err != nil {
|
||||
t.Fatalf("ListModelCatalog: %v", err)
|
||||
}
|
||||
byModel := map[string]CatalogModel{}
|
||||
for _, m := range got {
|
||||
byModel[m.Model] = m
|
||||
}
|
||||
if !byModel["second"].Allowed || byModel["second"].Rank != 0 {
|
||||
t.Errorf("second 应为 rank 0 的已选项：%+v", byModel["second"])
|
||||
}
|
||||
if !byModel["first"].Allowed || byModel["first"].Rank != 1 {
|
||||
t.Errorf("first 应为 rank 1 的已选项：%+v", byModel["first"])
|
||||
}
|
||||
if byModel["unpicked"].Allowed {
|
||||
t.Error("unpicked 不该被标为已选")
|
||||
}
|
||||
if byModel["first"].DisplayName != "第一" {
|
||||
t.Errorf("display_name 未带出：%q", byModel["first"].DisplayName)
|
||||
}
|
||||
}
|
||||
|
||||
// 顺序就是插件的降级顺序，必须原样保存。
|
||||
func TestSetAllowedModelsPreservesOrder(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
want := []ModelRef{
|
||||
{Provider: "c", Model: "3"},
|
||||
{Provider: "a", Model: "1"},
|
||||
{Provider: "b", Model: "2"},
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "dsh", want); err != nil {
|
||||
t.Fatalf("SetAllowedModels: %v", err)
|
||||
}
|
||||
got, err := ListAllowedModels(ctx, "dsh")
|
||||
if err != nil {
|
||||
t.Fatalf("ListAllowedModels: %v", err)
|
||||
}
|
||||
if len(got) != len(want) {
|
||||
t.Fatalf("数量不符：%d vs %d", len(got), len(want))
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("第 %d 项顺序错：%+v，期望 %+v —— "+
|
||||
"顺序就是插件的降级顺序，不能按字典序重排", i, got[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 空列表表示「不限定」，是合法输入。
|
||||
// 报错会让「取消所有限定」变成一件做不到的事。
|
||||
func TestSetAllowedModelsAcceptsEmpty(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{{Provider: "p", Model: "m"}}); err != nil {
|
||||
t.Fatalf("先设一个: %v", err)
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{}); err != nil {
|
||||
t.Fatalf("清空应当合法: %v", err)
|
||||
}
|
||||
got, _ := ListAllowedModels(ctx, "dsh")
|
||||
if len(got) != 0 {
|
||||
t.Errorf("清空后应为空，实际 %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 重复项跳过而不报错：它对最终顺序没有影响，
|
||||
// 为一次无害的重复让整次保存失败只会让人以为配置没生效。
|
||||
func TestSetAllowedModelsSkipsDuplicates(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
err := SetAllowedModels(ctx, "dsh", []ModelRef{
|
||||
{Provider: "p", Model: "m"},
|
||||
{Provider: "p", Model: "m"},
|
||||
{Provider: "p", Model: "other"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("重复项不该报错: %v", err)
|
||||
}
|
||||
got, _ := ListAllowedModels(ctx, "dsh")
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("应保留 2 项，实际 %+v", got)
|
||||
}
|
||||
// rank 要连续：跳过重复项后不该在序号上留空洞
|
||||
if got[0].Model != "m" || got[1].Model != "other" {
|
||||
t.Errorf("顺序错：%+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 字段不全的项跳过：半条记录在配置页上是一个点不动的空复选框。
|
||||
func TestModelCatalogSkipsIncomplete(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "", Model: "m"},
|
||||
{Provider: "p", Model: ""},
|
||||
{Provider: " ", Model: " "},
|
||||
{Provider: "p", Model: "ok"},
|
||||
}); err != nil {
|
||||
t.Fatalf("上报: %v", err)
|
||||
}
|
||||
got, _ := ListModelCatalog(ctx, "dsh")
|
||||
if len(got) != 1 || got[0].Model != "ok" {
|
||||
t.Fatalf("应只留 1 项：%+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 目录里重复的 provider/model 不该让整次事务失败（主键冲突）。
|
||||
func TestReplaceModelCatalogDedupes(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{
|
||||
{Provider: "p", Model: "m", DisplayName: "第一次"},
|
||||
{Provider: "p", Model: "m", DisplayName: "第二次"},
|
||||
}); err != nil {
|
||||
t.Fatalf("重复不该报错: %v", err)
|
||||
}
|
||||
got, _ := ListModelCatalog(ctx, "dsh")
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("应去重到 1 项：%+v", got)
|
||||
}
|
||||
if got[0].DisplayName != "第一次" {
|
||||
t.Errorf("应保留第一条：%q", got[0].DisplayName)
|
||||
}
|
||||
}
|
||||
|
||||
// 各 Agent 的目录与选择互不影响。
|
||||
func TestModelScopeIsolatedPerAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", []CatalogModel{{Provider: "p", Model: "dsh-only"}}); err != nil {
|
||||
t.Fatalf("dsh 上报: %v", err)
|
||||
}
|
||||
if err := ReplaceModelCatalog(ctx, "opencode", []CatalogModel{{Provider: "p", Model: "oc-only"}}); err != nil {
|
||||
t.Fatalf("opencode 上报: %v", err)
|
||||
}
|
||||
if err := SetAllowedModels(ctx, "dsh", []ModelRef{{Provider: "p", Model: "dsh-only"}}); err != nil {
|
||||
t.Fatalf("dsh 选择: %v", err)
|
||||
}
|
||||
|
||||
ocCatalog, _ := ListModelCatalog(ctx, "opencode")
|
||||
if len(ocCatalog) != 1 || ocCatalog[0].Model != "oc-only" {
|
||||
t.Fatalf("opencode 的目录被污染：%+v", ocCatalog)
|
||||
}
|
||||
if ocCatalog[0].Allowed {
|
||||
t.Error("dsh 的选择串到 opencode 上了")
|
||||
}
|
||||
ocAllowed, _ := ListAllowedModels(ctx, "opencode")
|
||||
if len(ocAllowed) != 0 {
|
||||
t.Errorf("opencode 不该有已选项：%+v", ocAllowed)
|
||||
}
|
||||
}
|
||||
|
||||
// 上报数量超上限时截断而不报错：平台把上游几千个模型全列出来是它的自由，
|
||||
// 但配置页上几千个复选框对人没有用。
|
||||
func TestReplaceModelCatalogCaps(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
many := make([]CatalogModel, maxCatalogModels+50)
|
||||
for i := range many {
|
||||
many[i] = CatalogModel{Provider: "p", Model: string(rune('a'+i%26)) + itoaTest(i)}
|
||||
}
|
||||
if err := ReplaceModelCatalog(ctx, "dsh", many); err != nil {
|
||||
t.Fatalf("超量上报不该报错: %v", err)
|
||||
}
|
||||
got, _ := ListModelCatalog(ctx, "dsh")
|
||||
if len(got) != maxCatalogModels {
|
||||
t.Errorf("应截断到 %d，实际 %d", maxCatalogModels, len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func itoaTest(n int) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
var b []byte
|
||||
for n > 0 {
|
||||
b = append([]byte{byte('0' + n%10)}, b...)
|
||||
n /= 10
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
119
server/internal/repo/nearesthuman_test.go
Normal file
119
server/internal/repo/nearesthuman_test.go
Normal file
@ -0,0 +1,119 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// 权限询问的决策者必须是人：Agent 收不到 SendToUser，桥的 await 永不 resolve。
|
||||
// 生产事故：pi 把任务派给自己的另一条会话 → 那条会话要跑 bash → 权限邮件发给 "pi"
|
||||
// → pi 不是人类用户 → 整条会话永久阻塞。
|
||||
//
|
||||
// 修法是顺着会话的邮件链上溯找最近的人类 —— 派活的人才是该点头的人。
|
||||
func TestNearestHumanInThread(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
mustExec(t, ctx, `INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ('alice','Alice','x','user')`)
|
||||
for _, a := range []string{"opencode", "dsh", "pi"} {
|
||||
if err := CreateOrUpdateAgent(ctx, a, "s", "test", nil); err != nil {
|
||||
t.Fatalf("注册 %s: %v", a, err)
|
||||
}
|
||||
}
|
||||
|
||||
t.Run("沿链上溯找到派活的人", func(t *testing.T) {
|
||||
sid, err := CreateSession(ctx, nil, "opencode", "任务链", "/home")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
// alice → opencode → dsh，dsh 触发权限询问
|
||||
m1, err := CreateMail(ctx, sid, nil, "alice", "", "opencode", "", "任务", "请帮忙", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("封1: %v", err)
|
||||
}
|
||||
m2, err := CreateMail(ctx, sid, &m1, "opencode", "", "dsh", "", "转派", "你来看", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("封2: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, sid, &m2, "dsh", "", "opencode", "", "进展", "做了一半", nil); err != nil {
|
||||
t.Fatalf("封3: %v", err)
|
||||
}
|
||||
|
||||
for _, agent := range []string{"dsh", "opencode"} {
|
||||
got, err := NearestHumanInThread(ctx, sid, agent)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", agent, err)
|
||||
}
|
||||
if got != "alice" {
|
||||
t.Errorf("%s 触发权限时应路由到 alice，得到 %q", agent, got)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("全 Agent 链返回空串", func(t *testing.T) {
|
||||
// 没有人类参与的链条：调用方据此拒绝请求，而不是转给一个
|
||||
// 对上下文一无所知的管理员。
|
||||
sid, err := CreateSession(ctx, nil, "opencode", "纯 Agent", "/tmp")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
m1, err := CreateMail(ctx, sid, nil, "opencode", "", "dsh", "", "干活", "go", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("封1: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, sid, &m1, "dsh", "", "opencode", "", "好", "ok", nil); err != nil {
|
||||
t.Fatalf("封2: %v", err)
|
||||
}
|
||||
|
||||
got, err := NearestHumanInThread(ctx, sid, "dsh")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Errorf("链上没有人类时应返回空串，得到 %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("skipAgent 是自己时不会把自己当人", func(t *testing.T) {
|
||||
// pi 给自己的另一条会话派活正是事故场景：链上只有 pi 一个名字。
|
||||
sid, err := CreateSession(ctx, nil, "pi", "自派", "/home")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
if _, err := CreateMail(ctx, sid, nil, "pi", "", "pi", "", "拆任务", "自己干", nil); err != nil {
|
||||
t.Fatalf("封1: %v", err)
|
||||
}
|
||||
|
||||
got, err := NearestHumanInThread(ctx, sid, "pi")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Errorf("自派链上没有人类，应返回空串，得到 %q", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("空会话不报错", func(t *testing.T) {
|
||||
sid, err := CreateSession(ctx, nil, "opencode", "空", "/tmp")
|
||||
if err != nil {
|
||||
t.Fatalf("建会话: %v", err)
|
||||
}
|
||||
got, err := NearestHumanInThread(ctx, sid, "opencode")
|
||||
if err != nil {
|
||||
t.Fatalf("空会话应返回空串而非报错，得到 err=%v", err)
|
||||
}
|
||||
if got != "" {
|
||||
t.Errorf("空会话应返回空串，得到 %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func mustExec(t *testing.T, ctx context.Context, q string, args ...any) {
|
||||
t.Helper()
|
||||
if _, err := db.DB.ExecContext(ctx, q, args...); err != nil {
|
||||
t.Fatalf("exec %s: %v", q, err)
|
||||
}
|
||||
}
|
||||
135
server/internal/repo/participants.go
Normal file
135
server/internal/repo/participants.go
Normal file
@ -0,0 +1,135 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sort"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Participant 是一条会话里的一个参与方。
|
||||
//
|
||||
// Path 是该参与方**自己那个地址的 path 位**，不是别人的：一封主发给 dsh@/b、
|
||||
// 抄送给 opencode@/a 的邮件里，两人的工作目录不同，混用会让对方在别人的目录里
|
||||
// 开会话（生产上已发生过一次，见 PLUGIN-CONTRACT 9.3）。
|
||||
type Participant struct {
|
||||
Name string `json:"name"`
|
||||
Path string `json:"path"`
|
||||
// Roles 是该参与方在这条会话里出现过的全部身份，from / to / cc 的并集。
|
||||
// 用集合而非单值：同一个人常常既发过信也被抄送过，只留最后一个身份会让
|
||||
// 「谁是这件事的负责人」这个判断出错。
|
||||
Roles []string `json:"roles"`
|
||||
// MailCount 是该参与方作为发件人的邮件数。用来回答「谁还没回」——
|
||||
// 参与方列表里 from 计数为 0 的那个就是还没开口的人。
|
||||
MailCount int `json:"mail_count"`
|
||||
}
|
||||
|
||||
// SessionParticipants 列出会话的全部参与方及各自的地址素材。
|
||||
//
|
||||
// 为什么要逐封扫而不是看 sessions 表：**参与方是随往来增长的**。会话建立时
|
||||
// 只有发件人与收件人，一封抄送、一次转发都会带进新的人。sessions 表里只有
|
||||
// from_agent 一个名字，回答不了「这条线索上现在有谁」。
|
||||
//
|
||||
// 排序：按首次出现顺序（created_at）。这让主收件人稳定排在抄送方之前，
|
||||
// 模型据此判断「谁是负责人、谁是配合方」——按名字排序会丢掉这个信息。
|
||||
func SessionParticipants(ctx context.Context, sessionID uuid.UUID) ([]Participant, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT from_name, COALESCE(from_workspace,''),
|
||||
to_name, COALESCE(to_workspace,''),
|
||||
cc_list
|
||||
FROM mails
|
||||
WHERE session_id = $1
|
||||
ORDER BY created_at ASC, mail_id ASC
|
||||
`, sessionID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
type acc struct {
|
||||
p Participant
|
||||
roles map[string]bool
|
||||
order int
|
||||
}
|
||||
seen := map[string]*acc{}
|
||||
next := 0
|
||||
|
||||
// note 记录一次「某人以某身份出现」。
|
||||
//
|
||||
// path 只在**当前为空且新值非空**时补写：同一个人可能在不同邮件里带不同
|
||||
// path（先被抄送到 /a，后被主发到 /b）。保留首个非空值而不是最后一个，
|
||||
// 与排序口径一致（首次出现顺序），也避免一封转发把地址改指到别处。
|
||||
note := func(name, path, role string, isSender bool) {
|
||||
if name == "" {
|
||||
return
|
||||
}
|
||||
a, ok := seen[name]
|
||||
if !ok {
|
||||
a = &acc{
|
||||
p: Participant{Name: name, Path: path},
|
||||
roles: map[string]bool{},
|
||||
order: next,
|
||||
}
|
||||
next++
|
||||
seen[name] = a
|
||||
}
|
||||
if a.p.Path == "" && path != "" {
|
||||
a.p.Path = path
|
||||
}
|
||||
a.roles[role] = true
|
||||
if isSender {
|
||||
a.p.MailCount++
|
||||
}
|
||||
}
|
||||
|
||||
for rows.Next() {
|
||||
var fromName, fromWS, toName, toWS string
|
||||
var ccRaw []byte
|
||||
if err := rows.Scan(&fromName, &fromWS, &toName, &toWS, &ccRaw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// **发件人一侧不取 from_workspace 当 path。** Agent 回信时那一列存的是
|
||||
// Agent 名而不是路径（历史遗留，FindOrCreateDefaultSession 的注释里也提到
|
||||
// 同一个坑）。拿它拼地址会得到 `dsh@dsh.alias` 这种投不出去的东西。
|
||||
note(fromName, "", "from", true)
|
||||
note(toName, toWS, "to", false)
|
||||
|
||||
if len(ccRaw) > 0 {
|
||||
var cc []models.Address
|
||||
if json.Unmarshal(ccRaw, &cc) == nil {
|
||||
for _, c := range cc {
|
||||
note(c.Name, c.Path, "cc", false)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
out := make([]Participant, 0, len(seen))
|
||||
for _, a := range seen {
|
||||
a.p.Roles = sortedKeys(a.roles)
|
||||
out = append(out, a.p)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
return seen[out[i].Name].order < seen[out[j].Name].order
|
||||
})
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sortedKeys 给出稳定顺序的角色列表。
|
||||
// map 迭代顺序随机，不排序的话同一条会话每次返回的 roles 顺序都不同，
|
||||
// 插件侧做 diff 或缓存时会误判为「参与方变了」。
|
||||
func sortedKeys(m map[string]bool) []string {
|
||||
out := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
out = append(out, k)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
170
server/internal/repo/participants_test.go
Normal file
170
server/internal/repo/participants_test.go
Normal file
@ -0,0 +1,170 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 参与方列表要回答的是「这条线索上现在有谁、用什么地址找到他、谁还没开口」。
|
||||
// 三个问题里每一个都曾经答错过：
|
||||
// - 有谁：sessions 表只有 from_agent 一个名字，抄送方与转发引入的人都不在里面
|
||||
// - 什么地址：拿 from_workspace 当 path 会拼出 dsh@dsh.alias 这种投不出去的东西
|
||||
// - 谁还没回：只留最后一个身份的话，既发过信又被抄送过的人会被算成纯配合方
|
||||
|
||||
func TestSessionParticipantsIncludesCC(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "抄收联调", "/home/program/llmsproxy")
|
||||
// 复刻线上那封：admin 主发 dsh，抄送 opencode
|
||||
mustMail(t, sid, "admin", "", "dsh", "/home/program/llmsproxy",
|
||||
[]models.Address{{Name: "opencode", Path: "/home", Session: "new", Raw: "opencode@/home.new"}})
|
||||
|
||||
parts, err := SessionParticipants(ctx, sid)
|
||||
if err != nil {
|
||||
t.Fatalf("列参与方: %v", err)
|
||||
}
|
||||
|
||||
byName := map[string]Participant{}
|
||||
for _, p := range parts {
|
||||
byName[p.Name] = p
|
||||
}
|
||||
for _, want := range []string{"admin", "dsh", "opencode"} {
|
||||
if _, ok := byName[want]; !ok {
|
||||
t.Errorf("参与方缺 %s，实得 %+v", want, parts)
|
||||
}
|
||||
}
|
||||
// 抄送方的 path 必须是它自己那个地址的 path 位，不是主收件人的。
|
||||
// 用错的后果：对方在别人的工作目录里开会话。
|
||||
if got := byName["opencode"].Path; got != "/home" {
|
||||
t.Errorf("opencode 的 path = %q，应为 /home（它自己地址的 path 位）", got)
|
||||
}
|
||||
if got := byName["dsh"].Path; got != "/home/program/llmsproxy" {
|
||||
t.Errorf("dsh 的 path = %q，应为 /home/program/llmsproxy", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsSenderPathStaysEmpty(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "回信", "/tmp/ws")
|
||||
// Agent 回信时 from_workspace 存的是 Agent 名而非路径（历史遗留）。
|
||||
// 若把它当 path，地址会拼成 dsh@dsh.alias —— 投不出去。
|
||||
mustMail(t, sid, "dsh", "dsh", "admin", "", nil)
|
||||
|
||||
parts, _ := SessionParticipants(ctx, sid)
|
||||
for _, p := range parts {
|
||||
if p.Name == "dsh" && p.Path == "dsh" {
|
||||
t.Fatal("发件人的 path 取了 from_workspace（那列存的是 Agent 名），会拼出无效地址")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsMergesRoles(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "往返", "/tmp/ws")
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/ws", nil) // admin=from, dsh=to
|
||||
mustMail(t, sid, "dsh", "dsh", "admin", "", nil) // dsh=from, admin=to
|
||||
|
||||
parts, _ := SessionParticipants(ctx, sid)
|
||||
for _, p := range parts {
|
||||
if len(p.Roles) != 2 {
|
||||
t.Errorf("%s 的 roles = %v，双方都该同时有 from 与 to", p.Name, p.Roles)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsCountsOnlySends(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "谁还没回", "/tmp/ws")
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/ws",
|
||||
[]models.Address{{Name: "opencode", Path: "/tmp/ws", Raw: "opencode@/tmp/ws"}})
|
||||
mustMail(t, sid, "dsh", "dsh", "admin", "", nil)
|
||||
|
||||
parts, _ := SessionParticipants(ctx, sid)
|
||||
got := map[string]int{}
|
||||
for _, p := range parts {
|
||||
got[p.Name] = p.MailCount
|
||||
}
|
||||
// MailCount 只数「作为发件人」的邮件：抄送方 opencode 一封都没发，
|
||||
// 计数为 0 正是「还没开口的人」这个判断的依据。
|
||||
if got["opencode"] != 0 {
|
||||
t.Errorf("opencode 只被抄送未发信，MailCount 应为 0，实为 %d", got["opencode"])
|
||||
}
|
||||
if got["admin"] != 1 || got["dsh"] != 1 {
|
||||
t.Errorf("admin/dsh 各发过一封，实为 %d/%d", got["admin"], got["dsh"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsKeepsFirstSeenOrder(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "顺序", "/tmp/ws")
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/ws",
|
||||
[]models.Address{{Name: "opencode", Path: "/home", Raw: "opencode@/home"}})
|
||||
|
||||
parts, _ := SessionParticipants(ctx, sid)
|
||||
// 按首次出现排序，让主收件人稳定排在抄送方之前 ——
|
||||
// 模型据此判断谁是负责人、谁是配合方；按名字排序会丢掉这个信息。
|
||||
want := []string{"admin", "dsh", "opencode"}
|
||||
if len(parts) != len(want) {
|
||||
t.Fatalf("参与方数量 %d，期望 %d: %+v", len(parts), len(want), parts)
|
||||
}
|
||||
for i, w := range want {
|
||||
if parts[i].Name != w {
|
||||
t.Errorf("第 %d 位是 %s，期望 %s", i, parts[i].Name, w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsPrefersFirstNonEmptyPath(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "改指", "/tmp/a")
|
||||
// 同一个人先被抄送到 /home，后被主发到 /tmp/b。
|
||||
// 保留首个非空值，与排序口径一致，也避免一封转发把地址改指到别处。
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/a",
|
||||
[]models.Address{{Name: "opencode", Path: "/home", Raw: "opencode@/home"}})
|
||||
mustMail(t, sid, "admin", "", "opencode", "/tmp/b", nil)
|
||||
|
||||
parts, _ := SessionParticipants(ctx, sid)
|
||||
for _, p := range parts {
|
||||
if p.Name == "opencode" && p.Path != "/home" {
|
||||
t.Fatalf("opencode 的 path = %q，应保持首次出现的 /home", p.Path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionParticipantsEmptySession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 会话刚建、还没有邮件。返回空列表而不是报错：
|
||||
// 调用方拿到空表能正常渲染「暂无参与方」，拿到 error 只能整个失败。
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "空会话", "/tmp/ws")
|
||||
parts, err := SessionParticipants(ctx, sid)
|
||||
if err != nil {
|
||||
t.Fatalf("空会话应正常返回: %v", err)
|
||||
}
|
||||
if len(parts) != 0 {
|
||||
t.Fatalf("空会话应无参与方，实得 %+v", parts)
|
||||
}
|
||||
}
|
||||
|
||||
func mustMail(t *testing.T, sid uuid.UUID, from, fromWS, to, toWS string, cc []models.Address) {
|
||||
t.Helper()
|
||||
if _, err := CreateMail(context.Background(), sid, nil,
|
||||
from, fromWS, to, toWS, "主题", "正文", cc); err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
}
|
||||
148
server/internal/repo/permission_mode.go
Normal file
148
server/internal/repo/permission_mode.go
Normal file
@ -0,0 +1,148 @@
|
||||
package repo
|
||||
|
||||
// 会话级权限档位的读写。
|
||||
//
|
||||
// ## 为什么档位挂在会话上而不是每封邮件上
|
||||
//
|
||||
// 与配额同一个理由（见 quota.go 的注释）：档位是**任务**的属性。
|
||||
// 「这件事只许你看不许你动」描述的是任务性质，不是某一封信的性质。
|
||||
//
|
||||
// 如果续谈的邮件也能带档位，每封新信都会悄悄改掉对方正在遵守的规则 ——
|
||||
// 而 plan 档的会话里模型已经被告知「只许看」，第二封信把它改成 full，
|
||||
// 是在一段已有上下文里换规则。人不一定意识到自己改了。
|
||||
//
|
||||
// 所以：**新建会话时设，续谈时忽略该字段，在对话页里显式编辑。**
|
||||
//
|
||||
// ## 为什么 Agent 不能自己指定档位
|
||||
//
|
||||
// 否则 Agent 发一封 mode=full 的信就给自己提权了。Agent 派活时子会话的档位
|
||||
// 由 InheritedMode 从父会话推导，且**只能同档或更严**（models.ModeAtMost）。
|
||||
// 这保证 plan 档的任务派不出 full 档的子任务 —— 与 hop_limit 一个形状。
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// SessionPermission 是一条会话的档位与实际强制力。
|
||||
//
|
||||
// 两个字段必须一起返回：档位是「要求什么」，强制力是「平台实际做到了什么」。
|
||||
// 只给前者会让人以为 plan 档管住了 homeagent（它的核心没有工具调用拦截点）。
|
||||
type SessionPermission struct {
|
||||
Mode string `json:"permission_mode"`
|
||||
Enforcement string `json:"permission_enforcement"`
|
||||
}
|
||||
|
||||
// GetSessionPermission 读一条会话的档位与强制力。
|
||||
//
|
||||
// 读出来的值一律过 Normalize：库里可能有历史脏数据（手工改库、旧版本写入），
|
||||
// 而调用方拿到一个认不出的档位时的行为无法预期。归一化在这里做一次，
|
||||
// 后续所有判断就都能假定值是合法的。
|
||||
func GetSessionPermission(ctx context.Context, id uuid.UUID) (SessionPermission, error) {
|
||||
var p SessionPermission
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(NULLIF(permission_mode, ''), 'workspace'),
|
||||
COALESCE(NULLIF(permission_enforcement, ''), 'advisory')
|
||||
FROM sessions WHERE session_id = $1`, id).Scan(&p.Mode, &p.Enforcement)
|
||||
if err != nil {
|
||||
return SessionPermission{}, err
|
||||
}
|
||||
p.Mode = models.NormalizePermissionMode(p.Mode)
|
||||
p.Enforcement = models.NormalizeEnforcement(p.Enforcement)
|
||||
return p, nil
|
||||
}
|
||||
|
||||
// SessionPermissionMode 只取档位，读不到时回落默认档。
|
||||
//
|
||||
// 供投递路径使用：那里拿不到档位也得继续走（不能因为查询失败就拒收邮件），
|
||||
// 但回落必须是默认档而不是 full —— 查询失败不该换来更大的权限。
|
||||
func SessionPermissionMode(ctx context.Context, id uuid.UUID) string {
|
||||
p, err := GetSessionPermission(ctx, id)
|
||||
if err != nil {
|
||||
return models.DefaultPermissionMode
|
||||
}
|
||||
return p.Mode
|
||||
}
|
||||
|
||||
// SetSessionPermissionMode 设置会话档位。
|
||||
//
|
||||
// 非法档位一律收敛成默认档而不是报错：这个函数的调用方包括人在界面上操作，
|
||||
// 而界面传来一个拼错的值时，静默用默认档比让整次操作失败更合理 ——
|
||||
// 默认档本身是安全的。
|
||||
func SetSessionPermissionMode(ctx context.Context, id uuid.UUID, mode string) (SessionPermission, error) {
|
||||
m := models.NormalizePermissionMode(mode)
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET permission_mode = $2, updated_at = NOW() WHERE session_id = $1`,
|
||||
id, m)
|
||||
if err != nil {
|
||||
return SessionPermission{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return SessionPermission{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
return GetSessionPermission(ctx, id)
|
||||
}
|
||||
|
||||
// SetSessionEnforcement 记录接收平台实际做到的强制力。
|
||||
//
|
||||
// 由投递路径在建会话时按收件 Agent 的自报能力写入 —— 它是**事实快照**
|
||||
// 而不是配置：插件升级后能力会变，但已结束的会话不该被改写成「其实当时
|
||||
// 是被强制的」。所以不跟着 agents.mode_enforcement 走，而是建会话时定死。
|
||||
func SetSessionEnforcement(ctx context.Context, id uuid.UUID, enforcement string) error {
|
||||
e := models.NormalizeEnforcement(enforcement)
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET permission_enforcement = $2 WHERE session_id = $1`, id, e)
|
||||
return err
|
||||
}
|
||||
|
||||
// AgentModeEnforcement 取某个 Agent 自报的档位强制力。
|
||||
//
|
||||
// Agent 不存在或没自报过时返回 advisory：不能替一个没说过话的插件宣称
|
||||
// 「档位在它那里是被强制的」。保守方向是承认做不到。
|
||||
func AgentModeEnforcement(ctx context.Context, agentName string) string {
|
||||
var e string
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(NULLIF(mode_enforcement, ''), 'advisory')
|
||||
FROM agents WHERE agent_name = $1`, agentName).Scan(&e)
|
||||
if err != nil {
|
||||
return models.EnforcementAdvisory
|
||||
}
|
||||
return models.NormalizeEnforcement(e)
|
||||
}
|
||||
|
||||
// SetAgentModeEnforcement 落库 Agent 心跳自报的档位强制力。
|
||||
//
|
||||
// 走心跳而不是注册：注册只在插件启动时发生一次，而能力可能因为配置变化
|
||||
// （比如 DSH 的 sandbox 被换成 danger-full-access）而改变。与模型目录上报
|
||||
// 同一条通道 —— I-1：平台自己说的才算。
|
||||
func SetAgentModeEnforcement(ctx context.Context, agentName, enforcement string) error {
|
||||
e := models.NormalizeEnforcement(enforcement)
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET mode_enforcement = $2 WHERE agent_name = $1`, agentName, e)
|
||||
return err
|
||||
}
|
||||
|
||||
// InheritedMode 推导子会话应当继承的档位。
|
||||
//
|
||||
// parentSessionID 为 nil（人直接发起、或没有父会话可依据）时返回 requested
|
||||
// 归一化后的值；有父会话时取**父档位与请求档位里更严的那一个**。
|
||||
//
|
||||
// 为什么必须取更严：Agent 派活时若能给子会话一个更宽松的档位，plan 档的
|
||||
// 任务就能通过「派给自己一条 full 档子会话」来提权，档位形同虚设。
|
||||
// 这与 hop_limit 防自激的形状一样 —— 约束必须沿着链条传递下去。
|
||||
func InheritedMode(ctx context.Context, parentSessionID *uuid.UUID, requested string) string {
|
||||
req := models.NormalizePermissionMode(requested)
|
||||
if parentSessionID == nil {
|
||||
return req
|
||||
}
|
||||
parent, err := GetSessionPermission(ctx, *parentSessionID)
|
||||
if err != nil {
|
||||
// 父会话查不到时按默认档与请求档取更严 —— 不能因为查询失败而放宽。
|
||||
return models.ModeAtMost(models.DefaultPermissionMode, req)
|
||||
}
|
||||
return models.ModeAtMost(parent.Mode, req)
|
||||
}
|
||||
369
server/internal/repo/permission_mode_test.go
Normal file
369
server/internal/repo/permission_mode_test.go
Normal file
@ -0,0 +1,369 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ─── InheritedMode：继承、收紧、不存在的父会话 ───
|
||||
|
||||
// parent nil → 返回 requested 的规范化值
|
||||
func TestInheritedMode_NilParent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
got := InheritedMode(ctx, nil, "plan")
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("nil parent + plan: got %q, want plan", got)
|
||||
}
|
||||
|
||||
got = InheritedMode(ctx, nil, "workspace")
|
||||
if got != models.ModeWorkspace {
|
||||
t.Errorf("nil parent + workspace: got %q, want workspace", got)
|
||||
}
|
||||
|
||||
got = InheritedMode(ctx, nil, "full")
|
||||
if got != models.ModeFull {
|
||||
t.Errorf("nil parent + full: got %q, want full", got)
|
||||
}
|
||||
}
|
||||
|
||||
// parent plan → 子会话只能 plan（不能提权）
|
||||
func TestInheritedMode_ParentPlan_Tightens(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id := createTestSession(t, ctx, "pi", "/ws")
|
||||
perm, err := SetSessionPermissionMode(ctx, id, models.ModePlan)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if perm.Mode != models.ModePlan {
|
||||
t.Fatal("expected plan")
|
||||
}
|
||||
|
||||
// 请求 workspace（更宽松）→ 应被收紧为 plan
|
||||
got := InheritedMode(ctx, &id, models.ModeWorkspace)
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("plan parent + workspace request: got %q, want plan", got)
|
||||
}
|
||||
|
||||
// 请求 full → 同样收紧
|
||||
got = InheritedMode(ctx, &id, models.ModeFull)
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("plan parent + full request: got %q, want plan", got)
|
||||
}
|
||||
|
||||
// 请求 plan → 保持 plan
|
||||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("plan parent + plan request: got %q, want plan", got)
|
||||
}
|
||||
}
|
||||
|
||||
// parent workspace → 子会话 workspace 或更严
|
||||
func TestInheritedMode_ParentWorkspace(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, err := SetSessionPermissionMode(ctx, id, models.ModeWorkspace)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// 请求 full → 收紧为 workspace（子不能比父更松）
|
||||
got := InheritedMode(ctx, &id, models.ModeFull)
|
||||
if got != models.ModeWorkspace {
|
||||
t.Errorf("workspace parent + full: got %q, want workspace", got)
|
||||
}
|
||||
|
||||
// 请求 plan → 保留 plan（比父更严，允许）
|
||||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("workspace parent + plan: got %q, want plan", got)
|
||||
}
|
||||
}
|
||||
|
||||
// parent full → 子会话可请求任意档位
|
||||
func TestInheritedMode_ParentFull(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, err := SetSessionPermissionMode(ctx, id, models.ModeFull)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got := InheritedMode(ctx, &id, models.ModeWorkspace)
|
||||
if got != models.ModeWorkspace {
|
||||
t.Errorf("full parent + workspace: got %q, want workspace", got)
|
||||
}
|
||||
got = InheritedMode(ctx, &id, models.ModePlan)
|
||||
if got != models.ModePlan {
|
||||
t.Errorf("full parent + plan: got %q, want plan", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 脏值 fallback：非法的 requested 值在 InheritedMode 里被规范化为默认档
|
||||
func TestInheritedMode_InvalidRequested(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
got := InheritedMode(ctx, nil, "elephant")
|
||||
if got != models.DefaultPermissionMode {
|
||||
t.Errorf("invalid requested: got %q, want %q", got, models.DefaultPermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
// parent 不存在时（查不到）回落：ModeAtMost(default, req)
|
||||
func TestInheritedMode_InvalidParent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
fakeID := uuid.New()
|
||||
got := InheritedMode(ctx, &fakeID, "full")
|
||||
want := models.ModeAtMost(models.DefaultPermissionMode, "full")
|
||||
if got != want {
|
||||
t.Errorf("invalid parent + full: got %q, want %q (modeAtMost(default, full))", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── SetSessionPermissionMode roundtrip + dirty value normalization ───
|
||||
|
||||
func TestSetSessionPermissionMode_Roundtrip(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "dsh", 20)
|
||||
|
||||
id := createTestSession(t, ctx, "dsh", "/ws")
|
||||
perm, err := SetSessionPermissionMode(ctx, id, models.ModeFull)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if perm.Mode != models.ModeFull || perm.Enforcement != "advisory" {
|
||||
t.Errorf("full mode: got mode=%q enforcement=%q", perm.Mode, perm.Enforcement)
|
||||
}
|
||||
|
||||
// 读出来一致
|
||||
got := SessionPermissionMode(ctx, id)
|
||||
if got != models.ModeFull {
|
||||
t.Errorf("read back: got %q, want full", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetSessionPermissionMode_DirtyValue_FailClosed(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id := createTestSession(t, ctx, "pi", "/ws")
|
||||
perm, err := SetSessionPermissionMode(ctx, id, "BOGUS")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if perm.Mode != models.DefaultPermissionMode {
|
||||
t.Errorf("dirty value: got %q, want %q (fail-closed to default)", perm.Mode, models.DefaultPermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 脏值归一化测试（覆盖 NormalizePermissionMode 本身） ───
|
||||
|
||||
func TestNormalizePermissionMode_Inputs(t *testing.T) {
|
||||
tests := []struct {
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{"plan", "plan"},
|
||||
{"workspace", "workspace"},
|
||||
{"full", "full"},
|
||||
// 大小写/空白不归一：NormalizePermissionMode 只接受精确匹配的合法档位，
|
||||
// 其余一律 fail-closed 到默认档（workspace）—— 不 trim 不 lowercase，
|
||||
// 避免「我以为给了 plan 实际拿到别的」这种隐式转换造成的安全错觉。
|
||||
{"Plan", models.DefaultPermissionMode},
|
||||
{" PLAN ", models.DefaultPermissionMode},
|
||||
{"", models.DefaultPermissionMode},
|
||||
{"bogus", models.DefaultPermissionMode},
|
||||
{"F ULL", models.DefaultPermissionMode},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
got := models.NormalizePermissionMode(tt.input)
|
||||
if got != tt.want {
|
||||
t.Errorf("NormalizePermissionMode(%q) = %q, want %q", tt.input, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 负向对照：plan 档不能提权 ───
|
||||
|
||||
func TestInheritedMode_PlanCannotEscalate(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
parent := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, _ = SetSessionPermissionMode(ctx, parent, models.ModePlan)
|
||||
|
||||
child := InheritedMode(ctx, &parent, models.ModeFull)
|
||||
if child != models.ModePlan {
|
||||
t.Errorf("SECURITY FAIL: plan session escalated to %q via InheritedMode", child)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 三层继承链 ───
|
||||
|
||||
func TestInheritedMode_ThreeLevelChain(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
root := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, _ = SetSessionPermissionMode(ctx, root, models.ModeFull)
|
||||
|
||||
child := InheritedMode(ctx, &root, models.ModeWorkspace) // workspace < full → workspace
|
||||
childID := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, _ = SetSessionPermissionMode(ctx, childID, child)
|
||||
|
||||
grandchild := InheritedMode(ctx, &childID, models.ModeFull) // full vs workspace → workspace
|
||||
if grandchild != models.ModeWorkspace {
|
||||
t.Errorf("grandchild: got %q, want workspace", grandchild)
|
||||
}
|
||||
|
||||
// plan → workspace → plan chain
|
||||
planChild := InheritedMode(ctx, &root, models.ModePlan) // plan < full → plan
|
||||
planChildID := createTestSession(t, ctx, "pi", "/ws")
|
||||
_, _ = SetSessionPermissionMode(ctx, planChildID, planChild)
|
||||
|
||||
grandchild2 := InheritedMode(ctx, &planChildID, models.ModeFull) // full vs plan → plan
|
||||
if grandchild2 != models.ModePlan {
|
||||
t.Errorf("plan chain grandchild: got %q, want plan", grandchild2)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── calendar_events permission_mode roundtrip ───
|
||||
|
||||
func TestCalendarEventPermissionMode_Roundtrip(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
// CreateCalendarEvent 应规范化档位
|
||||
e := &models.CalendarEvent{
|
||||
Title: "测试日程",
|
||||
AgentName: "pi",
|
||||
ToAddress: "pi@/home/program/agentmail",
|
||||
PermissionMode: "full",
|
||||
Status: "active",
|
||||
CreatedBy: "jianf",
|
||||
}
|
||||
created, err := CreateCalendarEvent(ctx, e)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created.PermissionMode != "full" {
|
||||
t.Errorf("created event permission_mode: got %q, want full", created.PermissionMode)
|
||||
}
|
||||
|
||||
// 读回来一致
|
||||
got, err := GetCalendarEvent(ctx, created.EventID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.PermissionMode != "full" {
|
||||
t.Errorf("read back: got %q, want full", got.PermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalendarEventPermissionMode_DirtyValue_Normalized(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
e := &models.CalendarEvent{
|
||||
Title: "脏值日程",
|
||||
AgentName: "pi",
|
||||
ToAddress: "pi@/home/program/agentmail",
|
||||
PermissionMode: "INVALID",
|
||||
Status: "active",
|
||||
CreatedBy: "jianf",
|
||||
}
|
||||
created, err := CreateCalendarEvent(ctx, e)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if created.PermissionMode != models.DefaultPermissionMode {
|
||||
t.Errorf("dirty value: got %q, want %q", created.PermissionMode, models.DefaultPermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCalendarEventPermissionMode_UpdateRoundtrip(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
e := &models.CalendarEvent{
|
||||
Title: "更新日程",
|
||||
AgentName: "pi",
|
||||
ToAddress: "pi@/home/program/agentmail",
|
||||
PermissionMode: "workspace",
|
||||
Status: "active",
|
||||
CreatedBy: "jianf",
|
||||
}
|
||||
created, err := CreateCalendarEvent(ctx, e)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
created.PermissionMode = "plan"
|
||||
if err := UpdateCalendarEvent(ctx, created.EventID, created); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := GetCalendarEvent(ctx, created.EventID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.PermissionMode != "plan" {
|
||||
t.Errorf("after update: got %q, want plan", got.PermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── adopt 接管时会话档位必须写入 ───
|
||||
|
||||
func TestAdoptPlatformSession_WritesDefaultMode(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "pi", 20)
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "plat-123", "my-proj", "/ws", "接管测试")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 接管会话应显式写入默认档位（不是靠 DB 默认值）
|
||||
mode := SessionPermissionMode(ctx, id)
|
||||
if mode != models.DefaultPermissionMode {
|
||||
t.Errorf("adopt session mode: got %q, want %q", mode, models.DefaultPermissionMode)
|
||||
}
|
||||
}
|
||||
|
||||
// ─── helpers ───
|
||||
|
||||
func createTestSession(t *testing.T, ctx context.Context, agent, workspace string) uuid.UUID {
|
||||
t.Helper()
|
||||
id, err := CreateSession(ctx, nil, agent, "test subject", workspace)
|
||||
if err != nil {
|
||||
t.Fatalf("create session: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
// init 确保每个 test 函数执行前 DB 足够干净
|
||||
func init() {
|
||||
// 空 —— setupTestDB 在每个测试函数内调用
|
||||
}
|
||||
78
server/internal/repo/platform_owner_test.go
Normal file
78
server/internal/repo/platform_owner_test.go
Normal file
@ -0,0 +1,78 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// platform_session_id 必须只发给归属方 —— 生产上 pi 的会话 id 被推给了抄送方 dsh，
|
||||
// DSH 在自己磁盘上找不到那个文件，按 N-8 抛错，邮件静默消失。
|
||||
func TestPlatformSessionFor_ReturnsOwnerFromMirror(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedPlatformAgent(t, "pi")
|
||||
seedPlatformAgent(t, "dsh")
|
||||
|
||||
// pi 上报一条平台会话
|
||||
if err := ReplacePlatformSessions(ctx, "pi", []PlatformSession{
|
||||
{PlatformID: "pid-pi-1", Workspace: "/w", Slug: "项目定位", Title: "项目定位"},
|
||||
}); err != nil {
|
||||
t.Fatalf("ReplacePlatformSessions: %v", err)
|
||||
}
|
||||
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pid-pi-1", "项目定位", "/w", "项目定位")
|
||||
if err != nil {
|
||||
t.Fatalf("AdoptPlatformSession: %v", err)
|
||||
}
|
||||
|
||||
pid, owner := PlatformSessionFor(ctx, id)
|
||||
if pid != "pid-pi-1" {
|
||||
t.Errorf("platformID = %q, want pid-pi-1", pid)
|
||||
}
|
||||
if owner != "pi" {
|
||||
t.Errorf("owner = %q, want pi（镜像里 agent_name=pi）", owner)
|
||||
}
|
||||
}
|
||||
|
||||
// 未接管的普通会话不该返回任何 platform id。
|
||||
func TestPlatformSessionFor_PlainSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedAgent(t, "dsh", 20)
|
||||
id, err := CreateSession(ctx, nil, "dsh", "普通会话", "/w")
|
||||
if err != nil {
|
||||
t.Fatalf("CreateSession: %v", err)
|
||||
}
|
||||
pid, owner := PlatformSessionFor(ctx, id)
|
||||
if pid != "" || owner != "" {
|
||||
t.Errorf("got (%q,%q), want ('','')", pid, owner)
|
||||
}
|
||||
}
|
||||
|
||||
// 镜像那行被整表替换掉（平台侧删了会话）时退回 sessions.from_agent，
|
||||
// 而不是让 owner 变空 —— 变空会让归属方也收不到 platform_session_id。
|
||||
func TestPlatformSessionFor_MirrorGoneFallsBackToFromAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
seedPlatformAgent(t, "pi")
|
||||
if err := ReplacePlatformSessions(ctx, "pi", []PlatformSession{
|
||||
{PlatformID: "pid-pi-2", Workspace: "/w", Slug: "s2", Title: "t2"},
|
||||
}); err != nil {
|
||||
t.Fatalf("ReplacePlatformSessions: %v", err)
|
||||
}
|
||||
id, err := AdoptPlatformSession(ctx, "pi", "pid-pi-2", "s2", "/w", "t2")
|
||||
if err != nil {
|
||||
t.Fatalf("AdoptPlatformSession: %v", err)
|
||||
}
|
||||
// 平台侧删了这条会话 → 心跳整表替换成空
|
||||
if err := ReplacePlatformSessions(ctx, "pi", []PlatformSession{}); err != nil {
|
||||
t.Fatalf("ReplacePlatformSessions(empty): %v", err)
|
||||
}
|
||||
pid, owner := PlatformSessionFor(ctx, id)
|
||||
if pid != "pid-pi-2" {
|
||||
t.Errorf("platformID = %q, want pid-pi-2", pid)
|
||||
}
|
||||
if owner != "pi" {
|
||||
t.Errorf("owner = %q, want pi（退回 sessions.from_agent）", owner)
|
||||
}
|
||||
}
|
||||
385
server/internal/repo/platform_sessions.go
Normal file
385
server/internal/repo/platform_sessions.go
Normal file
@ -0,0 +1,385 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 平台会话镜像 ----------
|
||||
//
|
||||
// Agent 平台自己也在开会话：有些经由邮件驱动，有些是人直接在平台界面上开的。
|
||||
// 写信时想续谈某条会话，得先知道那个工作区下有哪些会话可续 —— 而 Gateway
|
||||
// 只看得见邮件驱动的那部分。
|
||||
//
|
||||
// **由插件在心跳里上报，Gateway 不反向拉取。**
|
||||
// 当前架构是单向的（Agent 持密钥主动连 Gateway，Gateway 从不外呼）；
|
||||
// 让 Gateway 去调平台接口需要它保存各平台的地址与凭证，那是另一套信任模型。
|
||||
// 代价是插件没运行时同步不了 —— 但插件没运行时邮件本来也投不进去。
|
||||
|
||||
// PlatformSession 是插件上报的一条平台侧会话。
|
||||
type PlatformSession struct {
|
||||
PlatformID string `json:"platform_id"`
|
||||
Workspace string `json:"workspace"`
|
||||
Slug string `json:"slug,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
MailDriven bool `json:"mail_driven"`
|
||||
UpdatedAt *time.Time `json:"updated_at,omitempty"`
|
||||
}
|
||||
|
||||
// maxPlatformSessions 限制单次上报的会话数。
|
||||
//
|
||||
// 一个长期运行的平台可以累积上千条会话，而候选列表上千项对人没有意义。
|
||||
// 插件按最近活跃排序后上报前 N 条即可。
|
||||
const maxPlatformSessions = 200
|
||||
|
||||
// ReplacePlatformSessions 整表替换某 Agent 的平台会话镜像。
|
||||
//
|
||||
// 整表替换而非增量合并：镜像是平台当前状态的快照。增量合并会让已经删掉的
|
||||
// 平台会话永远留在候选列表里，而那正是「选了却送不到」的来源
|
||||
// —— session 位是三态语义，指向一条不存在的会话会直接 404。
|
||||
func ReplacePlatformSessions(ctx context.Context, agentName string, list []PlatformSession) error {
|
||||
agentName = strings.TrimSpace(agentName)
|
||||
if agentName == "" {
|
||||
return nil
|
||||
}
|
||||
if len(list) > maxPlatformSessions {
|
||||
list = list[:maxPlatformSessions]
|
||||
}
|
||||
|
||||
tx, err := db.DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`DELETE FROM agent_platform_sessions WHERE agent_name = $1`, agentName); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
seen := map[string]bool{}
|
||||
for _, ps := range list {
|
||||
id := strings.TrimSpace(ps.PlatformID)
|
||||
if id == "" || seen[id] {
|
||||
continue
|
||||
}
|
||||
seen[id] = true
|
||||
driven := 0
|
||||
if ps.MailDriven {
|
||||
driven = 1
|
||||
}
|
||||
if _, err := tx.ExecContext(ctx, `
|
||||
INSERT INTO agent_platform_sessions
|
||||
(agent_name, platform_id, workspace, slug, title, mail_driven, updated_at, reported_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, NOW())
|
||||
`, agentName, id, strings.TrimSpace(ps.Workspace), strings.TrimSpace(ps.Slug),
|
||||
strings.TrimSpace(ps.Title), driven, ps.UpdatedAt); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// SessionCandidate 是「续谈某条会话」的一个候选项。
|
||||
type SessionCandidate struct {
|
||||
// Alias 是填进 session 位的值 —— 候选项的实际用途就是它
|
||||
Alias string `json:"alias"`
|
||||
// Title 给人看，用来分辨两条别名相似的会话在谈什么
|
||||
Title string `json:"title,omitempty"`
|
||||
// Source 说明这条候选从哪来：
|
||||
// mail 本侧邮件线索（可直接送达）
|
||||
// platform 平台侧会话镜像（本侧还没有对应线索）
|
||||
Source string `json:"source"`
|
||||
// Unread 仅 mail 来源有意义
|
||||
Unread int `json:"unread,omitempty"`
|
||||
}
|
||||
|
||||
// SuggestSessionCandidates 汇总某 name@path 下可续谈的会话。
|
||||
//
|
||||
// 两个来源合并：
|
||||
// 1. 本侧邮件线索（sessions.workspace 匹配，或历史数据里靠 mails 反推）
|
||||
// 2. 平台会话镜像里带 slug 的那些
|
||||
//
|
||||
// 本侧优先：邮件线索是「这个别名一定送得到」的保证，而镜像只是平台的说法。
|
||||
// 同名时保留本侧那条，并把镜像的标题补上去（镜像通常有更新的标题）。
|
||||
func SuggestSessionCandidates(ctx context.Context, forUser, peerName, path string) ([]SessionCandidate, error) {
|
||||
out := []SessionCandidate{}
|
||||
seen := map[string]int{} // alias -> out 下标
|
||||
|
||||
// ---- 来源 1：本侧邮件线索 ----
|
||||
//
|
||||
// sessions.workspace 是权威来源。它是新加的列，历史会话为空串，
|
||||
// 因此保留 mails 反推作为兜底：`s.workspace = $2 OR (s.workspace = '' AND <mails 反推>)`。
|
||||
// 反推只看 to_workspace —— Agent 回信时 from_workspace 存的是 Agent 名而非路径，
|
||||
// 拿它比路径永远匹配不上。
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT s.session_alias,
|
||||
COALESCE(s.subject, ''),
|
||||
COALESCE(s.platform_id, ''),
|
||||
(SELECT COUNT(*) FROM mails u
|
||||
WHERE u.session_id = s.session_id AND u.status = 'unread')
|
||||
FROM sessions s
|
||||
WHERE s.session_alias IS NOT NULL AND s.session_alias <> ''
|
||||
AND s.status <> 'archived'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM mails m
|
||||
WHERE m.session_id = s.session_id
|
||||
AND (m.to_name = $1 OR m.from_name = $1 OR `+db.CCHas("m.cc_list", 1)+`)
|
||||
)
|
||||
AND ($2 = ''
|
||||
OR s.workspace = $2
|
||||
OR (s.workspace = '' AND EXISTS (
|
||||
SELECT 1 FROM mails w
|
||||
WHERE w.session_id = s.session_id
|
||||
AND COALESCE(w.to_workspace,'') = $2
|
||||
)))
|
||||
AND ($3 = '' OR s.owner_user_id = (SELECT user_id FROM users WHERE username = $3)
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM mails mm
|
||||
WHERE mm.session_id = s.session_id
|
||||
AND (mm.from_name = $3 OR mm.to_name = $3
|
||||
OR `+db.CCHas("mm.cc_list", 3)+`)
|
||||
))
|
||||
ORDER BY s.updated_at DESC
|
||||
`, peerName, path, forUser)
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
for rows.Next() {
|
||||
var alias, title, pid string
|
||||
var unread int
|
||||
if err := rows.Scan(&alias, &title, &pid, &unread); err != nil {
|
||||
return out, err
|
||||
}
|
||||
if alias == "" {
|
||||
continue
|
||||
}
|
||||
seen[alias] = len(out)
|
||||
out = append(out, SessionCandidate{
|
||||
Alias: alias, Title: title, Source: "mail", Unread: unread,
|
||||
})
|
||||
if pid != "" {
|
||||
seen["pid:"+pid] = len(out) - 1
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return out, err
|
||||
}
|
||||
|
||||
// ---- 来源 2：平台会话镜像 ----
|
||||
prows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT slug, title, platform_id
|
||||
FROM agent_platform_sessions
|
||||
WHERE agent_name = $1
|
||||
AND slug <> ''
|
||||
AND ($2 = '' OR workspace = $2)
|
||||
-- 不用 NULLS LAST：它要 SQLite 3.30+，而驱动自带的版本不由我们控制。
|
||||
-- COALESCE 在两个方言里都成立，语义也更直接：没有 updated_at 就用上报时间。
|
||||
ORDER BY COALESCE(updated_at, reported_at) DESC
|
||||
`, peerName, path)
|
||||
if err != nil {
|
||||
// 镜像查不到不该让整个补全失败：本侧线索已经够用了
|
||||
return out, nil
|
||||
}
|
||||
defer prows.Close()
|
||||
|
||||
for prows.Next() {
|
||||
var slug, title, pid string
|
||||
if err := prows.Scan(&slug, &title, &pid); err != nil {
|
||||
break
|
||||
}
|
||||
if slug == "" {
|
||||
continue
|
||||
}
|
||||
// 已被接管的平台会话不再单独列：选它也会落进已有的那条本侧线索，
|
||||
// 但候选列表出现两次会让人以为有两条不同的会话（项目定位 x2 的场景）。
|
||||
if pid != "" {
|
||||
if _, dup := seen["pid:"+pid]; dup {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if i, ok := seen[slug]; ok {
|
||||
// 本侧已有同名线索：保留 mail 来源（它保证送得到），
|
||||
// 但补上镜像的标题 —— 平台侧标题通常比会话建立时的主题更贴切
|
||||
if out[i].Title == "" && title != "" {
|
||||
out[i].Title = title
|
||||
}
|
||||
continue
|
||||
}
|
||||
seen[slug] = len(out)
|
||||
out = append(out, SessionCandidate{Alias: slug, Title: title, Source: "platform"})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// SetSessionWorkspace 记下会话所属的工作目录。
|
||||
//
|
||||
// 只在为空时写入：会话的工作区在建立时就定下了，之后不该被一封发往
|
||||
// 别处的邮件改掉 —— 那会让这条会话在候选列表里凭空换一个工作区。
|
||||
func SetSessionWorkspace(ctx context.Context, sessionID interface{ String() string }, workspace string) error {
|
||||
ws := strings.TrimSpace(workspace)
|
||||
if ws == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET workspace = $1 WHERE session_id = $2 AND workspace = ''`,
|
||||
ws, sessionID.String())
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- 接管平台会话 ----------
|
||||
//
|
||||
// TUI 与邮箱是同一个 Agent 的**两个入口**，不是两套隔离的世界。
|
||||
// 人在平台界面上开的会话，应该也能被邮件投进去 —— 补全早就把它们列为候选，
|
||||
// 缺的只是投递侧这一跳。
|
||||
//
|
||||
// 「接管」= 在本侧建一条会话并把 platform_id 记上。之后：
|
||||
// - 这条会话在 sessions 表里有正式身份（可寻址、有预算、能归档）
|
||||
// - 插件收到投递事件时看到 platform_id，就去 resume 那条平台会话
|
||||
// 而不是新建一条
|
||||
//
|
||||
// 一条平台会话只能被接管一次：第二次投递复用第一次建的本侧会话，
|
||||
// 否则同一条 TUI 对话会在邮箱里裂成多条互不相干的线索。
|
||||
|
||||
// FindPlatformSession 按 (agent, slug, workspace) 找一条平台会话镜像。
|
||||
//
|
||||
// workspace 为空表示不限（地址省略 path 位时）。返回 platform_id 与它的
|
||||
// 真实 workspace —— 后者是权威的：**会话的 cwd 在它创建时就定了**，
|
||||
// 地址里的 path 位若与之不同，以会话为准。人是从候选列表里选的，
|
||||
// 他要的是「那条会话」而不是「那个目录」。
|
||||
func FindPlatformSession(ctx context.Context, agentName, slug, workspace string) (platformID, realWorkspace, title string, err error) {
|
||||
agentName = strings.TrimSpace(agentName)
|
||||
slug = strings.TrimSpace(slug)
|
||||
if agentName == "" || slug == "" {
|
||||
return "", "", "", ErrSessionNotFound
|
||||
}
|
||||
ws := strings.TrimSpace(workspace)
|
||||
err = db.DB.QueryRowContext(ctx, `
|
||||
SELECT platform_id, workspace, title
|
||||
FROM agent_platform_sessions
|
||||
WHERE agent_name = $1 AND slug = $2
|
||||
AND ($3 = '' OR workspace = $3)
|
||||
ORDER BY COALESCE(updated_at, reported_at) DESC
|
||||
LIMIT 1
|
||||
`, agentName, slug, ws).Scan(&platformID, &realWorkspace, &title)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", "", "", ErrSessionNotFound
|
||||
}
|
||||
return platformID, realWorkspace, title, err
|
||||
}
|
||||
|
||||
// FindSessionByPlatformID 找出已经接管了某条平台会话的本侧会话。
|
||||
//
|
||||
// 返回 ErrSessionNotFound 表示还没被接管。归档的也算 —— 让归档过的会话
|
||||
// 重新被接管会造出第二条本侧会话，同一条 TUI 对话在邮箱里就裂成两截。
|
||||
// 需要恢复的话人应该去取消归档。
|
||||
func FindSessionByPlatformID(ctx context.Context, agentName, platformID string) (uuid.UUID, error) {
|
||||
var id uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT s.session_id
|
||||
FROM sessions s
|
||||
WHERE s.platform_id = $1
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM mails m
|
||||
WHERE m.session_id = s.session_id
|
||||
AND (m.to_name = $2 OR m.from_name = $2 OR `+db.CCHas("m.cc_list", 2)+`)
|
||||
)
|
||||
ORDER BY s.updated_at DESC
|
||||
LIMIT 1
|
||||
`, platformID, agentName).Scan(&id)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return uuid.Nil, ErrSessionNotFound
|
||||
}
|
||||
return id, err
|
||||
}
|
||||
|
||||
// AdoptPlatformSession 接管一条平台会话：建本侧会话并绑定 platform_id。
|
||||
//
|
||||
// alias 用平台自己的 slug —— 「别名复用平台命名」是既定决策，而且人在补全里
|
||||
// 看到的就是那个 slug，投递后别名换成别的会让他找不到自己刚发的信。
|
||||
//
|
||||
// workspace 用平台会话的真实 cwd 而不是地址里的 path 位，理由见
|
||||
// FindPlatformSession 的注释。
|
||||
func AdoptPlatformSession(ctx context.Context, agentName, platformID, slug, workspace, subject string) (uuid.UUID, error) {
|
||||
// slug 可能与本侧某条无关会话撞名（别名全局唯一）。撞了就加后缀 ——
|
||||
// EnsureSessionAlias 已有这套逻辑，这里先建后命名即可。
|
||||
id, err := CreateSession(ctx, nil, agentName, subject, workspace)
|
||||
if err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET platform_id = $1 WHERE session_id = $2`,
|
||||
platformID, id); err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
// 显式写入档位与强制力：接管一条平台会话没有父会话，
|
||||
// 只靠 DB 默认值会在「schema 列定义变动」或「迁移补列给了不同默认」时
|
||||
// 静默偏离预期 —— 显式写 'workspace' 是唯一可靠表述「这条会话是新接管的，
|
||||
// 没有继承来源」的方式。与 me.go 新建会话那条路径一致。
|
||||
if _, err := SetSessionPermissionMode(ctx, id, models.DefaultPermissionMode); err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
_ = SetSessionEnforcement(ctx, id, AgentModeEnforcement(ctx, agentName))
|
||||
// 别名尽量用 slug；撞名时 EnsureSessionAlias 自动加后缀
|
||||
_, _ = EnsureSessionAlias(ctx, id, slug)
|
||||
return id, nil
|
||||
}
|
||||
|
||||
// PlatformIDOf 读一条本侧会话绑定的平台会话 id（空 = 不是接管来的）。
|
||||
//
|
||||
// 投递时要把它放进 SSE 事件：插件据此决定 resume 还是新建。
|
||||
//
|
||||
// 只在「已知这条会话只有一个参与方」时用它。有抄送时必须用
|
||||
// PlatformSessionFor 拿到归属方 —— 理由见那个函数。
|
||||
func PlatformIDOf(ctx context.Context, sessionID uuid.UUID) string {
|
||||
pid, _ := PlatformSessionFor(ctx, sessionID)
|
||||
return pid
|
||||
}
|
||||
|
||||
// PlatformSessionFor 返回一条本侧会话绑定的平台会话 id **及其归属 Agent**。
|
||||
//
|
||||
// # 为什么归属方是必须的
|
||||
//
|
||||
// `platform_id` 是**会话级**的一个值，而一封邮件可以有多个参与方。
|
||||
// 把它无差别推给所有人，收到的一方会拿它去自己的磁盘上找会话文件 ——
|
||||
// 那个 id 属于别的平台。
|
||||
//
|
||||
// 生产实测：会话 `16845133` 接管了 pi 的会话 `01a05a5e-…`，而那封邮件抄送了
|
||||
// `dsh@/home/program/agentmail.new`。DSH 收到同一个 platform_session_id，
|
||||
// 在 `~/.dsh/sessions/` 里查不到（那是 `/root/.pi/agent/sessions/` 下的文件），
|
||||
// 于是走进「平台侧会话已删」那条防线抛错。那道防线本身是对的（N-8：
|
||||
// 不能退回新建，否则人在界面上看不到这封邮件带来的对话），它拦下的却是
|
||||
// 「别人的会话」—— 邮件因此静默消失，而插件侧的日志走的是不进 journalctl
|
||||
// 的通道，连线索都没有。
|
||||
//
|
||||
// 归属方以镜像（`agent_platform_sessions.agent_name`，Agent 自己上报的）为准；
|
||||
// 镜像整表替换，平台侧删了会话那行就没了，此时退回 `sessions.from_agent` ——
|
||||
// `AdoptPlatformSession` 建会话时把归属 Agent 写在那里，是可靠的第二来源。
|
||||
func PlatformSessionFor(ctx context.Context, sessionID uuid.UUID) (platformID, owner string) {
|
||||
var pid, fromAgent string
|
||||
var mirrored *string
|
||||
if err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COALESCE(s.platform_id, ''), COALESCE(s.from_agent, ''), aps.agent_name
|
||||
FROM sessions s
|
||||
LEFT JOIN agent_platform_sessions aps
|
||||
ON aps.platform_id = s.platform_id AND COALESCE(s.platform_id, '') <> ''
|
||||
WHERE s.session_id = $1`, sessionID).Scan(&pid, &fromAgent, &mirrored); err != nil {
|
||||
return "", ""
|
||||
}
|
||||
if pid == "" {
|
||||
return "", ""
|
||||
}
|
||||
if mirrored != nil && *mirrored != "" {
|
||||
return pid, *mirrored
|
||||
}
|
||||
return pid, fromAgent
|
||||
}
|
||||
354
server/internal/repo/platform_sessions_test.go
Normal file
354
server/internal/repo/platform_sessions_test.go
Normal file
@ -0,0 +1,354 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
310
server/internal/repo/quota.go
Normal file
310
server/internal/repo/quota.go
Normal file
@ -0,0 +1,310 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 配额 ----------
|
||||
//
|
||||
// **配额是任务的属性，不是 Agent 的属性。**
|
||||
//
|
||||
// 真正的约束在 `sessions.max_rounds`（见本文件末尾的「会话级往返预算」）：
|
||||
// 每条会话独立计数，人在派活时给、在对话页里随时调。
|
||||
//
|
||||
// `agents` 表这边只剩两样东西：
|
||||
//
|
||||
// default_rounds —— 派给这个 Agent 的**新任务**默认多少个来回。
|
||||
// 不同 Agent 能力不同（跑测试的小工具 vs 重构整个模块），
|
||||
// 默认值分开设才合理。
|
||||
//
|
||||
// used_rounds —— 纯统计，累计发信数。**不再拦任何请求。**
|
||||
// 它原本是「终身额度」：跑满就得管理员手工重置才能再干活，
|
||||
// 而 Agent 是长期在线的 —— 终身额度是错的工具。
|
||||
// 保留是因为「这个 Agent 一共发了多少信」本身有观测价值。
|
||||
//
|
||||
// 防止 Agent 用 `.new` 开一串新会话绕过预算，靠的是**新建会话速率限制**
|
||||
// （见 sessionRateLimiter），而不是终身额度。
|
||||
// AgentStats 是一个 Agent 的配额默认值与累计统计。
|
||||
//
|
||||
// 没有 Remaining / Unlimited 字段：这里不再有「剩余额度」的概念 ——
|
||||
// 额度属于会话（SessionBudget），这里只有「新任务默认多少来回」与「一共发了多少信」。
|
||||
type AgentStats struct {
|
||||
AgentName string `json:"agent_name"`
|
||||
// DefaultRounds 派给该 Agent 的新任务默认多少个来回（0 = 不限）
|
||||
DefaultRounds int `json:"default_rounds"`
|
||||
// SentTotal 累计发信数（纯统计，不拦请求）
|
||||
SentTotal int `json:"sent_total"`
|
||||
// ActiveSessions 该 Agent 参与的未归档会话数，配合默认值判断设多少合适
|
||||
ActiveSessions int `json:"active_sessions"`
|
||||
// Status 是 agents.status：online / offline / disabled。
|
||||
// 管理页靠它决定显示「停用」还是「恢复」。
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
// DefaultRoundsFor 读取该 Agent 的新任务默认预算。
|
||||
//
|
||||
// Agent 不存在时返回全局兜底值而非报错：派活的人不该因为「对方还没注册」
|
||||
// 就拿不到一个合理的默认预算 —— 邮件本来就支持发给尚未上线的收件人。
|
||||
func DefaultRoundsFor(ctx context.Context, agentName string) int {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(default_rounds, 0) FROM agents WHERE agent_name = $1`,
|
||||
agentName).Scan(&n)
|
||||
if err != nil || n < 0 {
|
||||
return fallbackDefaultRounds
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// fallbackDefaultRounds 是 Agent 未注册时的兜底默认预算。
|
||||
// 与建表默认值保持一致；改这里要同时改两份 schema。
|
||||
const fallbackDefaultRounds = 20
|
||||
|
||||
// SetDefaultRounds 设置该 Agent 的新任务默认预算（0 = 不限）。
|
||||
func SetDefaultRounds(ctx context.Context, agentName string, n int) (AgentStats, error) {
|
||||
if n < 0 {
|
||||
n = 0
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET default_rounds = $2 WHERE agent_name = $1`, agentName, n)
|
||||
if err != nil {
|
||||
return AgentStats{}, err
|
||||
}
|
||||
if k, _ := tag.RowsAffected(); k == 0 {
|
||||
return AgentStats{}, fmt.Errorf("agent %q 不存在", agentName)
|
||||
}
|
||||
return GetAgentStats(ctx, agentName)
|
||||
}
|
||||
|
||||
// GetAgentStats 读取某 Agent 的默认预算与累计统计。
|
||||
func GetAgentStats(ctx context.Context, agentName string) (AgentStats, error) {
|
||||
var st AgentStats
|
||||
st.AgentName = agentName
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COALESCE(default_rounds, 0), COALESCE(used_rounds, 0)
|
||||
FROM agents WHERE agent_name = $1`, agentName,
|
||||
).Scan(&st.DefaultRounds, &st.SentTotal)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return AgentStats{}, fmt.Errorf("agent %q 不存在", agentName)
|
||||
}
|
||||
if err != nil {
|
||||
return AgentStats{}, err
|
||||
}
|
||||
st.ActiveSessions = countActiveSessionsFor(ctx, agentName)
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// countActiveSessionsFor 统计该 Agent 参与的未归档会话数。
|
||||
// 查不出来返回 0：这只是个展示用的数字，不该让整个统计接口失败。
|
||||
func countActiveSessionsFor(ctx context.Context, agentName string) int {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(DISTINCT s.session_id)
|
||||
FROM sessions s
|
||||
JOIN mails m ON m.session_id = s.session_id
|
||||
WHERE s.status <> 'archived'
|
||||
AND (m.from_name = $1 OR m.to_name = $1 OR `+db.CCHas("m.cc_list", 1)+`)
|
||||
`, agentName).Scan(&n)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// BumpSentCount 累加发信统计。
|
||||
//
|
||||
// **绝不拦请求**：它是观测数据，不是额度。返回值只有 error，
|
||||
// 而且调用方应当忽略它 —— 统计写失败不该让一封已经该发出的邮件失败。
|
||||
func BumpSentCount(ctx context.Context, agentName string) {
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET used_rounds = COALESCE(used_rounds, 0) + 1 WHERE agent_name = $1`,
|
||||
agentName)
|
||||
}
|
||||
|
||||
// ListAgentStats 列出所有 Agent 的默认预算与统计（管理员视图）。
|
||||
func ListAgentStats(ctx context.Context) ([]AgentStats, error) {
|
||||
// 带上 status：管理页靠它区分「在线 / 离线 / 已停用」并决定显示
|
||||
// 「停用」还是「恢复」按钮。不过滤 disabled —— 这里是唯一能把已停用的
|
||||
// Agent 恢复回来的地方，过滤掉就再也找不到它了。
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT agent_name, COALESCE(default_rounds, 0), COALESCE(used_rounds, 0),
|
||||
COALESCE(status, 'offline')
|
||||
FROM agents ORDER BY agent_name`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []AgentStats{}
|
||||
for rows.Next() {
|
||||
var st AgentStats
|
||||
if err := rows.Scan(&st.AgentName, &st.DefaultRounds, &st.SentTotal, &st.Status); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, st)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// 会话数逐个查：Agent 数量是个位数到几十，不值得为它写一个 GROUP BY 的联合查询
|
||||
for i := range out {
|
||||
out[i].ActiveSessions = countActiveSessionsFor(ctx, out[i].AgentName)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ---------- 转发 ----------
|
||||
|
||||
// ForwardSource 是被转发邮件的必要信息。
|
||||
type ForwardSource struct {
|
||||
Mail *models.Mail
|
||||
Session uuid.UUID
|
||||
}
|
||||
|
||||
// LoadForwardSource 读取待转发的邮件，并校验转发者确实参与过该邮件
|
||||
//（收件人、发件人或被抄送方之一）。防止凭 mail_id 转发别人的邮件。
|
||||
func LoadForwardSource(ctx context.Context, mailID uuid.UUID, actor string) (*models.Mail, error) {
|
||||
m, err := GetMailByID(ctx, mailID)
|
||||
if err != nil {
|
||||
return nil, ErrMailNotFound
|
||||
}
|
||||
|
||||
if m.FromName == actor || m.ToName == actor {
|
||||
return m, nil
|
||||
}
|
||||
for _, cc := range m.CCList {
|
||||
if cc.Name == actor {
|
||||
return m, nil
|
||||
}
|
||||
}
|
||||
return nil, ErrForwardNotAllowed
|
||||
}
|
||||
|
||||
var (
|
||||
// ErrMailNotFound 待转发的邮件不存在
|
||||
ErrMailNotFound = errors.New("mail not found")
|
||||
// ErrForwardNotAllowed 转发者未参与该邮件
|
||||
ErrForwardNotAllowed = errors.New("not a participant of that mail")
|
||||
)
|
||||
|
||||
// ---------- 会话级往返预算 ----------
|
||||
//
|
||||
// 配额的真实语义是「这件事值得多少个来回」——那是**任务**的属性，不是 Agent 的属性。
|
||||
// 只有 agents.max_rounds 一个全局计数器时有两个问题：
|
||||
// 1. 两个并行任务互相抢额度：给紧急任务留的份被另一条线索吃掉
|
||||
// 2. used_rounds 单调递增，跑满就得管理员手工重置才能再干活
|
||||
// 因此预算下沉到会话，由人在写信时给、在对话页里随时调。
|
||||
//
|
||||
// **两层都要过**：会话预算 + Agent 全局配额。少了后者，Agent 自己 `.new` 开一串会话
|
||||
// 每条都是全新预算，全局上限形同虚设；少了前者，就回到抢额度的老问题。
|
||||
|
||||
// SessionBudget 是一个会话的往返预算快照。
|
||||
type SessionBudget struct {
|
||||
SessionID uuid.UUID `json:"session_id"`
|
||||
Max int `json:"max_rounds"` // 0 = 本会话不限
|
||||
Used int `json:"used_rounds"`
|
||||
Remaining int `json:"remaining"` // 不限时为 -1
|
||||
Unlimited bool `json:"unlimited"`
|
||||
}
|
||||
|
||||
func makeSessionBudget(id uuid.UUID, max, used int) SessionBudget {
|
||||
b := SessionBudget{SessionID: id, Max: max, Used: used, Unlimited: max <= 0}
|
||||
if b.Unlimited {
|
||||
b.Remaining = -1
|
||||
return b
|
||||
}
|
||||
if r := max - used; r > 0 {
|
||||
b.Remaining = r
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// ErrSessionBudgetExhausted 表示该会话的往返预算已用尽。
|
||||
var ErrSessionBudgetExhausted = errors.New("session budget exhausted")
|
||||
|
||||
// GetSessionBudget 读取会话预算。
|
||||
func GetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
var max, used int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(max_rounds,0), COALESCE(used_rounds,0) FROM sessions WHERE session_id = $1`,
|
||||
id).Scan(&max, &used)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
return makeSessionBudget(id, max, used), nil
|
||||
}
|
||||
|
||||
// ConsumeSessionBudget 原子地占用会话的一次往返。
|
||||
//
|
||||
// 与 ConsumeQuota 同理：判断与自增必须在同一条 UPDATE 里（WHERE used_rounds < max_rounds），
|
||||
// 否则并发发信会双双通过检查再各自 +1，把预算刷穿。
|
||||
func ConsumeSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
tag, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) + 1
|
||||
WHERE session_id = $1
|
||||
AND (COALESCE(max_rounds,0) <= 0 OR COALESCE(used_rounds,0) < max_rounds)
|
||||
`, id)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
b, bErr := GetSessionBudget(ctx, id)
|
||||
if bErr != nil {
|
||||
return SessionBudget{}, bErr
|
||||
}
|
||||
return b, ErrSessionBudgetExhausted
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// SetSessionBudget 设置会话预算上限（0 = 不限）。
|
||||
//
|
||||
// 允许把上限调到低于已用次数：那表示「就到这里为止」，是人的合法意图，
|
||||
// 不该因为算不出正的剩余量就拒绝。此时 Remaining 为 0，下次发信即被拦。
|
||||
func SetSessionBudget(ctx context.Context, id uuid.UUID, max int) (SessionBudget, error) {
|
||||
if max < 0 {
|
||||
max = 0
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET max_rounds = $2, updated_at = NOW() WHERE session_id = $1`, id, max)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// ResetSessionBudget 把该会话的已用次数归零（上限不变）。
|
||||
func ResetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET used_rounds = 0, updated_at = NOW() WHERE session_id = $1`, id)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// RefundSessionBudget 退还一次往返。
|
||||
//
|
||||
// 会话预算先扣、Agent 全局配额后扣，全局那层拦下时必须把会话这次还回去，
|
||||
// 否则会话预算白掉一格 —— 那次往返实际上没有发生。
|
||||
func RefundSessionBudget(ctx context.Context, id uuid.UUID) {
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) - 1
|
||||
WHERE session_id = $1 AND COALESCE(used_rounds,0) > 0`, id)
|
||||
}
|
||||
125
server/internal/repo/quota_test.go
Normal file
125
server/internal/repo/quota_test.go
Normal file
@ -0,0 +1,125 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// setupTestDB 起一个临时 SQLite 库并建表，供配额测试使用。
|
||||
// 直接用真实的 SQLite 而非 mock：配额的正确性核心在于「判断与自增在同一条 UPDATE 里」，
|
||||
// 这正是只有真实数据库才能验证的部分。
|
||||
func setupTestDB(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
dir := t.TempDir()
|
||||
if err := db.Connect(context.Background(), filepath.Join(dir, "test.db")); err != nil {
|
||||
t.Fatalf("connect: %v", err)
|
||||
}
|
||||
if err := db.Migrate(context.Background()); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
t.Cleanup(db.Close)
|
||||
}
|
||||
|
||||
func seedAgent(t *testing.T, name string, max int) {
|
||||
t.Helper()
|
||||
_, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO agents (agent_name, secret, platform, default_rounds) VALUES ($1, 'x', 'test', $2)`,
|
||||
name, max)
|
||||
if err != nil {
|
||||
t.Fatalf("seed agent: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// default_rounds 是「派给这个 Agent 的新任务默认多少个来回」，
|
||||
// 不是会拦请求的终身额度 —— 真正的额度在 sessions.max_rounds 上。
|
||||
func TestDefaultRoundsRoundTrip(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "bot", 0)
|
||||
ctx := context.Background()
|
||||
|
||||
st, err := SetDefaultRounds(ctx, "bot", 15)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.DefaultRounds != 15 {
|
||||
t.Fatalf("default_rounds = %d，期望 15", st.DefaultRounds)
|
||||
}
|
||||
if got := DefaultRoundsFor(ctx, "bot"); got != 15 {
|
||||
t.Fatalf("DefaultRoundsFor = %d，期望 15", got)
|
||||
}
|
||||
|
||||
// 负数归一为 0（不限），而不是造出一个永远发不出信的默认值
|
||||
if st, err = SetDefaultRounds(ctx, "bot", -3); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.DefaultRounds != 0 {
|
||||
t.Fatalf("负数应归一为 0，实际 %d", st.DefaultRounds)
|
||||
}
|
||||
}
|
||||
|
||||
// 未注册的 Agent 取默认预算时给兜底值而不是报错：
|
||||
// 派活的人不该因为「对方还没上线」就拿不到一个合理默认值 ——
|
||||
// 邮件本来就支持发给尚未上线的收件人。
|
||||
func TestDefaultRoundsForUnknownAgentFallsBack(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
if got := DefaultRoundsFor(context.Background(), "ghost"); got != fallbackDefaultRounds {
|
||||
t.Fatalf("未注册 Agent 应回落到 %d，实际 %d", fallbackDefaultRounds, got)
|
||||
}
|
||||
}
|
||||
|
||||
// BumpSentCount 是纯统计：只累加，绝不拦请求，也绝不返回错误
|
||||
func TestBumpSentCountOnlyCounts(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "bot", 0)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
BumpSentCount(ctx, "bot")
|
||||
}
|
||||
st, err := GetAgentStats(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.SentTotal != 5 {
|
||||
t.Fatalf("累计发信 = %d，期望 5", st.SentTotal)
|
||||
}
|
||||
|
||||
// 不存在的 Agent 也不该 panic 或报错 —— 它只是没有行可更新
|
||||
BumpSentCount(ctx, "ghost")
|
||||
}
|
||||
|
||||
func TestGetAgentStatsUnknownAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
if _, err := GetAgentStats(context.Background(), "ghost"); err == nil {
|
||||
t.Fatal("不存在的 Agent 应报错")
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAgentStats(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "alpha", 0)
|
||||
seedAgent(t, "beta", 0)
|
||||
ctx := context.Background()
|
||||
SetDefaultRounds(ctx, "alpha", 5)
|
||||
|
||||
list, err := ListAgentStats(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(list) != 2 {
|
||||
t.Fatalf("应有 2 个 Agent，实际 %d", len(list))
|
||||
}
|
||||
// 按名字排序，alpha 在前
|
||||
if list[0].AgentName != "alpha" || list[0].DefaultRounds != 5 {
|
||||
t.Fatalf("alpha 的记录不对：%+v", list[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
64
server/internal/repo/ratelimit.go
Normal file
64
server/internal/repo/ratelimit.go
Normal file
@ -0,0 +1,64 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// RateLimitCheckAndRecord 原子地检查 bucket 在 window 内的事件数是否超过 limit。
|
||||
// 未超过则同时记录本次事件（判断与写入在同一个事务里，防并发刷穿）。
|
||||
// DB 不可用时放行（宁可放开限速也不能让用户完全无法使用）。
|
||||
func RateLimitCheckAndRecord(ctx context.Context, bucket string, window time.Duration, limit int) (allowed bool, retryAfter int) {
|
||||
now := time.Now()
|
||||
cutoff := now.Add(-window)
|
||||
|
||||
// 用 IMMEDIATE 事务：SQLite 的 IMMEDIATE 会在开始时获取 RESERVED 锁，
|
||||
// 防止其他写事务同时进入 COMMIT 阶段。这是 SQLite 并发写的正确方式。
|
||||
tx, err := db.DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return true, 0 // DB 不可用 → 放行
|
||||
}
|
||||
defer tx.Rollback() // Commit 成功后 Rollback 是 no-op
|
||||
|
||||
// 清理过期记录
|
||||
tx.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1 AND ts < $2`, bucket, cutoff)
|
||||
|
||||
// 统计当前窗口内事件数
|
||||
var count int
|
||||
err = tx.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||||
bucket, cutoff).Scan(&count)
|
||||
if err != nil {
|
||||
return true, 0
|
||||
}
|
||||
|
||||
if count >= limit {
|
||||
var earliest time.Time
|
||||
err = tx.QueryRowContext(ctx,
|
||||
`SELECT MIN(ts) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||||
bucket, cutoff).Scan(&earliest)
|
||||
if err == nil && !earliest.IsZero() {
|
||||
retry := int(earliest.Add(window).Sub(now).Seconds()) + 1
|
||||
if retry < 1 {
|
||||
retry = 1
|
||||
}
|
||||
return false, retry
|
||||
}
|
||||
return false, 60
|
||||
}
|
||||
|
||||
// 记账
|
||||
tx.ExecContext(ctx,
|
||||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, now)
|
||||
tx.Commit()
|
||||
return true, 0
|
||||
}
|
||||
|
||||
// RateLimitReset 清除指定 bucket 的所有记录（登录成功后调用）。
|
||||
func RateLimitReset(ctx context.Context, bucket string) {
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1`, bucket)
|
||||
}
|
||||
87
server/internal/repo/relay.go
Normal file
87
server/internal/repo/relay.go
Normal file
@ -0,0 +1,87 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 插件自动转发（免配额通道） ----------
|
||||
//
|
||||
// **基本原则：配额约束的是模型的自主发信，不是 harness 的转发。**
|
||||
//
|
||||
// 配额存在的意义是防止 Agent 无限自我循环。而插件代劳搬运的两类消息不属于此列：
|
||||
// 1. 平台原生的权限询问（opencode 的 permission.updated）—— 不转给人，人就看不到，
|
||||
// Agent 卡在那里等一个永远不会来的回答
|
||||
// 2. 本轮的最终总结（session.idle 时最后一条 assistant 消息）—— 模型已经把话说完了，
|
||||
// 插件只是把它搬到邮件里；对它收费会导致配额用尽时 Agent 连交代都做不了
|
||||
//
|
||||
// 防滥用不靠计数，靠**幂等键**：relay_key 是上游那条消息的稳定标识
|
||||
// （permission id / assistant message id）。唯一约束让同一条上游消息只能转一次，
|
||||
// 于是插件重试与 SSE 重放不会产生第二封，想多转就得拿出不同的上游消息 id ——
|
||||
// 而那些 id 由平台生成，模型伪造不出来。
|
||||
|
||||
// ErrRelayDuplicate 表示这条上游消息已经转发过了。
|
||||
var ErrRelayDuplicate = errors.New("relay already recorded")
|
||||
|
||||
// ClaimRelay 占用一次免配额转发名额。
|
||||
//
|
||||
// 判断与占用在同一条 INSERT 里（靠主键唯一约束），并发重试下只有一个能成功 ——
|
||||
// 分成「先查有没有、再插入」两步的话，插件的两次重试会双双通过检查各插一条。
|
||||
//
|
||||
// 返回 ErrRelayDuplicate 表示重复，调用方应当据此跳过发信而不是报错：
|
||||
// 重复转发是插件重试的正常结果，不是故障。
|
||||
func ClaimRelay(ctx context.Context, agentName, relayKey, kind string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO relayed_mails (agent_name, relay_key, kind) VALUES ($1, $2, $3)`,
|
||||
agentName, relayKey, kind)
|
||||
if err != nil {
|
||||
if db.IsUniqueViolation(err) {
|
||||
return ErrRelayDuplicate
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BindRelayMail 把已占用的名额关联到真正发出的邮件，便于事后审计
|
||||
// 「这封免配额的信是从哪条上游消息来的」。
|
||||
//
|
||||
// 关联失败不该让发信失败：邮件已经入库，缺一条审计关联不影响功能。
|
||||
func BindRelayMail(ctx context.Context, agentName, relayKey string, mailID uuid.UUID) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE relayed_mails SET mail_id = $1 WHERE agent_name = $2 AND relay_key = $3`,
|
||||
mailID, agentName, relayKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReleaseRelay 撤销名额占用。
|
||||
//
|
||||
// 占用成功但发信失败时必须还回去，否则那条上游消息永远转不出来了 ——
|
||||
// 幂等键会一直认为它已经转过。
|
||||
func ReleaseRelay(ctx context.Context, agentName, relayKey string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM relayed_mails WHERE agent_name = $1 AND relay_key = $2 AND mail_id IS NULL`,
|
||||
agentName, relayKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// RelayKeyForMail 反查某封邮件对应的上游消息 id。
|
||||
//
|
||||
// 人类决策一条权限请求后，插件需要知道该回复 opencode 的哪个 permission ——
|
||||
// 光有 AgentMail 的 mail_id 是不够的，两边的 id 空间不同。
|
||||
// 插件重启后内存映射会丢，所以这个映射必须在服务端持久化。
|
||||
//
|
||||
// 无记录时返回空串（例如旧数据，或压根没走 relay 通道的请求）。
|
||||
func RelayKeyForMail(ctx context.Context, mailID uuid.UUID) (string, string) {
|
||||
var key, kind string
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT relay_key, kind FROM relayed_mails WHERE mail_id = $1 LIMIT 1`,
|
||||
mailID).Scan(&key, &kind)
|
||||
if err != nil {
|
||||
return "", ""
|
||||
}
|
||||
return key, kind
|
||||
}
|
||||
83
server/internal/repo/relayhops.go
Normal file
83
server/internal/repo/relayhops.go
Normal file
@ -0,0 +1,83 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 连续 relay 跳数限制 —— 防止两个 Agent 靠自动转发互相唤醒到无穷。
|
||||
//
|
||||
// # 这是什么问题
|
||||
//
|
||||
// 每个插件都在「一轮结束时把模型最后那段话自动发回去」（契约 B-5）。
|
||||
// 当收件方也是一个装了同类插件的 Agent 时，这封信唤醒对方 → 对方跑一轮 →
|
||||
// 对方也自动回一封 → 循环。**双方都没有「决定继续」，因为双方都不在做决定** ——
|
||||
// 发信是插件代劳的。
|
||||
//
|
||||
// 生产上真实发生过：会话 f3d824ce（dsh 与 opencode 联调 llmsproxy）共 41 封，
|
||||
// 最后一封人类意图的邮件之后，**每一封都是 relay:summary**，
|
||||
// 间隔从 15 分钟一路缩到 5 秒，内容已无新增信息。
|
||||
//
|
||||
// # 为什么 relay_key 拦不住
|
||||
//
|
||||
// 它是幂等键，职责是「同一条上游消息不重复转发」，这一点它做对了。
|
||||
// 但每一轮都是**真正不同**的新消息：opencode 侧是 assistant message id
|
||||
// （msg_0655bbf6…、msg_0656cc7fc…），dsh 侧是事件计数（…:12220、…:13347）。
|
||||
// 每次 ClaimRelay 都合法通过。
|
||||
//
|
||||
// # 为什么需要两道防线
|
||||
//
|
||||
// 主防线是「免配额只给发往人类的 relay」（见 handler.SendMail）：
|
||||
// Agent→Agent 的自动转发转而消耗会话预算，max_rounds 会截断它。
|
||||
//
|
||||
// 但那还不够：预算给得大（比如 200）时，两个 Agent 仍能烧掉 200 个来回；
|
||||
// 而故障报告这类**必须**走 relay 的邮件也需要受约束。因此这里再加一道
|
||||
// 与预算无关的硬上限：一条会话里**连续**的 relay 邮件不得超过 maxRelayHops。
|
||||
//
|
||||
// 「连续」是关键：只要中间有一封自主发信（模型真的决定说什么）或人类插话，
|
||||
// 计数就归零。这让正常的「模型回一封、插件补一封总结」不受影响，
|
||||
// 只掐住「全程无人决策」的那种回路。
|
||||
//
|
||||
// hop_limit 列早就在 schema 里（DEFAULT 5）却从没有人读它 —— 它显然
|
||||
// 就是为这件事准备的。这里把它接上，取同一个默认值。
|
||||
const maxRelayHops = 5
|
||||
|
||||
// CountTrailingRelayHops 数会话尾部**连续**的 relay 邮件数。
|
||||
//
|
||||
// 从最新一封往前扫，遇到第一封非 relay 邮件即停。返回值即「若本次再发一封
|
||||
// relay，它会是第几跳」的前一个数。
|
||||
//
|
||||
// 判据用 relayed_mails 的存在性而不是 mails 上的某个标记：
|
||||
// relay 身份本来就记在那张表里，在 mails 上再存一份等于给同一事实留两个答案。
|
||||
func CountTrailingRelayHops(ctx context.Context, sessionID uuid.UUID) (int, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT CASE WHEN r.mail_id IS NULL THEN 0 ELSE 1 END AS is_relay
|
||||
FROM mails m
|
||||
LEFT JOIN relayed_mails r ON r.mail_id = m.mail_id
|
||||
WHERE m.session_id = $1
|
||||
ORDER BY m.created_at DESC, m.mail_id DESC
|
||||
`, sessionID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
hops := 0
|
||||
for rows.Next() {
|
||||
var isRelay int
|
||||
if err := rows.Scan(&isRelay); err != nil {
|
||||
return hops, err
|
||||
}
|
||||
if isRelay == 0 {
|
||||
// 遇到一封自主发信/人类邮件：链条到此为止
|
||||
break
|
||||
}
|
||||
hops++
|
||||
}
|
||||
return hops, rows.Err()
|
||||
}
|
||||
|
||||
// MaxRelayHops 暴露上限供错误文案与测试使用。
|
||||
func MaxRelayHops() int { return maxRelayHops }
|
||||
126
server/internal/repo/relayhops_test.go
Normal file
126
server/internal/repo/relayhops_test.go
Normal file
@ -0,0 +1,126 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 连续 relay 跳数上限守的是一个真实事故：会话 f3d824ce（dsh 与 opencode 联调
|
||||
// llmsproxy）共 41 封，最后一封人类意图的邮件之后每一封都是 relay:summary，
|
||||
// 间隔从 15 分钟一路缩到 5 秒。双方都没有「决定继续」，因为双方都不在做决定 ——
|
||||
// 发信是插件代劳的，而免配额通道让整个回路里没有任何一处在计数。
|
||||
|
||||
func TestTrailingRelayHopsEmptySession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "空会话", "/tmp/ws")
|
||||
n, err := CountTrailingRelayHops(ctx, sid)
|
||||
if err != nil {
|
||||
t.Fatalf("空会话应正常返回: %v", err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Fatalf("空会话跳数应为 0，实为 %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrailingRelayHopsCountsOnlyRelay(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "混合", "/tmp/ws")
|
||||
// 人类发一封（非 relay）
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/ws", nil)
|
||||
// 插件连续转发三封
|
||||
for i := 0; i < 3; i++ {
|
||||
seedRelayMail(t, ctx, sid, "dsh", "opencode")
|
||||
}
|
||||
|
||||
n, err := CountTrailingRelayHops(ctx, sid)
|
||||
if err != nil {
|
||||
t.Fatalf("数跳数: %v", err)
|
||||
}
|
||||
if n != 3 {
|
||||
t.Fatalf("尾部连续 relay 应为 3，实为 %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrailingRelayHopsResetsOnAutonomousSend(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 「连续」是这条规则的关键：中间只要有一封自主发信（模型真的决定说什么）
|
||||
// 或人类插话，计数就归零。否则正常的「模型回一封、插件补一封总结」
|
||||
// 会被误判成回路。
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "打断", "/tmp/ws")
|
||||
for i := 0; i < 4; i++ {
|
||||
seedRelayMail(t, ctx, sid, "dsh", "opencode")
|
||||
}
|
||||
// 模型亲手发了一封 —— 链条到此为止
|
||||
mustMail(t, sid, "dsh", "", "opencode", "/tmp/ws", nil)
|
||||
seedRelayMail(t, ctx, sid, "opencode", "dsh")
|
||||
|
||||
n, err := CountTrailingRelayHops(ctx, sid)
|
||||
if err != nil {
|
||||
t.Fatalf("数跳数: %v", err)
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("自主发信之后只剩 1 跳，实为 %d", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrailingRelayHopsReachesLimit(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 攒到上限：此时 handler 应当拒绝下一封 relay。
|
||||
sid, _ := CreateSession(ctx, nil, "admin", "到顶", "/tmp/ws")
|
||||
mustMail(t, sid, "admin", "", "dsh", "/tmp/ws", nil)
|
||||
for i := 0; i < MaxRelayHops(); i++ {
|
||||
seedRelayMail(t, ctx, sid, "dsh", "opencode")
|
||||
}
|
||||
|
||||
n, _ := CountTrailingRelayHops(ctx, sid)
|
||||
if n < MaxRelayHops() {
|
||||
t.Fatalf("应达到上限 %d，实为 %d", MaxRelayHops(), n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrailingRelayHopsIsPerSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 一条会话的回路不该影响另一条：两个 Agent 在 A 会话里刷爆了，
|
||||
// B 会话的正常自动转发仍应放行。
|
||||
a, _ := CreateSession(ctx, nil, "admin", "A", "/tmp/a")
|
||||
b, _ := CreateSession(ctx, nil, "admin", "B", "/tmp/b")
|
||||
for i := 0; i < 5; i++ {
|
||||
seedRelayMail(t, ctx, a, "dsh", "opencode")
|
||||
}
|
||||
seedRelayMail(t, ctx, b, "dsh", "admin")
|
||||
|
||||
na, _ := CountTrailingRelayHops(ctx, a)
|
||||
nb, _ := CountTrailingRelayHops(ctx, b)
|
||||
if na != 5 || nb != 1 {
|
||||
t.Fatalf("跳数应按会话独立计：A=%d（期望 5）B=%d（期望 1）", na, nb)
|
||||
}
|
||||
}
|
||||
|
||||
// seedRelayMail 建一封走 relay 通道的邮件（同时占幂等键并关联 mail_id），
|
||||
// 复刻 handler.SendMail 的真实写入顺序。
|
||||
func seedRelayMail(t *testing.T, ctx context.Context, sid uuid.UUID, from, to string) {
|
||||
t.Helper()
|
||||
key := "relay-" + uuid.NewString()
|
||||
if err := ClaimRelay(ctx, from, key, "summary"); err != nil {
|
||||
t.Fatalf("占幂等键: %v", err)
|
||||
}
|
||||
mid, err := CreateMail(ctx, sid, nil, from, "", to, "", "Re: 主题", "正文", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("建邮件: %v", err)
|
||||
}
|
||||
if err := BindRelayMail(ctx, from, key, mid); err != nil {
|
||||
t.Fatalf("关联 relay: %v", err)
|
||||
}
|
||||
}
|
||||
1502
server/internal/repo/repo.go
Normal file
1502
server/internal/repo/repo.go
Normal file
File diff suppressed because it is too large
Load Diff
228
server/internal/repo/repo_test.go
Normal file
228
server/internal/repo/repo_test.go
Normal file
@ -0,0 +1,228 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ListSessionsFor 的 Scan 列数必须与 SELECT 一致。
|
||||
//
|
||||
// 这个测试存在的理由：预算两列（max_rounds/used_rounds）加进了 SELECT 却忘了加进
|
||||
// Scan，于是 /me/sessions 整个 500 —— 联系人栏一条数据都拉不到，
|
||||
// 而错误信息只是 "Failed to list sessions"，看不出是列数不匹配。
|
||||
// 列数错位是纯结构问题，一个最小用例就能钉住。
|
||||
func TestListSessionsForScanMatchesSelect(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ('alice', 'Alice', 'x', 'user')`); err != nil {
|
||||
t.Fatalf("seed user: %v", err)
|
||||
}
|
||||
|
||||
sid := seedSessionRow(t, "list-scan")
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`UPDATE sessions SET max_rounds = 7, used_rounds = 3 WHERE session_id = $1`,
|
||||
sid); err != nil {
|
||||
t.Fatalf("set budget: %v", err)
|
||||
}
|
||||
seedMailIn(t, sid, "alice", "opencode", "hello")
|
||||
|
||||
// 无过滤（管理员 all=true 走这条）
|
||||
all, err := ListSessionsFor(context.Background(), "", 50)
|
||||
if err != nil {
|
||||
t.Fatalf("列出全部会话失败: %v", err)
|
||||
}
|
||||
if len(all) != 1 {
|
||||
t.Fatalf("应有 1 个会话，实际 %d", len(all))
|
||||
}
|
||||
// 预算两列要真的读出来，不是零值
|
||||
if all[0].MaxRounds != 7 || all[0].UsedRounds != 3 {
|
||||
t.Errorf("预算未读出：max=%d used=%d（期望 7/3）",
|
||||
all[0].MaxRounds, all[0].UsedRounds)
|
||||
}
|
||||
if all[0].MailCount != 1 {
|
||||
t.Errorf("邮件数应为 1，实际 %d —— 列顺序可能错位", all[0].MailCount)
|
||||
}
|
||||
|
||||
// 带用户过滤（普通用户走这条，SQL 分支不同，要分别验）
|
||||
mine, err := ListSessionsFor(context.Background(), "alice", 50)
|
||||
if err != nil {
|
||||
t.Fatalf("列出自己的会话失败: %v", err)
|
||||
}
|
||||
if len(mine) != 1 {
|
||||
t.Fatalf("alice 参与过该会话，应能看到，实际 %d 个", len(mine))
|
||||
}
|
||||
if mine[0].MaxRounds != 7 || mine[0].MailCount != 1 {
|
||||
t.Errorf("过滤分支的列顺序错位：%+v", mine[0])
|
||||
}
|
||||
|
||||
// 与自己无关的人看不到
|
||||
other, err := ListSessionsFor(context.Background(), "bob", 50)
|
||||
if err != nil {
|
||||
t.Fatalf("列出 bob 的会话失败: %v", err)
|
||||
}
|
||||
if len(other) != 0 {
|
||||
t.Errorf("bob 未参与该会话，不该看到，实际 %d 个", len(other))
|
||||
}
|
||||
|
||||
// 归档会话不出现在列表里
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`UPDATE sessions SET status = 'archived' WHERE session_id = $1`, sid); err != nil {
|
||||
t.Fatalf("archive: %v", err)
|
||||
}
|
||||
after, err := ListSessionsFor(context.Background(), "", 50)
|
||||
if err != nil {
|
||||
t.Fatalf("归档后列出失败: %v", err)
|
||||
}
|
||||
if len(after) != 0 {
|
||||
t.Errorf("归档会话不该出现在列表里，实际 %d 个", len(after))
|
||||
}
|
||||
}
|
||||
|
||||
// 工作列表卡片视图需要「这条线索在干什么 / 还剩几个来回 / 最新进展是什么」，
|
||||
// 这些都从 ListContactsFor 一次取回 —— 否则卡片要为每条会话再打一次库。
|
||||
func TestListContactsForCardFields(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO users (username, display_name, password_hash, role)
|
||||
VALUES ('alice', 'Alice', 'x', 'user')`); err != nil {
|
||||
t.Fatalf("seed user: %v", err)
|
||||
}
|
||||
|
||||
sid := seedSessionRow(t, "card-fields")
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`UPDATE sessions SET subject = '缓存层选型评估', max_rounds = 5, used_rounds = 2
|
||||
WHERE session_id = $1`, sid); err != nil {
|
||||
t.Fatalf("set session: %v", err)
|
||||
}
|
||||
|
||||
// 三封：最早一封决定联系人身份，最后一封决定「最新进展」
|
||||
seedMailIn(t, sid, "alice", "opencode", "第一封")
|
||||
seedMailIn(t, sid, "opencode", "alice", "第二封")
|
||||
last := seedMailIn(t, sid, "opencode", "alice", "第三封")
|
||||
if _, err := db.DB.ExecContext(context.Background(),
|
||||
`UPDATE mails SET body = '已经跑完压测，Redis 方案在这个负载下明显更稳。'
|
||||
WHERE mail_id = $1`, last); err != nil {
|
||||
t.Fatalf("set body: %v", err)
|
||||
}
|
||||
|
||||
got, err := ListContactsFor(context.Background(), "alice", false)
|
||||
if err != nil {
|
||||
t.Fatalf("列出联系人失败: %v", err)
|
||||
}
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("应有 1 个联系人，实际 %d", len(got))
|
||||
}
|
||||
c := got[0]
|
||||
|
||||
if c.Subject != "缓存层选型评估" {
|
||||
t.Errorf("主题未带回：%q", c.Subject)
|
||||
}
|
||||
if c.MaxRounds != 5 || c.UsedRounds != 2 {
|
||||
t.Errorf("预算未带回：%d/%d（期望 5/2）", c.UsedRounds, c.MaxRounds)
|
||||
}
|
||||
// 最新进展取的是【最后】一封，不是第一封
|
||||
if c.LastFrom != "opencode" {
|
||||
t.Errorf("最新发件人应为 opencode，实际 %q", c.LastFrom)
|
||||
}
|
||||
if !strings.Contains(c.LastPreview, "Redis 方案") {
|
||||
t.Errorf("最新摘要应来自最后一封，实际 %q", c.LastPreview)
|
||||
}
|
||||
// 联系人身份仍取最早一封的对端
|
||||
if c.AgentName != "opencode" {
|
||||
t.Errorf("联系人应为 opencode，实际 %q", c.AgentName)
|
||||
}
|
||||
if c.MailCount != 3 {
|
||||
t.Errorf("邮件数应为 3，实际 %d —— 列顺序可能错位", c.MailCount)
|
||||
}
|
||||
if c.Address != "opencode@.card-fields" && !strings.HasSuffix(c.Address, ".card-fields") {
|
||||
t.Errorf("地址应带会话别名，实际 %q", c.Address)
|
||||
}
|
||||
}
|
||||
|
||||
// 摘要按字符截断，不按字节 —— 中文一字三字节，裸切会留半个字符。
|
||||
func TestPreviewRunes(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
n int
|
||||
want string
|
||||
}{
|
||||
{"短文本", 10, "短文本"},
|
||||
{" 两边有空白 ", 10, "两边有空白"},
|
||||
{"", 5, ""},
|
||||
{"一二三四五六", 3, "一二三…"},
|
||||
{"abcdefgh", 3, "abc…"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := previewRunes(c.in, c.n); got != c.want {
|
||||
t.Errorf("previewRunes(%q, %d) = %q，期望 %q", c.in, c.n, got, c.want)
|
||||
}
|
||||
}
|
||||
|
||||
// 截断结果必须是合法 UTF-8（不含替换字符）
|
||||
long := strings.Repeat("汉字", 200)
|
||||
got := previewRunes(long, 90)
|
||||
if strings.ContainsRune(got, '\uFFFD') {
|
||||
t.Error("截断产生了 U+FFFD，说明按字节切了")
|
||||
}
|
||||
if n := len([]rune(got)); n != 91 { // 90 + 省略号
|
||||
t.Errorf("截断后应为 90 字符 + 省略号，实际 %d 字符", n)
|
||||
}
|
||||
}
|
||||
|
||||
// 时间戳精度回归：SQLite 的 CURRENT_TIMESTAMP 只有秒，同秒插入的多行排序不确定，
|
||||
// 「会话里最早那封」（决定联系人身份）与「最后那封」（决定最新进展）都会取错。
|
||||
// NOW() 现在返回毫秒精度，且 mails 的 INSERT 显式传它 —— 这两点都要钉住。
|
||||
func TestMailTimestampSubSecond(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "ts-precision")
|
||||
|
||||
// 连续插 8 封（不显式给时间戳，走 CreateMail 里的 NOW()）
|
||||
ids := make([]uuid.UUID, 0, 8)
|
||||
for i := 0; i < 8; i++ {
|
||||
id, err := CreateMail(context.Background(), sid, nil,
|
||||
"alice", "", "opencode", "", fmt.Sprintf("第%d封", i), "body", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("创建邮件 %d 失败: %v", i, err)
|
||||
}
|
||||
ids = append(ids, id)
|
||||
}
|
||||
|
||||
// 至少要出现亚秒差异，否则说明 NOW() 又退回秒精度
|
||||
var distinct int
|
||||
if err := db.DB.QueryRowContext(context.Background(),
|
||||
`SELECT COUNT(DISTINCT created_at) FROM mails WHERE session_id = $1`,
|
||||
sid).Scan(&distinct); err != nil {
|
||||
t.Fatalf("统计不同时间戳失败: %v", err)
|
||||
}
|
||||
if distinct < 2 {
|
||||
var sample string
|
||||
db.DB.QueryRowContext(context.Background(),
|
||||
`SELECT CAST(created_at AS TEXT) FROM mails WHERE session_id = $1 LIMIT 1`,
|
||||
sid).Scan(&sample)
|
||||
t.Fatalf("8 封邮件只有 %d 个不同时间戳（样例 %q）—— NOW() 精度不足，"+
|
||||
"同秒邮件的先后顺序会由随机 UUID 决定", distinct, sample)
|
||||
}
|
||||
|
||||
// GetSessionMails 按时间升序，顺序必须与插入顺序一致
|
||||
got, err := GetSessionMails(context.Background(), sid)
|
||||
if err != nil {
|
||||
t.Fatalf("取会话邮件失败: %v", err)
|
||||
}
|
||||
if len(got) != len(ids) {
|
||||
t.Fatalf("应有 %d 封，实际 %d", len(ids), len(got))
|
||||
}
|
||||
for i, m := range got {
|
||||
if m.ID != ids[i] {
|
||||
t.Errorf("第 %d 封顺序错位：期望 %s，实际 %s（主题 %q）",
|
||||
i, ids[i], m.ID, m.Subject)
|
||||
}
|
||||
}
|
||||
}
|
||||
94
server/internal/repo/sessionrate.go
Normal file
94
server/internal/repo/sessionrate.go
Normal file
@ -0,0 +1,94 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// ---------- 新建会话速率限制 ----------
|
||||
//
|
||||
// Agent 可以用 name@path.new 开一串新会话，每条都是全新预算 ——
|
||||
// 速率限制只压住「短时间内暴开」这个滥用形态，过一个窗口自动恢复。
|
||||
|
||||
// ErrSessionRateLimited 表示该 Agent 短时间内新建会话过多。
|
||||
// （目前未使用，直接返回 retryAfter 由 handler 构造 429 响应）
|
||||
|
||||
const (
|
||||
sessionRateWindow = time.Hour
|
||||
sessionRateLimit = 20
|
||||
)
|
||||
|
||||
// AllowNewSession 供 handler 调用：Agent 新建会话前先过速率限制。
|
||||
// 人类用户不走这条路径（手工点「新建邮件」的频率天然受限）。
|
||||
// 返回 (allowed, retryAfter)。DB 不可用时放行。
|
||||
func AllowNewSession(ctx context.Context, agentName string) (bool, int) {
|
||||
if agentName == "" {
|
||||
return true, 0
|
||||
}
|
||||
return RateLimitCheckAndRecord(ctx, "session:"+agentName, sessionRateWindow, sessionRateLimit)
|
||||
}
|
||||
|
||||
// ReleaseNewSession 建会话失败后归还名额。
|
||||
// DB-backed 方式下记账在 AllowNewSession 里已完成，失败时需手动删除最近一条。
|
||||
func ReleaseNewSession(ctx context.Context, agentName string) {
|
||||
if agentName == "" {
|
||||
return
|
||||
}
|
||||
bucket := "session:" + agentName
|
||||
// 删掉最近一条（建会话失败，那次不该占名额）
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1 AND ts = (
|
||||
SELECT MAX(ts) FROM rate_limits WHERE bucket = $1
|
||||
)`, bucket)
|
||||
}
|
||||
|
||||
// SessionRateLimit 暴露窗口内的新建上限，供错误文案使用。
|
||||
func SessionRateLimit() int { return sessionRateLimit }
|
||||
|
||||
// ---------- Agent 建日历事件的速率限制 ----------
|
||||
//
|
||||
// 与新建会话同一套机制、独立的桶。为什么必须限：
|
||||
//
|
||||
// 日历事件是**长效**的 —— 一条每日重复的提醒会一直发下去，直到有人去删。
|
||||
// 模型在循环里每轮建一个「10 分钟后提醒我检查」，攒出几十条定时任务后，
|
||||
// 即使那条会话早已归档，提醒仍会按时发出。这比 `.new` 洪泛更难收拾：
|
||||
// 后者只是多几条空会话，前者是持续产生新邮件的源头。
|
||||
//
|
||||
// 上限与新建会话一致（20 次/小时）：正常用法下 Agent 一次任务里建
|
||||
// 一两条日程，20 条足够宽松；而循环失控时一小时内就会撞上限。
|
||||
const (
|
||||
calendarRateWindow = time.Hour
|
||||
calendarRateLimit = 20
|
||||
)
|
||||
|
||||
// AllowAgentCalendarEvent 供 handler 调用：Agent 建日历事件前先过速率限制。
|
||||
//
|
||||
// 人类不走这条路径（在界面上手工填表的频率天然受限），
|
||||
// 因此桶名带 agent: 前缀，与人类操作完全隔离。
|
||||
// 返回 (allowed, retryAfter)。DB 不可用时放行 —— 限速不该成为可用性的单点。
|
||||
func AllowAgentCalendarEvent(ctx context.Context, agentName string) (bool, int) {
|
||||
if agentName == "" {
|
||||
return true, 0
|
||||
}
|
||||
return RateLimitCheckAndRecord(ctx, "calendar:"+agentName, calendarRateWindow, calendarRateLimit)
|
||||
}
|
||||
|
||||
// ReleaseAgentCalendarEvent 建事件失败后归还名额。
|
||||
//
|
||||
// 与 ReleaseNewSession 同理：记账发生在检查那一刻，
|
||||
// 后续的写库失败意味着「那次创建实际没有发生」，不该占名额。
|
||||
func ReleaseAgentCalendarEvent(ctx context.Context, agentName string) {
|
||||
if agentName == "" {
|
||||
return
|
||||
}
|
||||
bucket := "calendar:" + agentName
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1 AND ts = (
|
||||
SELECT MAX(ts) FROM rate_limits WHERE bucket = $1
|
||||
)`, bucket)
|
||||
}
|
||||
|
||||
// CalendarRateLimit 暴露窗口内的上限，供错误文案使用。
|
||||
func CalendarRateLimit() int { return calendarRateLimit }
|
||||
127
server/internal/repo/sessionrate_test.go
Normal file
127
server/internal/repo/sessionrate_test.go
Normal file
@ -0,0 +1,127 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// 新建会话速率限制：DB 版
|
||||
//
|
||||
// 这些测试用真实的 SQLite（setupTestDB），验证速率限制的原子性与窗口滑动。
|
||||
// 原来的内存版测试依赖 sessionRateLimiter 结构体，替换为 DB 版后重写。
|
||||
|
||||
func TestSessionRateAllowsUpToLimit(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 1; i <= sessionRateLimit; i++ {
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatalf("第 %d 次应放行（上限 %d）", i, sessionRateLimit)
|
||||
}
|
||||
}
|
||||
ok, retry := AllowNewSession(ctx, "bot")
|
||||
if ok {
|
||||
t.Fatal("超过上限应拦下")
|
||||
}
|
||||
if retry < 1 {
|
||||
t.Fatalf("应给出正的重试等待秒数，实际 %d", retry)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionRateIsPerAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
AllowNewSession(ctx, "busy")
|
||||
}
|
||||
if ok, _ := AllowNewSession(ctx, "busy"); ok {
|
||||
t.Fatal("busy 应已被拦")
|
||||
}
|
||||
if ok, _ := AllowNewSession(ctx, "idle"); !ok {
|
||||
t.Fatal("另一个 Agent 不该被牵连")
|
||||
}
|
||||
}
|
||||
|
||||
// 并发请求不能把上限刷穿（判断与记账必须原子）
|
||||
func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
passed := 0
|
||||
|
||||
for i := 0; i < sessionRateLimit*4; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||||
mu.Lock()
|
||||
passed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if passed != sessionRateLimit {
|
||||
t.Fatalf("%d 并发下放行 %d 次，期望恰好 %d 次",
|
||||
sessionRateLimit*4, passed, sessionRateLimit)
|
||||
}
|
||||
}
|
||||
|
||||
// 建会话失败时要还名额
|
||||
func TestSessionRateRelease(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
AllowNewSession(ctx, "bot")
|
||||
}
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||||
t.Fatal("应已刷满")
|
||||
}
|
||||
ReleaseNewSession(ctx, "bot")
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatal("归还名额后应能再开一条")
|
||||
}
|
||||
}
|
||||
|
||||
// 人类不走限速（空 agentName）
|
||||
func TestAllowNewSessionSkipsHumans(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit*3; i++ {
|
||||
if ok, _ := AllowNewSession(ctx, ""); !ok {
|
||||
t.Fatal("人类不该被限速")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 窗口滑过后自动恢复（过期记录自动清理）
|
||||
func TestSessionRateWindowSlides(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 插入过期记录（1小时前）
|
||||
cutoff := time.Now().Add(-sessionRateWindow - time.Minute)
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`,
|
||||
"session:bot", cutoff)
|
||||
if err != nil {
|
||||
t.Fatalf("插入过期记录失败: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 窗口外的记录应被清理，此次应放行
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatal("窗口外的记录应被清掉，此次应放行")
|
||||
}
|
||||
}
|
||||
246
server/internal/repo/thread.go
Normal file
246
server/internal/repo/thread.go
Normal file
@ -0,0 +1,246 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 对话树。
|
||||
//
|
||||
// **不另建 tree_nodes 表**：`mails.parent_mail_id` 已经完整编码了树结构 ——
|
||||
// 回复指向来信，转发指向被转发的原件。再维护一张 tree_nodes 就是第二份真相，
|
||||
// 两处不一致时无法判断谁对。这里直接用递归 CTE 在 mails 上查。
|
||||
//
|
||||
// 树可以跨会话：转发把线索引到新会话，但 parent 仍指向原件。这正是「对话树」比
|
||||
// 「会话内平铺」更有价值的地方 —— 能看出一条线索分叉去了哪里。
|
||||
// 也正因如此，读取时必须按会话逐个鉴权（见 handler）：
|
||||
// A 转发给 B 之后，B 与 C 在新会话里的往来不能回流给 A。
|
||||
//
|
||||
// **从根展开，而不是从锚点展开**：曾经的实现是「锚点的祖先链 + 锚点的子树」，
|
||||
// 于是兄弟节点整条分支都在盲区里 —— 一封抄送给两个 Agent 的邮件，两个回复
|
||||
// 互为兄弟，从其中一个看树看不到另一个；挂在原件上的转发同理。
|
||||
// 兄弟既不是锚点的祖先也不是它的子孙，只有先上溯到根、再整棵 BFS 才能覆盖。
|
||||
//
|
||||
// **分块加载而非截断**：线索可以有几百封，一次全取要把几 MB 预览塞给前端。
|
||||
// 从根 BFS 后只剩一个方向，游标就是「已取到的节点数」。
|
||||
|
||||
// TreeMail 是树里的一个节点。正文只带预览：整棵线索带全文可能几百 KB，
|
||||
// 前端点开某封时再单取全文与附件清单。
|
||||
type TreeMail struct {
|
||||
models.Mail
|
||||
// Depth 是**距线索根**的层级：0 = 根，1 = 它的直接回复。
|
||||
// 从根展开后根一定在结果里，绝对深度因此总是可知的（早先按相对锚点算，
|
||||
// 是因为那时根可能还没取到）。
|
||||
Depth int `json:"depth"`
|
||||
AttachmentCount int `json:"attachment_count"`
|
||||
}
|
||||
|
||||
// descendantDepthCap 只是数据损坏时的兜底。
|
||||
//
|
||||
// parent_mail_id 正常不成环（新邮件只能指向已存在的旧邮件），但一旦被外部工具改坏，
|
||||
// 无上限的递归 CTE 会把进程拖死。取得足够大，正常数据碰不到。
|
||||
const descendantDepthCap = 10000
|
||||
|
||||
const threadCols = `m.mail_id, m.session_id, m.parent_mail_id,
|
||||
m.from_name, m.from_workspace, m.to_name, m.to_workspace,
|
||||
m.cc_list, m.subject, m.body, m.mail_type,
|
||||
COALESCE(m.permission_result,'') AS permission_result,
|
||||
m.status, m.created_at, s.session_alias, s.workspace,
|
||||
(SELECT COUNT(*) FROM attachments a WHERE a.mail_id = m.mail_id) AS attach_count,
|
||||
EXISTS (SELECT 1 FROM users u WHERE u.username = m.from_name) AS from_human,
|
||||
EXISTS (SELECT 1 FROM users u WHERE u.username = m.to_name) AS to_human`
|
||||
|
||||
// ThreadRootOf 沿 parent_mail_id 上溯到线索的根，返回根的 mail_id 与锚点到根的层数。
|
||||
//
|
||||
// 「根」= 链条最上面那封：parent_mail_id 为 NULL，或者指向一封已被删掉的邮件
|
||||
// （JOIN 断掉，递归自然停在这一层）。锚点自己没有父时返回它自己、depth 0。
|
||||
//
|
||||
// **不做可见性过滤**：不可见的中间段必须能穿过 —— 转发把线索引进别人的会话，
|
||||
// 再往上却可能仍是自己参与的往来。只返回 id 与层数，不泄露任何内容。
|
||||
func ThreadRootOf(ctx context.Context, anchorID uuid.UUID) (uuid.UUID, int, error) {
|
||||
var rootID uuid.UUID
|
||||
var lvl int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
WITH RECURSIVE up(mail_id, parent_mail_id, lvl) AS (
|
||||
SELECT mail_id, parent_mail_id, 0 FROM mails WHERE mail_id = $1
|
||||
UNION ALL
|
||||
SELECT m.mail_id, m.parent_mail_id, up.lvl + 1
|
||||
FROM mails m JOIN up ON m.mail_id = up.parent_mail_id
|
||||
WHERE up.lvl < $2
|
||||
)
|
||||
SELECT mail_id, lvl FROM up ORDER BY lvl DESC LIMIT 1
|
||||
`, anchorID, descendantDepthCap).Scan(&rootID, &lvl)
|
||||
if err != nil {
|
||||
return uuid.Nil, 0, err
|
||||
}
|
||||
return rootID, lvl, nil
|
||||
}
|
||||
|
||||
// AncestorsRaw 沿 parent_mail_id 上溯，取第 offset+1 .. offset+limit 层的祖先。
|
||||
// 层号 1 = 父，2 = 祖父；返回的 Depth 为负数（相对锚点）。
|
||||
//
|
||||
// 从根 BFS 之后这个函数只在一处还有用：巨型线索里锚点没落在 BFS 首页时，
|
||||
// 用它把「根到锚点」这条路径单独补齐，保证点开的那封一定看得见。
|
||||
// 调用方需要自己把负 depth 换算成绝对深度（锚点绝对深度由 ThreadRootOf 给出）。
|
||||
//
|
||||
// **不做可见性过滤**，理由同 ThreadRootOf。过滤放在 handler 层（那里知道调用者是谁）。
|
||||
//
|
||||
// 第二个返回值表示 offset+limit 层之上还有节点。
|
||||
func AncestorsRaw(ctx context.Context, anchorID uuid.UUID, offset, limit int) ([]TreeMail, bool, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
WITH RECURSIVE up(mail_id, parent_mail_id, lvl) AS (
|
||||
SELECT mail_id, parent_mail_id, 0 FROM mails WHERE mail_id = $1
|
||||
UNION ALL
|
||||
SELECT m.mail_id, m.parent_mail_id, up.lvl + 1
|
||||
FROM mails m JOIN up ON m.mail_id = up.parent_mail_id
|
||||
WHERE up.lvl < $2
|
||||
)
|
||||
SELECT `+threadCols+`, u.lvl
|
||||
FROM up u
|
||||
JOIN mails m ON m.mail_id = u.mail_id
|
||||
JOIN sessions s ON m.session_id = s.session_id
|
||||
WHERE u.lvl > $3
|
||||
ORDER BY u.lvl ASC
|
||||
`, anchorID, offset+limit+1, offset)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
// 多取一层用来判断「上面还有没有」，不返回给调用方
|
||||
out, err := scanTreeRows(rows, true)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasMore := len(out) > limit
|
||||
if hasMore {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, hasMore, nil
|
||||
}
|
||||
|
||||
// DescendantsRaw 取给定节点及其全部子孙，BFS 顺序（同层按时间），按节点数分页。
|
||||
//
|
||||
// 传线索的根（见 ThreadRootOf）就能覆盖整棵树：兄弟、抄送产生的平行回复、
|
||||
// 挂在原件上的转发分支，全都是根的子孙。offset = 0 时结果第一个是起点自己（Depth 0）。
|
||||
//
|
||||
// 同样不做可见性过滤：不可见的子节点下面可能挂着可见的孙节点
|
||||
// （别人把线索转走又转回来给我）。
|
||||
//
|
||||
// 注意 CTE 每次都会走完整棵子树，LIMIT 只截断输出。一条邮件线索通常几十封，
|
||||
// 这个代价可以接受；真出现巨型线索时再加物化。
|
||||
func DescendantsRaw(ctx context.Context, anchorID uuid.UUID, offset, limit int) ([]TreeMail, bool, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
WITH RECURSIVE down(mail_id, lvl) AS (
|
||||
SELECT mail_id, 0 FROM mails WHERE mail_id = $1
|
||||
UNION ALL
|
||||
SELECT m.mail_id, down.lvl + 1
|
||||
FROM mails m JOIN down ON m.parent_mail_id = down.mail_id
|
||||
WHERE down.lvl < $2
|
||||
)
|
||||
SELECT `+threadCols+`, d.lvl
|
||||
FROM down d
|
||||
JOIN mails m ON m.mail_id = d.mail_id
|
||||
JOIN sessions s ON m.session_id = s.session_id
|
||||
ORDER BY d.lvl ASC, m.created_at ASC, m.mail_id ASC
|
||||
LIMIT $3 OFFSET $4
|
||||
`, anchorID, descendantDepthCap, limit+1, offset)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
out, err := scanTreeRows(rows, false)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasMore := len(out) > limit
|
||||
if hasMore {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, hasMore, nil
|
||||
}
|
||||
|
||||
// TreeMailByID 取单封邮件的树节点形式，深度由调用方给定。
|
||||
//
|
||||
// 补齐「根 → 锚点」路径时用得上：AncestorsRaw 从父开始，不含锚点自己。
|
||||
// 同样不做可见性过滤，由 handler 负责。
|
||||
func TreeMailByID(ctx context.Context, id uuid.UUID, depth int) (*TreeMail, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT `+threadCols+`, $2
|
||||
FROM mails m
|
||||
JOIN sessions s ON m.session_id = s.session_id
|
||||
WHERE m.mail_id = $1
|
||||
`, id, depth)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out, err := scanTreeRows(rows, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, sql.ErrNoRows
|
||||
}
|
||||
return &out[0], nil
|
||||
}
|
||||
|
||||
// scanTreeRows 读出节点。negate 为真时把层号取负（祖先方向）。
|
||||
func scanTreeRows(rows interface {
|
||||
Next() bool
|
||||
Scan(...interface{}) error
|
||||
Err() error
|
||||
Close() error
|
||||
}, negate bool) ([]TreeMail, error) {
|
||||
defer rows.Close()
|
||||
|
||||
out := []TreeMail{}
|
||||
for rows.Next() {
|
||||
var t TreeMail
|
||||
var alias *string
|
||||
var ccJSON []byte
|
||||
var lvl int
|
||||
if err := rows.Scan(&t.ID, &t.SessionID, &t.ParentMailID,
|
||||
&t.FromName, &t.FromWorkspace, &t.ToName, &t.ToWorkspace,
|
||||
&ccJSON, &t.Subject, &t.Body, &t.MailType, &t.PermResult,
|
||||
&t.Status, &t.CreatedAt, &alias, &t.SessionWorkspace, &t.AttachmentCount,
|
||||
&t.FromHuman, &t.ToHuman, &lvl); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(ccJSON) > 0 {
|
||||
json.Unmarshal(ccJSON, &t.CCList)
|
||||
}
|
||||
if t.CCList == nil {
|
||||
t.CCList = []models.Address{}
|
||||
}
|
||||
if alias != nil {
|
||||
t.SessionAlias = *alias
|
||||
}
|
||||
if negate {
|
||||
t.Depth = -lvl
|
||||
} else {
|
||||
t.Depth = lvl
|
||||
}
|
||||
t.BodyPreview = preview(t.Body, 240)
|
||||
t.Body = "" // 树视图只要预览，全文按需单取
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// preview 按 UTF-8 边界截断正文。
|
||||
// 直接切字节会把多字节字符切成半个，前端渲染出 U+FFFD 替换符。
|
||||
func preview(s string, max int) string {
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
cut := max
|
||||
for cut > 0 && !utf8Start(s[cut]) {
|
||||
cut--
|
||||
}
|
||||
return s[:cut] + "..."
|
||||
}
|
||||
|
||||
// utf8Start 判断某字节是否为一个 UTF-8 序列的首字节
|
||||
func utf8Start(b byte) bool { return b&0xC0 != 0x80 }
|
||||
286
server/internal/repo/thread_test.go
Normal file
286
server/internal/repo/thread_test.go
Normal file
@ -0,0 +1,286 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
func TestPreviewTruncatesOnUTF8Boundary(t *testing.T) {
|
||||
// 「巡」是 3 字节；在 max=4 处切会切进第 2 个字符中间
|
||||
s := "巡检报告"
|
||||
got := preview(s, 4)
|
||||
if got != "巡..." {
|
||||
t.Fatalf("按 UTF-8 边界截断失败：%q", got)
|
||||
}
|
||||
for i, r := range got {
|
||||
if r == 0xFFFD {
|
||||
t.Fatalf("位置 %d 出现替换符，说明切在了字符中间", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreviewKeepsShortBodyIntact(t *testing.T) {
|
||||
if got := preview("短正文", 240); got != "短正文" {
|
||||
t.Fatalf("未超长却被改动：%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// seedReply 插一封回复：parent 指向来信。
|
||||
func seedReply(t *testing.T, sessionID uuid.UUID, parent uuid.UUID, from, to, subject string) uuid.UUID {
|
||||
t.Helper()
|
||||
var id uuid.UUID
|
||||
err := db.DB.QueryRowContext(context.Background(), `
|
||||
INSERT INTO mails (session_id, parent_mail_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, parent, from, to, subject, nextSeedTime()).Scan(&id)
|
||||
if err != nil {
|
||||
t.Fatalf("seed reply: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
// 线索根定位：整棵树从根展开，所以这一步错了后面全错。
|
||||
func TestThreadRootOf(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "root-of")
|
||||
|
||||
root := seedMailIn(t, sid, "admin", "dsh", "原件")
|
||||
mid := seedReply(t, sid, root, "dsh", "admin", "Re: 原件")
|
||||
leaf := seedReply(t, sid, mid, "admin", "dsh", "Re: Re: 原件")
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
from uuid.UUID
|
||||
depth int
|
||||
}{
|
||||
{"从根本身出发", root, 0},
|
||||
{"从中间一封出发", mid, 1},
|
||||
{"从叶子出发", leaf, 2},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
gotRoot, gotDepth, err := ThreadRootOf(context.Background(), c.from)
|
||||
if err != nil {
|
||||
t.Fatalf("ThreadRootOf: %v", err)
|
||||
}
|
||||
if gotRoot != root {
|
||||
t.Errorf("根定位错误：得到 %s，期望 %s", gotRoot, root)
|
||||
}
|
||||
if gotDepth != c.depth {
|
||||
t.Errorf("层数错误：得到 %d，期望 %d", gotDepth, c.depth)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 这个测试是对话树那次故障的回归：
|
||||
// 一封抄送给两个 Agent 的邮件收到两个回复，它们互为**兄弟**。
|
||||
// 旧实现从锚点分「祖先方向 + 子孙方向」两路展开，兄弟既不是锚点的祖先
|
||||
// 也不是它的子孙，于是整条分支在树里根本不出现。
|
||||
// 从线索根 BFS 之后，兄弟都是根的子孙，必须一次全出来。
|
||||
func TestDescendantsFromRootIncludesSiblings(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "siblings")
|
||||
|
||||
root := seedMailIn(t, sid, "admin", "dsh", "测试抄送")
|
||||
replyA := seedReply(t, sid, root, "dsh", "admin", "Re: 测试抄送")
|
||||
replyB := seedReply(t, sid, root, "opencode", "admin", "Re: 测试抄送")
|
||||
|
||||
// 从 replyA 出发定位根，再从根整树展开
|
||||
gotRoot, anchorDepth, err := ThreadRootOf(context.Background(), replyA)
|
||||
if err != nil {
|
||||
t.Fatalf("ThreadRootOf: %v", err)
|
||||
}
|
||||
if gotRoot != root || anchorDepth != 1 {
|
||||
t.Fatalf("根定位错误：root=%s depth=%d", gotRoot, anchorDepth)
|
||||
}
|
||||
|
||||
nodes, hasMore, err := DescendantsRaw(context.Background(), gotRoot, 0, 50)
|
||||
if err != nil {
|
||||
t.Fatalf("DescendantsRaw: %v", err)
|
||||
}
|
||||
if hasMore {
|
||||
t.Error("三封邮件不该报 hasMore")
|
||||
}
|
||||
|
||||
byID := map[uuid.UUID]TreeMail{}
|
||||
for _, n := range nodes {
|
||||
byID[n.ID] = n
|
||||
}
|
||||
for name, id := range map[string]uuid.UUID{"根": root, "回复A": replyA, "回复B": replyB} {
|
||||
if _, ok := byID[id]; !ok {
|
||||
t.Errorf("%s 不在树里 —— 兄弟分支又丢了", name)
|
||||
}
|
||||
}
|
||||
if byID[root].Depth != 0 {
|
||||
t.Errorf("根的深度应为 0，实际 %d", byID[root].Depth)
|
||||
}
|
||||
if byID[replyA].Depth != 1 || byID[replyB].Depth != 1 {
|
||||
t.Errorf("两个回复都应在深度 1：A=%d B=%d", byID[replyA].Depth, byID[replyB].Depth)
|
||||
}
|
||||
}
|
||||
|
||||
// 转发落在**另一个会话**里，但 parent 仍指向原件。
|
||||
// 树必须跨会话展开，否则「这条线索转发给谁了」就看不见了。
|
||||
func TestDescendantsCrossSession(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
srcSession := seedSessionRow(t, "fwd-src")
|
||||
dstSession := seedSessionRow(t, "fwd-dst")
|
||||
|
||||
root := seedMailIn(t, srcSession, "admin", "dsh", "原件")
|
||||
// 转发：新会话，parent 仍指原件
|
||||
fwd := seedReply(t, dstSession, root, "admin", "opencode", "Fwd: 原件")
|
||||
// 转发的下游回复，还在新会话里
|
||||
fwdReply := seedReply(t, dstSession, fwd, "opencode", "admin", "Re: Fwd: 原件")
|
||||
|
||||
nodes, _, err := DescendantsRaw(context.Background(), root, 0, 50)
|
||||
if err != nil {
|
||||
t.Fatalf("DescendantsRaw: %v", err)
|
||||
}
|
||||
found := map[uuid.UUID]int{}
|
||||
for _, n := range nodes {
|
||||
found[n.ID] = n.Depth
|
||||
}
|
||||
if _, ok := found[fwd]; !ok {
|
||||
t.Error("转发不在树里 —— 跨会话展开失效")
|
||||
}
|
||||
if _, ok := found[fwdReply]; !ok {
|
||||
t.Error("转发的下游回复不在树里")
|
||||
}
|
||||
if found[fwd] != 1 || found[fwdReply] != 2 {
|
||||
t.Errorf("跨会话深度错误：fwd=%d fwdReply=%d（期望 1/2）", found[fwd], found[fwdReply])
|
||||
}
|
||||
// 会话不同 → session_id 必须如实反映，否则前端无法标出「线索去了别的会话」
|
||||
for _, n := range nodes {
|
||||
if n.ID == fwd && n.SessionID != dstSession {
|
||||
t.Errorf("转发的 session_id 错误：%s，期望 %s", n.SessionID, dstSession)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 分页：BFS 顺序稳定，两页拼起来等于一次全取。
|
||||
func TestDescendantsPaginationStable(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "paging")
|
||||
|
||||
root := seedMailIn(t, sid, "admin", "dsh", "原件")
|
||||
for i := 0; i < 5; i++ {
|
||||
seedReply(t, sid, root, "dsh", "admin", "Re: 原件")
|
||||
}
|
||||
|
||||
full, hasMoreFull, err := DescendantsRaw(context.Background(), root, 0, 50)
|
||||
if err != nil {
|
||||
t.Fatalf("全取: %v", err)
|
||||
}
|
||||
if hasMoreFull {
|
||||
t.Error("6 封邮件一次取完不该报 hasMore")
|
||||
}
|
||||
if len(full) != 6 {
|
||||
t.Fatalf("应有 6 个节点，实际 %d", len(full))
|
||||
}
|
||||
|
||||
page1, hasMore1, err := DescendantsRaw(context.Background(), root, 0, 4)
|
||||
if err != nil {
|
||||
t.Fatalf("第一页: %v", err)
|
||||
}
|
||||
if !hasMore1 {
|
||||
t.Error("还有 2 封没取，hasMore 应为真")
|
||||
}
|
||||
page2, hasMore2, err := DescendantsRaw(context.Background(), root, 4, 4)
|
||||
if err != nil {
|
||||
t.Fatalf("第二页: %v", err)
|
||||
}
|
||||
if hasMore2 {
|
||||
t.Error("第二页已取完，hasMore 应为假")
|
||||
}
|
||||
|
||||
joined := append(append([]TreeMail{}, page1...), page2...)
|
||||
if len(joined) != len(full) {
|
||||
t.Fatalf("两页拼接 %d 个，全取 %d 个", len(joined), len(full))
|
||||
}
|
||||
for i := range full {
|
||||
if joined[i].ID != full[i].ID {
|
||||
t.Fatalf("第 %d 个节点顺序不一致：分页 %s，全取 %s —— BFS 顺序不稳定，"+
|
||||
"分页加载会重复或漏掉节点", i, joined[i].ID, full[i].ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TreeMailByID 是「锚点没落进 BFS 首页」时的回填手段，深度由调用方给。
|
||||
func TestTreeMailByIDUsesGivenDepth(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "by-id")
|
||||
id := seedMailIn(t, sid, "admin", "dsh", "某封")
|
||||
|
||||
got, err := TreeMailByID(context.Background(), id, 7)
|
||||
if err != nil {
|
||||
t.Fatalf("TreeMailByID: %v", err)
|
||||
}
|
||||
if got.ID != id {
|
||||
t.Errorf("取错了邮件：%s", got.ID)
|
||||
}
|
||||
if got.Depth != 7 {
|
||||
t.Errorf("深度应取调用方给的 7，实际 %d", got.Depth)
|
||||
}
|
||||
// 树视图只要预览，全文必须被清空 —— 否则整条线索会把几百 KB 正文塞给前端
|
||||
if got.Body != "" {
|
||||
t.Errorf("Body 应清空，实际 %q", got.Body)
|
||||
}
|
||||
if got.BodyPreview == "" {
|
||||
t.Error("BodyPreview 应有内容")
|
||||
}
|
||||
}
|
||||
|
||||
// 抄送列表必须原样带出来：树上两个兄弟节点为什么并列，
|
||||
// 唯一的解释就是父邮件抄送给了两个人。丢了 cc_list 前端就没法说明。
|
||||
func TestTreeCarriesCCList(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "cc-carry")
|
||||
|
||||
var root 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, 'admin', '', 'dsh', '', '抄送两人', 'body',
|
||||
'[{"name":"opencode","path":"/home","session":"new","raw":"opencode@/home.new"}]', $2)
|
||||
RETURNING mail_id
|
||||
`, sid, nextSeedTime()).Scan(&root)
|
||||
if err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
|
||||
nodes, _, err := DescendantsRaw(context.Background(), root, 0, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("DescendantsRaw: %v", err)
|
||||
}
|
||||
if len(nodes) != 1 {
|
||||
t.Fatalf("应有 1 个节点，实际 %d", len(nodes))
|
||||
}
|
||||
if len(nodes[0].CCList) != 1 {
|
||||
t.Fatalf("抄送应有 1 人，实际 %d —— cc_list 没带出来", len(nodes[0].CCList))
|
||||
}
|
||||
if nodes[0].CCList[0].Raw != "opencode@/home.new" {
|
||||
t.Errorf("抄送 raw 错误：%q", nodes[0].CCList[0].Raw)
|
||||
}
|
||||
}
|
||||
|
||||
// 无抄送时 cc_list 必须是空数组而不是 null：
|
||||
// Go 的 nil slice 会序列化成 null，前端 node.cc_list.length 直接抛异常。
|
||||
func TestTreeEmptyCCIsArrayNotNull(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
sid := seedSessionRow(t, "cc-empty")
|
||||
root := seedMailIn(t, sid, "admin", "dsh", "无抄送")
|
||||
|
||||
nodes, _, err := DescendantsRaw(context.Background(), root, 0, 10)
|
||||
if err != nil {
|
||||
t.Fatalf("DescendantsRaw: %v", err)
|
||||
}
|
||||
if nodes[0].CCList == nil {
|
||||
t.Error("cc_list 为 nil，会序列化成 null")
|
||||
}
|
||||
}
|
||||
625
server/internal/repo/users.go
Normal file
625
server/internal/repo/users.go
Normal file
@ -0,0 +1,625 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
const (
|
||||
bcryptCost = 12
|
||||
sessionTTL = 7 * 24 * time.Hour
|
||||
userSelectCols = `user_id, username, display_name, password_hash, role, status, created_at, last_login, allowed_agents, allowed_paths`
|
||||
)
|
||||
|
||||
var (
|
||||
ErrUserNotFound = errors.New("user not found")
|
||||
ErrBadCredentials = errors.New("invalid username or password")
|
||||
ErrUserDisabled = errors.New("user disabled")
|
||||
ErrNameTaken = errors.New("name already taken by an agent or user")
|
||||
ErrSessionInvalid = errors.New("session invalid or expired")
|
||||
ErrInvalidUsername = errors.New("username must be 2-64 chars of [a-z0-9._-]")
|
||||
ErrAlreadySetup = errors.New("system already initialized")
|
||||
)
|
||||
|
||||
// ---------- 命名空间校验 ----------
|
||||
|
||||
// 三维地址的 name 位由人类用户与 Agent 共用，因此必须全局唯一
|
||||
func nameTaken(ctx context.Context, name string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT (SELECT COUNT(*) FROM users WHERE username = $1)
|
||||
+ (SELECT COUNT(*) FROM agents WHERE agent_name = $1)
|
||||
`, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// AgentNameAvailable 供 Agent 注册前校验（不与人类用户重名）
|
||||
func AgentNameAvailable(ctx context.Context, agentName string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE username = $1`, agentName).Scan(&n)
|
||||
return n == 0, err
|
||||
}
|
||||
|
||||
func validUsername(name string) bool {
|
||||
if len(name) < 2 || len(name) > 64 {
|
||||
return false
|
||||
}
|
||||
for _, r := range name {
|
||||
ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '.' || r == '_' || r == '-'
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
// 保留字：human 是兼容别名，不能被真实用户占用
|
||||
return name != "human"
|
||||
}
|
||||
|
||||
// ---------- User CRUD ----------
|
||||
|
||||
func scanUser(row *sql.Row) (*models.User, error) {
|
||||
var u models.User
|
||||
var agentsJSON, pathsJSON []byte
|
||||
err := row.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash,
|
||||
&u.Role, &u.Status, &u.CreatedAt, &u.LastLogin, &agentsJSON, &pathsJSON)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrUserNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
u.AllowedAgents = decodeStrList(agentsJSON)
|
||||
u.AllowedPaths = decodeStrList(pathsJSON)
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
func decodeStrList(raw []byte) []string {
|
||||
out := []string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
if out == nil {
|
||||
out = []string{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func CreateUser(ctx context.Context, username, password, displayName, role string,
|
||||
allowedAgents, allowedPaths []string) (*models.User, error) {
|
||||
username = strings.ToLower(strings.TrimSpace(username))
|
||||
if !validUsername(username) {
|
||||
return nil, ErrInvalidUsername
|
||||
}
|
||||
if role != "admin" {
|
||||
role = "user"
|
||||
}
|
||||
if displayName == "" {
|
||||
displayName = username
|
||||
}
|
||||
|
||||
taken, err := nameTaken(ctx, username)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if taken {
|
||||
return nil, ErrNameTaken
|
||||
}
|
||||
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(password), bcryptCost)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
agentsJSON, _ := json.Marshal(normalizeList(allowedAgents))
|
||||
pathsJSON, _ := json.Marshal(normalizeList(allowedPaths))
|
||||
|
||||
row := db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO users (username, display_name, password_hash, role, allowed_agents, allowed_paths)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
RETURNING `+userSelectCols,
|
||||
username, displayName, string(hash), role, agentsJSON, pathsJSON)
|
||||
return scanUser(row)
|
||||
}
|
||||
|
||||
func GetUserByName(ctx context.Context, username string) (*models.User, error) {
|
||||
return scanUser(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users WHERE username = $1`,
|
||||
strings.ToLower(strings.TrimSpace(username))))
|
||||
}
|
||||
|
||||
func GetUserByID(ctx context.Context, id uuid.UUID) (*models.User, error) {
|
||||
return scanUser(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users WHERE user_id = $1`, id))
|
||||
}
|
||||
|
||||
func ListUsers(ctx context.Context) ([]models.User, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users ORDER BY created_at ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
users := []models.User{}
|
||||
for rows.Next() {
|
||||
var u models.User
|
||||
var agentsJSON, pathsJSON []byte
|
||||
if err := rows.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash,
|
||||
&u.Role, &u.Status, &u.CreatedAt, &u.LastLogin, &agentsJSON, &pathsJSON); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
u.AllowedAgents = decodeStrList(agentsJSON)
|
||||
u.AllowedPaths = decodeStrList(pathsJSON)
|
||||
users = append(users, u)
|
||||
}
|
||||
return users, nil
|
||||
}
|
||||
|
||||
// UserUpdate 描述一次用户更新；nil 字段表示不改
|
||||
type UserUpdate struct {
|
||||
DisplayName *string
|
||||
Role *string
|
||||
Status *string
|
||||
AllowedAgents *[]string
|
||||
AllowedPaths *[]string
|
||||
}
|
||||
|
||||
func UpdateUser(ctx context.Context, id uuid.UUID, up UserUpdate) (*models.User, error) {
|
||||
var agentsJSON, pathsJSON *string
|
||||
if up.AllowedAgents != nil {
|
||||
b, _ := json.Marshal(normalizeList(*up.AllowedAgents))
|
||||
s := string(b)
|
||||
agentsJSON = &s
|
||||
}
|
||||
if up.AllowedPaths != nil {
|
||||
b, _ := json.Marshal(normalizeList(*up.AllowedPaths))
|
||||
s := string(b)
|
||||
pathsJSON = &s
|
||||
}
|
||||
|
||||
row := db.DB.QueryRowContext(ctx, `
|
||||
UPDATE users SET
|
||||
display_name = COALESCE($2, display_name),
|
||||
role = COALESCE($3, role),
|
||||
status = COALESCE($4, status),
|
||||
allowed_agents = COALESCE($5`+db.JSONCast()+`, allowed_agents),
|
||||
allowed_paths = COALESCE($6`+db.JSONCast()+`, allowed_paths)
|
||||
WHERE user_id = $1
|
||||
RETURNING `+userSelectCols,
|
||||
id, up.DisplayName, up.Role, up.Status, agentsJSON, pathsJSON)
|
||||
return scanUser(row)
|
||||
}
|
||||
|
||||
// normalizeList 去空白、去空项、去重，保持顺序
|
||||
func normalizeList(in []string) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
seen := map[string]bool{}
|
||||
for _, s := range in {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" || seen[s] {
|
||||
continue
|
||||
}
|
||||
seen[s] = true
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func SetPassword(ctx context.Context, id uuid.UUID, newPassword string) error {
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcryptCost)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE users SET password_hash = $2 WHERE user_id = $1`, id, string(hash))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
// 改密后踢掉该用户所有会话
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE user_id = $1`, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func DisableUser(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE users SET status = 'disabled' WHERE user_id = $1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE user_id = $1`, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func CountAdmins(ctx context.Context) (int, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE role = 'admin' AND status = 'active'`).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// EnsureAdminUser 首次启动时创建默认管理员（幂等）
|
||||
func EnsureAdminUser(ctx context.Context, username, password string) (*models.User, bool, error) {
|
||||
if n, err := CountAdmins(ctx); err != nil {
|
||||
return nil, false, err
|
||||
} else if n > 0 {
|
||||
u, err := GetUserByName(ctx, username)
|
||||
if err != nil && !errors.Is(err, ErrUserNotFound) {
|
||||
return nil, false, err
|
||||
}
|
||||
return u, false, nil
|
||||
}
|
||||
u, err := CreateUser(ctx, username, password, "管理员", "admin", nil, nil)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return u, true, nil
|
||||
}
|
||||
|
||||
// ---------- 登录 / 会话令牌 ----------
|
||||
|
||||
func Authenticate(ctx context.Context, username, password string) (*models.User, error) {
|
||||
u, err := GetUserByName(ctx, username)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrUserNotFound) {
|
||||
// 统一错误，避免暴露用户是否存在
|
||||
return nil, ErrBadCredentials
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrUserDisabled
|
||||
}
|
||||
if bcrypt.CompareHashAndPassword([]byte(u.PasswordHash), []byte(password)) != nil {
|
||||
return nil, ErrBadCredentials
|
||||
}
|
||||
_, _ = db.DB.ExecContext(ctx, `UPDATE users SET last_login = NOW() WHERE user_id = $1`, u.ID)
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func newToken() (string, error) {
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
func CreateUserSession(ctx context.Context, userID uuid.UUID, userAgent string) (string, time.Time, error) {
|
||||
token, err := newToken()
|
||||
if err != nil {
|
||||
return "", time.Time{}, err
|
||||
}
|
||||
expires := time.Now().Add(sessionTTL)
|
||||
if len(userAgent) > 256 {
|
||||
userAgent = userAgent[:256]
|
||||
}
|
||||
_, err = db.DB.ExecContext(ctx, `
|
||||
INSERT INTO user_sessions (token, user_id, expires_at, user_agent)
|
||||
VALUES ($1, $2, $3, $4)`, token, userID, expires, userAgent)
|
||||
if err != nil {
|
||||
return "", time.Time{}, err
|
||||
}
|
||||
// 顺手清理过期令牌
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE expires_at < NOW()`)
|
||||
return token, expires, nil
|
||||
}
|
||||
|
||||
// ResolveUserSession 校验令牌并滑动续期
|
||||
func ResolveUserSession(ctx context.Context, token string) (*models.User, error) {
|
||||
if token == "" {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
var userID uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT user_id FROM user_sessions
|
||||
WHERE token = $1 AND expires_at > NOW()`, token).Scan(&userID)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
u, err := GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrUserDisabled
|
||||
}
|
||||
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`UPDATE user_sessions SET expires_at = $2 WHERE token = $1`,
|
||||
token, time.Now().Add(sessionTTL))
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func DeleteUserSession(ctx context.Context, token string) error {
|
||||
_, err := db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE token = $1`, token)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- 人类用户候选（供地址补全） ----------
|
||||
|
||||
func ListActiveUsernames(ctx context.Context) ([]string, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT username FROM users WHERE status = 'active' ORDER BY username`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []string{}
|
||||
for rows.Next() {
|
||||
var s string
|
||||
if err := rows.Scan(&s); err == nil {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ---------- 会话归属 ----------
|
||||
|
||||
func SetSessionOwner(ctx context.Context, sessionID, userID uuid.UUID) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET owner_user_id = $2 WHERE session_id = $1 AND owner_user_id IS NULL`,
|
||||
sessionID, userID)
|
||||
return err
|
||||
}
|
||||
|
||||
// SessionOwnerUsername 返回会话归属人类用户名；无归属时返回空串
|
||||
func SessionOwnerUsername(ctx context.Context, sessionID uuid.UUID) (string, error) {
|
||||
var name *string
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT u.username
|
||||
FROM sessions s LEFT JOIN users u ON u.user_id = s.owner_user_id
|
||||
WHERE s.session_id = $1`, sessionID).Scan(&name)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", fmt.Errorf("session %s not found", sessionID)
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if name == nil {
|
||||
return "", nil
|
||||
}
|
||||
return *name, nil
|
||||
}
|
||||
|
||||
// UserCanAccessSession 判断用户能否访问该会话：owner、或在邮件收发/抄送中出现，或 admin
|
||||
func UserCanAccessSession(ctx context.Context, u *models.User, sessionID uuid.UUID) (bool, error) {
|
||||
if u.IsAdmin() {
|
||||
return true, nil
|
||||
}
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*) FROM sessions s
|
||||
WHERE s.session_id = $1
|
||||
AND (s.owner_user_id = $2
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM mails m
|
||||
WHERE m.session_id = s.session_id
|
||||
AND (m.from_name = $3 OR m.to_name = $3
|
||||
OR `+db.CCHas("m.cc_list", 3)+`)
|
||||
))
|
||||
`, sessionID, u.ID, u.Username).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// RandomPassword 生成一个随机初始密码（首次启动无 ADMIN_PASSWORD 时使用）
|
||||
func RandomPassword(n int) string {
|
||||
const charset = "abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789"
|
||||
b := make([]byte, n)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "ChangeMe" + fmt.Sprint(time.Now().Unix())
|
||||
}
|
||||
for i := range b {
|
||||
b[i] = charset[int(b[i])%len(charset)]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// ---------- Setup（首次初始化管理员） ----------
|
||||
|
||||
// NeedsSetup 返回系统是否尚未初始化（没有任何用户）
|
||||
func NeedsSetup(ctx context.Context) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `SELECT COUNT(*) FROM users`).Scan(&n)
|
||||
return n == 0, err
|
||||
}
|
||||
|
||||
// SetupFirstAdmin 在系统尚无任何用户时创建首个管理员。
|
||||
// 已初始化时返回 ErrAlreadySetup，避免被用作后门。
|
||||
func SetupFirstAdmin(ctx context.Context, username, password, displayName string) (*models.User, error) {
|
||||
empty, err := NeedsSetup(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !empty {
|
||||
return nil, ErrAlreadySetup
|
||||
}
|
||||
if displayName == "" {
|
||||
displayName = username
|
||||
}
|
||||
return CreateUser(ctx, username, password, displayName, "admin", nil, nil)
|
||||
}
|
||||
|
||||
// ---------- 可选目录候选（供权限设置界面） ----------
|
||||
|
||||
// AllWorkspaceNames 汇总所有 Agent 注册过的工作区名，供管理员挑选可访问目录
|
||||
func AllWorkspaceNames(ctx context.Context) ([]string, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT DISTINCT ws->>'name' AS name
|
||||
FROM agents, jsonb_array_elements(workspaces) AS ws
|
||||
WHERE COALESCE(ws->>'name', '') <> ''
|
||||
ORDER BY name`)
|
||||
if err != nil {
|
||||
return []string{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []string{}
|
||||
for rows.Next() {
|
||||
var s string
|
||||
if err := rows.Scan(&s); err == nil {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// IsHumanUser 判断某个三维地址 name 位是否为人类用户
|
||||
func IsHumanUser(ctx context.Context, name string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE username = $1`, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// NearestHumanInThread 找出该会话上「最该为它点头的人」。
|
||||
//
|
||||
// 用途：权限询问的决策者是 Agent 时的救援路径。Agent 收不到 SendToUser
|
||||
// （那是人类的 SSE 通道），桥的 await 也就永不 resolve —— 会话永久阻塞。
|
||||
// 生产事故：pi 把任务派给自己的另一条会话，那条会话要跑 bash，
|
||||
// 权限邮件发给了 "pi" 自己，整条线索卡死，只能改数据库救回来。
|
||||
//
|
||||
// 三级查找，按「谁最了解这件事」排序：
|
||||
//
|
||||
// 1. 会话 owner —— 人在界面上开的会话，归属明确
|
||||
// 2. 最近一个往这条线索里**发过信**的人类 —— 派活的人
|
||||
// 3. 最近一个作为**收件人或抄送**出现的人类 —— 至少他知道这件事在进行
|
||||
//
|
||||
// 找不到时返回空串（不是错误）：调用方据此拒绝请求。这比转给一个对上下文
|
||||
// 一无所知的管理员好 —— 他既不知道这个 bash 命令在做什么，
|
||||
// 也不知道拒绝之后 Agent 该怎么绕过去。
|
||||
//
|
||||
// skipSelf 是发起询问的 Agent 名，永不作为决策者返回：它正是被卡住的那一方。
|
||||
// 名字与人类用户名共用命名空间，所以这里也顺手挡住「Agent 名恰好等于某人类名」
|
||||
// 这种配置错误。
|
||||
func NearestHumanInThread(ctx context.Context, sessionID uuid.UUID, skipSelf string) (string, error) {
|
||||
// 1. 会话 owner
|
||||
if owner, err := SessionOwnerUsername(ctx, sessionID); err == nil && owner != "" && owner != skipSelf {
|
||||
return owner, nil
|
||||
}
|
||||
|
||||
// 2/3. 扫这条会话的邮件。发件人优先于收件人/抄送方：
|
||||
// 发过信的人是主动参与者，被抄送的人可能只是旁观。
|
||||
//
|
||||
// 不用递归 CTE 沿 parent_mail_id 上溯：调用方只持有 session_id，
|
||||
// 没有触发询问的那封锚点邮件，所谓「链」的起点本来就得靠猜。
|
||||
// 而按会话扫还能覆盖分叉分支与断链（父邮件被删）的情形。
|
||||
//
|
||||
// **必须先把行读完再判定人类身份**：SQLite 连接池在测试与单文件库下
|
||||
// 常常只有一条连接，rows 未关闭时再发一条查询会自我死锁（实测挂死 60s）。
|
||||
type participants struct {
|
||||
from string
|
||||
others []string
|
||||
}
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT m.from_name, m.to_name, m.cc_list
|
||||
FROM mails m
|
||||
WHERE m.session_id = $1
|
||||
ORDER BY m.created_at DESC, m.mail_id DESC
|
||||
`, sessionID)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
var scanned []participants
|
||||
for rows.Next() {
|
||||
var from, to string
|
||||
var ccRaw []byte
|
||||
if err := rows.Scan(&from, &to, &ccRaw); err != nil {
|
||||
rows.Close()
|
||||
return "", err
|
||||
}
|
||||
scanned = append(scanned, participants{
|
||||
from: from,
|
||||
others: append([]string{to}, ccNames(ccRaw)...),
|
||||
})
|
||||
}
|
||||
err = rows.Err()
|
||||
rows.Close()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// 人类身份查询结果缓存：一条会话里同一个名字会出现很多次
|
||||
human := map[string]bool{}
|
||||
isHuman := func(name string) (bool, error) {
|
||||
if v, ok := human[name]; ok {
|
||||
return v, nil
|
||||
}
|
||||
v, err := IsHumanUser(ctx, name)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
human[name] = v
|
||||
return v, nil
|
||||
}
|
||||
|
||||
var fallback string // 收件人/抄送里的人类，仅在没有人类发件人时使用
|
||||
for _, p := range scanned {
|
||||
if p.from != skipSelf {
|
||||
if ok, err := isHuman(p.from); err != nil {
|
||||
return "", err
|
||||
} else if ok {
|
||||
return p.from, nil // 最近的人类发件人，直接定案
|
||||
}
|
||||
}
|
||||
if fallback != "" {
|
||||
continue
|
||||
}
|
||||
for _, cand := range p.others {
|
||||
if cand == "" || cand == skipSelf {
|
||||
continue
|
||||
}
|
||||
if ok, err := isHuman(cand); err != nil {
|
||||
return "", err
|
||||
} else if ok {
|
||||
fallback = cand
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return fallback, nil
|
||||
}
|
||||
|
||||
// ccNames 从 cc_list 的 JSON 里取出 name 位。
|
||||
//
|
||||
// 解析失败返回空切片而不是报错：抄送列表读不出来只该让这一封少几个候选，
|
||||
// 不该让整个决策者查找失败 —— 那会把「会话卡死」换成「权限请求 500」。
|
||||
func ccNames(raw []byte) []string {
|
||||
if len(raw) == 0 {
|
||||
return nil
|
||||
}
|
||||
var list []models.Address
|
||||
if err := json.Unmarshal(raw, &list); err != nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(list))
|
||||
for _, a := range list {
|
||||
if a.Name != "" {
|
||||
out = append(out, a.Name)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
195
server/internal/repo/usersession_time_test.go
Normal file
195
server/internal/repo/usersession_time_test.go
Normal file
@ -0,0 +1,195 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
/**
|
||||
* 会话令牌的时间比较。
|
||||
*
|
||||
* 这些用例锁的是一个真实的安全 bug：SQLite 驱动默认把 time.Time 写成
|
||||
* Go 的 t.String()（"2026-09-10 15:10:36.12 +0800 HKT m=+607182.88"），
|
||||
* 而 NOW() 返回 "2006-01-02 15:04:05.000000" 的 UTC 串。
|
||||
* 两者做字符串比较时 'T'/'+' 与数字的字典序、以及 +0800 与 UTC 的偏移
|
||||
* 双重错位，结果是：
|
||||
* - expires_at > NOW() 恒为真 → 令牌永不过期
|
||||
* - DELETE WHERE expires_at < NOW() 删 0 行 → 过期令牌永久堆积
|
||||
* 生产库实测两条 user_sessions 都是这个状态。
|
||||
*
|
||||
* 修法在 db.sqliteDSN 的 _time_format=sqlite&_timezone=UTC，
|
||||
* 所以这里必须用真实的 db.Connect 路径来验（setupTestDB 就是）。
|
||||
*/
|
||||
|
||||
func seedUserForSession(t *testing.T, ctx context.Context, username string) uuid.UUID {
|
||||
t.Helper()
|
||||
id := uuid.New()
|
||||
_, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO users (user_id, username, password_hash, display_name, role, status)
|
||||
VALUES ($1, $2, 'x', $2, 'user', 'active')`, id, username)
|
||||
if err != nil {
|
||||
t.Fatalf("seed user %s: %v", username, err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
// 存进去的时间必须能被 SQLite 的时间函数解析。
|
||||
// datetime() 返回 NULL 意味着所有 SQL 侧时间运算（过期判断、日历到点判断）全废。
|
||||
func TestTimeBindingIsParsableBySQLite(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
uid := seedUserForSession(t, ctx, "alice")
|
||||
|
||||
if _, _, err := CreateUserSession(ctx, uid, "probe"); err != nil {
|
||||
t.Fatalf("create session: %v", err)
|
||||
}
|
||||
|
||||
var bad int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM user_sessions WHERE datetime(expires_at) IS NULL`).Scan(&bad)
|
||||
if err != nil {
|
||||
t.Fatalf("probe: %v", err)
|
||||
}
|
||||
if bad != 0 {
|
||||
var raw string
|
||||
db.DB.QueryRowContext(ctx,
|
||||
`SELECT CAST(expires_at AS TEXT) FROM user_sessions LIMIT 1`).Scan(&raw)
|
||||
t.Fatalf("expires_at 无法被 datetime() 解析（存储为 %q）——"+
|
||||
"所有 SQL 侧时间比较都会静默失效", raw)
|
||||
}
|
||||
}
|
||||
|
||||
// 新令牌必须被认作有效，且滑动续期后仍然有效。
|
||||
func TestFreshSessionResolves(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
uid := seedUserForSession(t, ctx, "bob")
|
||||
|
||||
token, expires, err := CreateUserSession(ctx, uid, "ua")
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
if !expires.After(time.Now()) {
|
||||
t.Fatalf("新令牌的过期时间应在将来，得到 %v", expires)
|
||||
}
|
||||
|
||||
u, err := ResolveUserSession(ctx, token)
|
||||
if err != nil {
|
||||
t.Fatalf("resolve: %v", err)
|
||||
}
|
||||
if u.Username != "bob" {
|
||||
t.Fatalf("解析出的用户应为 bob，得到 %q", u.Username)
|
||||
}
|
||||
|
||||
// 滑动续期写回的也是 time.Time，格式同样必须正确
|
||||
var bad int
|
||||
db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM user_sessions WHERE datetime(expires_at) IS NULL`).Scan(&bad)
|
||||
if bad != 0 {
|
||||
t.Fatal("滑动续期写回的 expires_at 格式不可解析")
|
||||
}
|
||||
}
|
||||
|
||||
// 过期令牌必须被拒。这是 bug 的核心症状：格式错位时 expires_at > NOW()
|
||||
// 恒为真，一个 2020 年就过期的令牌照样能登录。
|
||||
func TestExpiredSessionRejected(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
uid := seedUserForSession(t, ctx, "carol")
|
||||
|
||||
token, _, err := CreateUserSession(ctx, uid, "ua")
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
|
||||
// 手工把过期时间推到过去（走同一条 time.Time 绑定路径）
|
||||
past := time.Now().Add(-1 * time.Hour)
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE user_sessions SET expires_at = $2 WHERE token = $1`, token, past); err != nil {
|
||||
t.Fatalf("backdate: %v", err)
|
||||
}
|
||||
|
||||
if _, err := ResolveUserSession(ctx, token); err == nil {
|
||||
var raw string
|
||||
db.DB.QueryRowContext(ctx,
|
||||
`SELECT CAST(expires_at AS TEXT) FROM user_sessions WHERE token=$1`, token).Scan(&raw)
|
||||
t.Fatalf("过期令牌竟然解析成功（expires_at=%q）—— 会话永不过期", raw)
|
||||
}
|
||||
}
|
||||
|
||||
// 顺手清理必须真的删掉过期行。删 0 行意味着 user_sessions 无限增长，
|
||||
// 且被盗令牌永远有效。
|
||||
func TestExpiredSessionsGetPurged(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
uid := seedUserForSession(t, ctx, "dave")
|
||||
|
||||
stale, _, err := CreateUserSession(ctx, uid, "ua")
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
past := time.Now().Add(-48 * time.Hour)
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE user_sessions SET expires_at = $2 WHERE token = $1`, stale, past); err != nil {
|
||||
t.Fatalf("backdate: %v", err)
|
||||
}
|
||||
|
||||
// 再建一个会话，CreateUserSession 内部会顺手清理过期令牌
|
||||
if _, _, err := CreateUserSession(ctx, uid, "ua2"); err != nil {
|
||||
t.Fatalf("second create: %v", err)
|
||||
}
|
||||
|
||||
var n int
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM user_sessions WHERE token = $1`, stale).Scan(&n); err != nil {
|
||||
t.Fatalf("count: %v", err)
|
||||
}
|
||||
if n != 0 {
|
||||
t.Fatal("过期令牌没被清理 —— DELETE ... WHERE expires_at < NOW() 匹配不到行")
|
||||
}
|
||||
|
||||
// 没过期的那条不能被误删
|
||||
var alive int
|
||||
db.DB.QueryRowContext(ctx, `SELECT COUNT(*) FROM user_sessions`).Scan(&alive)
|
||||
if alive != 1 {
|
||||
t.Fatalf("应剩 1 条有效会话，得到 %d", alive)
|
||||
}
|
||||
}
|
||||
|
||||
// 时间能原样读回，且亚秒精度不丢。
|
||||
//
|
||||
// 精度是有代价地保住的：_time_format=datetime 也能让 datetime() 正常工作，
|
||||
// 但它把时间截断到秒 —— 同秒插入的多封邮件排序就不确定，
|
||||
// 「会话里最早那封」（决定联系人身份）会取错行。
|
||||
func TestTimeRoundTripKeepsSubSecond(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO agents (agent_name, secret, platform) VALUES ('probe','x','test')`)
|
||||
if err != nil {
|
||||
t.Fatalf("seed agent: %v", err)
|
||||
}
|
||||
|
||||
want := time.Date(2026, 9, 3, 8, 45, 38, 123456000, time.UTC)
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET last_seen = $1 WHERE agent_name = 'probe'`, want); err != nil {
|
||||
t.Fatalf("update: %v", err)
|
||||
}
|
||||
|
||||
var got time.Time
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT last_seen FROM agents WHERE agent_name = 'probe'`).Scan(&got); err != nil {
|
||||
t.Fatalf("scan: %v", err)
|
||||
}
|
||||
if !got.UTC().Equal(want) {
|
||||
t.Fatalf("时间往返不一致：写入 %v，读回 %v", want, got.UTC())
|
||||
}
|
||||
if got.UTC().Nanosecond() == 0 {
|
||||
t.Fatal("亚秒精度被截断 —— 同秒插入的多行排序会不确定")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user