MailUI4Agents/server/internal/repo/agent_key_provision_test.go

package repo

/*
admin 预建 Agent 身份的判据（2026-10-04）。

★ 这一格存在的理由是一次真实的断裂：
  admin 调 `POST /admin/agent-keys` 给一个不存在的名字签 key → 签发成功，
  key 也正确绑定了 agent_name。但 `agents` 表里**没有那一行** ⇒
  `VerifyAgentKey` 返回空名字 ⇒ MCP 客户端连上来得到

      {"error":"密钥尚未绑定 Agent，请先调用 /agent/register 完成注册"}

  而它**永远不会**去调 /agent/register（MCP 宿主只按配置发 tools/list）。
  于是「admin 建身份 → 填进 MCP 配置 → 用」这条路是断的，
  而断点没有任何报错指向 admin 侧。
*/

import (
	"context"
	"testing"

	"github.com/agentmail/gateway/internal/db"
	"github.com/google/uuid"
)

func setupKeyProvisionDB(t *testing.T) {
	t.Helper()
	db.Close()
	db.DB = nil
	path := t.TempDir() + "/provision.db"
	if err := db.Connect(context.Background(), path); err != nil {
		t.Fatalf("连接测试库: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("迁移测试库: %v", err)
	}
	t.Cleanup(func() {
		db.Close()
		db.DB = nil
	})
}

// 签 key 之后 agents 行必须存在，且 key 能通过 AgentAuth 的那条查询。
func TestCreateAgentKeyAlsoCreatesAgentRow(t *testing.T) {
	setupKeyProvisionDB(t)
	ctx := context.Background()

	const name = "mcp-provision-test"
	if _, err := CreateAgentKey(ctx, name, "permanent", "MCP 接入端", 0, testUserID(t), ""); err != nil {
		t.Fatalf("签发 key: %v", err)
	}

	// agents 行必须存在
	var n int
	if err := db.DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM agents WHERE agent_name = $1`, name).Scan(&n); err != nil {
		t.Fatalf("查 agents: %v", err)
	}
	if n != 1 {
		t.Fatalf("★ 签了 key 却没建 agents 行 —— VerifyAgentKey 会返回空名字，" +
			"接入端拿到「密钥尚未绑定 Agent」而它永远不会去注册")
	}

	// 关键：VerifyAgentKey 必须能解出这个身份（这正是 MCP 连不上的那一步）
	agent, err := VerifyAgentKey(ctx, tokenOf(t, name))
	if err != nil {
		t.Fatalf("VerifyAgentKey: %v", err)
	}
	if agent == "" {
		t.Fatal("★ VerifyAgentKey 返回空身份 —— AgentAuth 会拒掉这把 key")
	}
	if agent != name {
		t.Errorf("身份应是 %s，实际 %s", name, agent)
	}
}

// 预建身份**不得**是 online：它此刻还没连上。
//
// 写成 online 的症状很具体：WebUI 联系人里出现一个「在线」却
// 永远不回信的幽灵，而排查会先怀疑「是不是它卡住了」。
func TestProvisionedAgentIsNotMarkedOnline(t *testing.T) {
	setupKeyProvisionDB(t)
	ctx := context.Background()

	const name = "mcp-offline-test"
	if _, err := CreateAgentKey(ctx, name, "permanent", "", 0, testUserID(t), ""); err != nil {
		t.Fatalf("签发 key: %v", err)
	}

	var status string
	if err := db.DB.QueryRowContext(ctx,
		`SELECT status FROM agents WHERE agent_name = $1`, name).Scan(&status); err != nil {
		t.Fatalf("查 status: %v", err)
	}
	if status == "online" {
		t.Error("★ 预建身份被写成 online —— 它还没连上，WebUI 会显示一个永远不回信的幽灵")
	}

	var lastSeen *string
	if err := db.DB.QueryRowContext(ctx,
		`SELECT CAST(last_seen AS TEXT) FROM agents WHERE agent_name = $1`, name).Scan(&lastSeen); err != nil {
		t.Fatalf("查 last_seen: %v", err)
	}
	if lastSeen != nil && *lastSeen != "" {
		t.Errorf("★ 预建身份不该有 last_seen（它还没握过手），实际 %q", *lastSeen)
	}
}

// 已有身份不能被 key 签发改坏（status/workspaces 属管理员与插件）。
func TestCreateAgentKeyDoesNotClobberExistingAgent(t *testing.T) {
	setupKeyProvisionDB(t)
	ctx := context.Background()

	const name = "existing-agent"
	if err := CreateOrUpdateAgent(ctx, name, "s3cret", "pi", nil); err != nil {
		t.Fatalf("建 agent: %v", err)
	}
	if _, err := CreateAgentKey(ctx, name, "permanent", "", 0, testUserID(t), ""); err != nil {
		t.Fatalf("再签一把 key: %v", err)
	}

	var secret, platform string
	if err := db.DB.QueryRowContext(ctx,
		`SELECT secret, platform FROM agents WHERE agent_name = $1`, name).Scan(&secret, &platform); err != nil {
		t.Fatalf("查 agent: %v", err)
	}
	if secret != "s3cret" || platform != "pi" {
		t.Errorf("★ 签 key 覆盖了已有 agent 的字段：secret=%q platform=%q", secret, platform)
	}
}

// testUserID 建一个 admin 用户（created_by 是外键）。
func testUserID(t *testing.T) uuid.UUID {
	t.Helper()
	var id uuid.UUID
	if err := db.DB.QueryRowContext(context.Background(),
		`INSERT INTO users (username, password_hash, role) VALUES ('provision-admin','x','admin')
		 RETURNING user_id`).Scan(&id); err != nil {
		t.Fatalf("建用户: %v", err)
	}
	return id
}

// tokenOf 取某 agent 名下任一把 key 的 token。
func tokenOf(t *testing.T, agentName string) string {
	t.Helper()
	var tok string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT key_token FROM agent_keys WHERE agent_name = $1 LIMIT 1`, agentName).Scan(&tok); err != nil {
		t.Fatalf("取 key: %v", err)
	}
	return tok
}