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 }