package notify import ( "context" "encoding/json" "net/http" "net/http/httptest" "path/filepath" "strings" "testing" "github.com/agentmail/gateway/internal/db" "github.com/agentmail/gateway/internal/models" "github.com/agentmail/gateway/internal/repo" "github.com/agentmail/gateway/internal/sse" "github.com/google/uuid" ) 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) { 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) } } // attach 挂一个真实的 SSE 客户端并返回「读出这个 Agent 收到的 new_mail payload」的闭包。 // // 走真实的 sse.Default 而不是替换发送函数:要验的正是「谁收到什么」, // 而分发逻辑就在 Manager 里 —— 把它换掉等于不验。 func attach(t *testing.T, agentName string) func() map[string]any { t.Helper() rec := httptest.NewRecorder() req := httptest.NewRequest(http.MethodGet, "/api/v1/events/stream", nil) c := sse.Default.AddClient(rec, req, agentName, "") if c == nil { t.Fatalf("AddClient(%s) 返回 nil", agentName) } t.Cleanup(func() { sse.Default.RemoveClient(c.ID) }) return func() map[string]any { // SSE 帧形如 `id: N\nevent: new_mail\ndata: {…}\n\n` for _, frame := range strings.Split(rec.Body.String(), "\n\n") { if !strings.Contains(frame, "event: new_mail") { continue } for _, line := range strings.Split(frame, "\n") { if !strings.HasPrefix(line, "data: ") { continue } var m map[string]any if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &m); err == nil { return m } } } return nil } } // seedAdopted 建一条接管了 owner 的平台会话的本侧会话。 func seedAdopted(t *testing.T, owner, platformID, workspace string) uuid.UUID { t.Helper() ctx := context.Background() if err := repo.ReplacePlatformSessions(ctx, owner, []repo.PlatformSession{ {PlatformID: platformID, Workspace: workspace, Slug: "项目定位", Title: "项目定位"}, }); err != nil { t.Fatalf("ReplacePlatformSessions: %v", err) } id, err := repo.AdoptPlatformSession(ctx, owner, platformID, "项目定位", workspace, "项目定位") if err != nil { t.Fatalf("AdoptPlatformSession: %v", err) } return id } // platform_session_id 只该发给归属方。 // // 生产事故:会话接管了 pi 的 `01a05a5e-…`,而那封邮件抄送了 dsh。DSH 收到同一个 // id,在 `~/.dsh/sessions/` 里查不到(那是 `/root/.pi/agent/sessions/` 下的文件), // 于是按 N-8 抛「平台侧会话已删」——邮件静默消失,日志里一个字都没有。 func TestRecipients_PlatformIDOnlyToOwner(t *testing.T) { setupTestDB(t) seedAgent(t, "pi") seedAgent(t, "dsh") readPi, readDsh := attach(t, "pi"), attach(t, "dsh") sessionID := seedAdopted(t, "pi", "pid-pi-1", "/w") Recipients(context.Background(), Mail{ SessionID: sessionID, MailID: uuid.New(), From: "jianf", To: models.Address{Name: "pi", Path: "/w"}, CC: []models.Address{{Name: "dsh", Path: "/w"}}, Subject: "任务", }) pi, dsh := readPi(), readDsh() if pi == nil { t.Fatal("归属方 pi 没收到 new_mail") } if dsh == nil { t.Fatal("抄送方 dsh 没收到 new_mail(抄送方必须单独推)") } if v := pi["platform_session_id"]; v != "pid-pi-1" { t.Errorf("归属方 pi 的 platform_session_id = %v, want pid-pi-1", v) } if v := dsh["platform_session_id"]; v != "" { t.Errorf("抄送方 dsh 的 platform_session_id = %v, want 空串(那是 pi 的会话文件)", v) } } // 归属方在抄送位上也要拿到:归属与收件角色无关。 func TestRecipients_PlatformIDToOwnerEvenAsCC(t *testing.T) { setupTestDB(t) seedAgent(t, "pi") seedAgent(t, "dsh") readPi, readDsh := attach(t, "pi"), attach(t, "dsh") sessionID := seedAdopted(t, "pi", "pid-pi-2", "/w") Recipients(context.Background(), Mail{ SessionID: sessionID, MailID: uuid.New(), From: "jianf", To: models.Address{Name: "dsh", Path: "/w"}, CC: []models.Address{{Name: "pi", Path: "/w"}}, Subject: "任务", }) if v := readPi()["platform_session_id"]; v != "pid-pi-2" { t.Errorf("抄送位上的归属方 pi = %v, want pid-pi-2", v) } if v := readDsh()["platform_session_id"]; v != "" { t.Errorf("主收件人 dsh = %v, want 空串", v) } } // 普通(非接管)会话:谁都不该拿到 platform id。 func TestRecipients_PlainSessionNoPlatformID(t *testing.T) { setupTestDB(t) seedAgent(t, "pi") readPi := attach(t, "pi") id, err := repo.CreateSession(context.Background(), nil, "pi", "普通", "/w") if err != nil { t.Fatalf("CreateSession: %v", err) } Recipients(context.Background(), Mail{ SessionID: id, MailID: uuid.New(), From: "jianf", To: models.Address{Name: "pi", Path: "/w"}, Subject: "任务", }) if v := readPi()["platform_session_id"]; v != "" { t.Errorf("普通会话 = %v, want 空串", v) } }