MailUI4Agents/server/internal/repo/agent_disable_test.go

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
}