MailUI4Agents/gateway/internal/handler/agents.go

package handler

import (
	"database/sql"
	"errors"
	"net/http"
	"strings"

	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/go-chi/chi/v5"
)

// ---------- Agent ----------

type registerRequest struct {
	Name       string             `json:"name"`
	Secret     string             `json:"secret"`
	Workspaces []models.Workspace `json:"workspaces"`
	Platform   string             `json:"platform"`
}

// heartbeatRequest 是心跳可选带的上报体。
//
// 字段全可省：旧插件发空心跳，不能因为新增了上报就把它们报错。
type heartbeatRequest struct {
	// PlatformSessions 是平台侧当前的会话快照（按最近活跃排序）。
	//
	// 为什么让插件上报而不是 Gateway 反向拉取：当前架构是单向的
	// （Agent 持密钥主动连 Gateway，Gateway 从不外呼）。反向拉取需要 Gateway
	// 保存各平台的地址与凭证，那是另一套信任模型。
	//
	// nil 与空数组语义不同：nil = 本次不上报（保留现有镜像），
	// 空数组 = 平台侧确实一条会话都没有（清空镜像）。
	// 拿不到会话列表的插件应当省略该字段，而不是传空数组把镜像抹掉。
	PlatformSessions []repo.PlatformSession `json:"platform_sessions"`

	// Models 是平台当前看得见的模型目录，供配置页勾选。
	//
	// 随心跳上报而不是只在注册时上报：模型清单会在运行中变
	// （换 provider 配置、上游上下线、换了 API key）。只在注册时报一次的话，
	// 目录会静静变陈，而管理员在配置页上看到的是上次重启时的快照 ——
	// 选中一个平台已经调不到的模型，失败要到真发邮件时才暴露。
	//
	// 与 PlatformSessions 同一约定：nil = 本次不上报（保留现有目录），
	// 空数组 = 平台确实一个模型都拿不到。拿不到目录时必须省略：
	// 清空目录会让配置页变成空白，管理员以为该平台没有任何可用模型。
	Models []repo.CatalogModel `json:"models"`
}

// POST /api/v1/agent/register
//
// 两种认证方式：
//   1. Authorization: Bearer <agent_key_token> —— 密钥认证（推荐）。
//      密钥未绑定时用本请求的 name 落定；已绑定时 name 必须与之一致，
//      否则等于拿别人的密钥冒充新身份。
//   2. body 里带 secret —— 旧方式，兼容保留。
func RegisterAgent(w http.ResponseWriter, r *http.Request) {
	var req registerRequest
	if !DecodeBody(w, r, &req) {
		return
	}
	if req.Name == "" {
		Error(w, http.StatusBadRequest, "Missing name")
		return
	}

	keyToken := middleware.BearerToken(r)
	if keyToken == "" && req.Secret == "" {
		Error(w, http.StatusBadRequest, "需要 Authorization: Bearer <密钥> 或 body 里的 secret")
		return
	}

	if keyToken != "" {
		bound, err := repo.VerifyAgentKey(r.Context(), keyToken)
		if err != nil {
			writeKeyErr(w, err)
			return
		}
		if bound != "" && bound != req.Name {
			Error(w, http.StatusForbidden,
				"该密钥已绑定到 Agent \""+bound+"\"，不能用于注册 \""+req.Name+"\"")
			return
		}
	}

	if req.Platform == "" {
		req.Platform = "pi"
	}

	// 三维地址的 name 位与人类用户名共用命名空间，不得重名
	if ok, err := repo.AgentNameAvailable(r.Context(), req.Name); err != nil {
		Error(w, http.StatusInternalServerError, "Failed to validate agent name")
		return
	} else if !ok {
		Error(w, http.StatusConflict, "该名称已被人类用户占用")
		return
	}
	if req.Name == "human" {
		Error(w, http.StatusBadRequest, "human 是保留别名，不能作为 Agent 名")
		return
	}

	// 密钥认证时不需要 secret，但 agents.secret 非空约束仍在；
	// 存密钥本身作占位，旧的 name/secret 路径不受影响。
	secret := req.Secret
	if secret == "" {
		secret = keyToken
	}

	if err := repo.CreateOrUpdateAgent(r.Context(), req.Name, secret, req.Platform, req.Workspaces); err != nil {
		// 已停用的 Agent 不得靠重新注册复活。回 403 而不是 500：
		// 这是一个明确的策略拒绝，插件应当停止重试并把原因打出来。
		if errors.Is(err, repo.ErrAgentDisabled) {
			Error(w, http.StatusForbidden,
				"Agent \""+req.Name+"\" 已被管理员停用，无法注册。"+
					"如需重新启用，请在管理页「默认预算」里恢复它。")
			return
		}
		Error(w, http.StatusInternalServerError, "Failed to register agent")
		return
	}

	// 待绑定密钥在首次注册成功后落定到该 Agent
	if keyToken != "" {
		if err := repo.ClaimAgentKey(r.Context(), keyToken, req.Name); err != nil {
			Error(w, http.StatusInternalServerError, "Failed to bind key")
			return
		}
	}

	JSON(w, http.StatusOK, map[string]string{
		"status":     "registered",
		"agent_name": req.Name,
	})
}

// POST /api/v1/agent/heartbeat
func HeartbeatAgent(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}

	pending, err := repo.HeartbeatAgent(r.Context(), agentName)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to heartbeat")
		return
	}

	// 可选的平台会话快照。解不开就当作没带：心跳的主职责是「我还活着」，
	// 不该因为上报体格式不对就把 Agent 判成离线。
	var req heartbeatRequest
	if r.ContentLength > 0 {
		_ = Decode(r, &req)
	}
	syncedSessions := -1 // -1 = 本次未上报
	if req.PlatformSessions != nil {
		if err := repo.ReplacePlatformSessions(r.Context(), agentName, req.PlatformSessions); err != nil {
			// 镜像写失败只影响候选补全，不影响投递，因此不报错
			syncedSessions = -1
		} else {
			syncedSessions = len(req.PlatformSessions)
		}
	}

	// 模型目录同理：写失败只让配置页看到的目录陈一轮，下一次心跳会补上。
	syncedModels := -1
	if req.Models != nil {
		if err := repo.ReplaceModelCatalog(r.Context(), agentName, req.Models); err == nil {
			syncedModels = len(req.Models)
		}
	}

	// 心跳回传该 Agent 的累计统计与新任务默认预算。
	//
	// 不再回传「剩余额度」：额度属于具体任务（会话）而不属于 Agent，
	// 剩余往返随每次发信响应（budget_remaining）回传，在那里才有意义。
	stats, sErr := repo.GetAgentStats(r.Context(), agentName)
	if sErr != nil {
		// 统计读不到不影响心跳本身
		stats = repo.AgentStats{AgentName: agentName}
	}

	resp := map[string]interface{}{
		"status":        "ok",
		"pending_mails": pending,
		"stats":         stats,
	}
	if syncedSessions >= 0 {
		resp["platform_sessions_synced"] = syncedSessions
	}
	if syncedModels >= 0 {
		resp["models_synced"] = syncedModels
	}
	// 回传当前生效的模型范围，插件无需另起一个请求去读。
	//
	// 随心跳回传而不是让插件自己轮询：管理员在配置页改了范围后，
	// 插件最多一个心跳周期（30 秒）就能看到新值，不需要重启。
	if allowed, aErr := repo.ListAllowedModels(r.Context(), agentName); aErr == nil {
		resp["allowed_models"] = allowed
		resp["models_unrestricted"] = len(allowed) == 0
	}
	JSON(w, http.StatusOK, resp)
}

// GET /api/v1/agents
func ListAgents(w http.ResponseWriter, r *http.Request) {
	statusFilter := r.URL.Query().Get("status")

	agents, err := repo.ListAgents(r.Context(), statusFilter)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list agents")
		return
	}

	JSON(w, http.StatusOK, map[string]interface{}{
		"agents": emptySlice(agents),
	})
}

// ---------- 停用 / 恢复 ----------

type setAgentStatusRequest struct {
	// Disabled true = 停用，false = 恢复
	Disabled bool `json:"disabled"`
}

// PUT /api/v1/admin/agents/{name}/status —— 停用或恢复一个 Agent
//
// 停用是可逆的「归档」，不是删除：
//   - 邮件、会话、权限记录、转发幂等键全部保留（往来里有一半是人自己写的）
//   - 从地址补全、GET /agents、可授权范围里消失
//   - 全部密钥被撤销，插件拿不到新任务也发不出信
//   - 重新注册会被拒（否则插件下次启动就把它复活了）
//
// 不提供彻底删除：Agent 名与人类用户名共用命名空间，删掉之后历史邮件的
// from_name 指向一个不存在的名字，此时有人注册同名 Agent（或人类账号），
// 那些旧邮件会看起来像是他发的。
func AdminSetAgentStatus(w http.ResponseWriter, r *http.Request) {
	name := strings.TrimSpace(chi.URLParam(r, "name"))
	if name == "" {
		Error(w, http.StatusBadRequest, "Missing agent name")
		return
	}

	var req setAgentStatusRequest
	if !DecodeBody(w, r, &req) {
		return
	}

	revoked, err := repo.SetAgentDisabled(r.Context(), name, req.Disabled)
	if errors.Is(err, sql.ErrNoRows) {
		Error(w, http.StatusNotFound, "Agent 不存在: "+name)
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to update agent status")
		return
	}

	resp := map[string]any{
		"agent_name": name,
		"disabled":   req.Disabled,
	}
	if req.Disabled {
		resp["keys_revoked"] = revoked
		resp["detail"] = "已停用。邮件与会话保留；该 Agent 的密钥已全部撤销，" +
			"恢复后需要重新签发。"
	} else {
		resp["detail"] = "已恢复为离线状态。需要重新签发密钥，插件连上后自动转为在线。"
	}
	JSON(w, http.StatusOK, resp)
}