MailUI4Agents/server/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 {
	// Workspace 是调用方当前所处的工作区（绝对路径）——**必需**。
	//
	// 心跳返回的 pending_mails 按它算：那是桥的补投判据，混了别的工作区
	// 就会让桥反复重放（见 handler 里那段说明）。
	Workspace string `json:"workspace"`
	// 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"`

	// ModeEnforcement 是插件自报的权限档位强制能力：native / advisory。
	//
	// 为什么走心跳而不是注册：能力会在运行中变。DSH 的沙箱模式被改成
	// danger-full-access 时，它就从 native 退化成了 advisory（实测：
	// approval:"never" 会在 waterfall 之前短路，approval/request 根本不触发）。
	// 只在注册时报一次的话，发件人看到的是上次重启时的能力快照。
	//
	// 与模型目录同一条通道（I-1：平台自己说的才算）。
	// 省略 = 本次不上报，保留现有值（与 PlatformSessions / Models 同约定）。
	ModeEnforcement string `json:"mode_enforcement"`
}

// 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
	}

	// 已退役的名字不可重建 —— 历史邮件的署名由此不会被冒用
	if retired, err := repo.IsRetiredAgentName(r.Context(), req.Name); err != nil {
		Error(w, http.StatusInternalServerError, "Failed to check agent name")
		return
	} else if retired {
		Error(w, http.StatusConflict, "该名字已退役，不可重建（历史邮件署名保护）")
		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
	}

	/*
		★ pending_mails 是**全局**未读数（跨工作区），而收件箱接口是 worker 级
		（要求 workspace）。两者不矛盾，是不同层的东西 —— 见 UnreadWorkspaces 注释。

		但只有总数，桥不知道去**哪个**工作区补投（它此前调的是不带收窄的
		`/mail/inbox`，于是会跨工作区重放）。⇒ 一并给出有未读的工作区清单，
		桥逐个去读。这也是"心跳返回什么才够用"的答案。
	*/
	pendingWorkspaces, _ := repo.UnreadWorkspaces(r.Context(), agentName)

	// 宽容解码（心跳是全站唯一一处容忍未知字段的端点，理由见下面的说明）。
	var req heartbeatRequest
	var unknownFields []string
	if r.ContentLength > 0 {
		unknownFields, _ = DecodeLenient(r, &req)
	}

	// 可选的平台会话快照。解不开就当作没带：心跳的主职责是「我还活着」，
	// 不该因为上报体格式不对就把 Agent 判成离线。
	//
	// 但**未知字段必须回报**（resp["unknown_fields"]）：这是全站唯一一处宽容
	// 解码的端点，若还静默忽略，插件把 `models` 拼成 `modles` 就永远没人知道 ——
	// 而那与 `attachments` vs `attachment_ids` 是同一种事故形状。
	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)
		}
	}

	// 档位强制能力：省略时不动（保留现有值）。
	// 写失败不影响心跳本身 —— 心跳的主职责是「我还活着」。
	if req.ModeEnforcement != "" {
		_ = repo.SetAgentModeEnforcement(r.Context(), agentName, req.ModeEnforcement)
	}

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

	resp := map[string]interface{}{
		"pending_workspaces": pendingWorkspaces,
		"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
	}
	// 未知字段回报：只有真的出现时才带这一项，正常心跳的响应不多一个空数组。
	// 插件看到它就知道自己上报的某个字段服务端根本没收。
	if len(unknownFields) > 0 {
		resp["unknown_fields"] = unknownFields
	}
	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
//
// 停用是可逆的「归档」，比 DELETE 轻一档：
//   - 邮件、会话、权限记录、转发幂等键全部保留（往来里有一半是人自己写的）
//   - 从地址补全、GET /agents、可授权范围里消失
//   - 全部密钥被撤销，插件拿不到新任务也发不出信
//   - 重新注册会被拒（否则插件下次启动就把它复活了）
//   - 别人发信给它得到 409（见 repo.RecipientDeliverable）
//
// 与 DELETE 的分工：停用留着运行态随时可恢复，删除清掉运行态且名字退役。
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 的密钥已全部撤销，" +
			"恢复后需要重新签发。停用期间别人发信给它会收到 409。"
	} else {
		resp["needs_new_key"] = true
		resp["detail"] = "已恢复为离线状态。停用时撤销的密钥不会自动回来 —— " +
			"必须在「密钥」面板重新签发一把并写进该插件的配置，" +
			"否则它会一直拿旧密钥重试并被拒（401）。"
	}
	JSON(w, http.StatusOK, resp)
}

// DELETE /api/v1/admin/agents/{name}
//
// 删除 Agent 的全部运行态，保留邮件历史。
//
// 取舍（见 PLUGIN-CONTRACT.md）：
//   - 邮件与会话不删（是审计凭据，且往来里有一半是人自己写的）
//   - 它建的日历事件置 cancelled（留着会由调度器一直触发，发信人却已不存在）
//   - 名字立即不可重建（Agent 名与人类用户名共用命名空间，
//     否则下一个同名注册者会看起来像是历史邮件的发信人）
func AdminDeleteAgent(w http.ResponseWriter, r *http.Request) {
	name := strings.TrimSpace(chi.URLParam(r, "name"))
	if name == "" {
		Error(w, http.StatusBadRequest, "Missing agent name")
		return
	}
	// 人类账号不走这个端点。硬编码某个用户名会在换管理员时失效，
	// 所以按「是不是人类用户」判定。
	if isHuman, hErr := repo.IsHumanUser(r.Context(), name); hErr != nil {
		Error(w, http.StatusInternalServerError, "Failed to check recipient")
		return
	} else if isHuman {
		Error(w, http.StatusForbidden,
			"\""+name+"\" 是人类用户，不能用这个端点删除（请去用户管理）")
		return
	}

	keysRevoked, err := repo.DeleteAgent(r.Context(), name)
	if errors.Is(err, sql.ErrNoRows) {
		Error(w, http.StatusNotFound, "Agent 不存在: "+name)
		return
	}
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to delete agent")
		return
	}

	JSON(w, http.StatusOK, map[string]any{
		"agent_name":   name,
		"keys_revoked": keysRevoked,
		"detail": "已删除。邮件与会话保留（审计凭据）；" +
			"密钥、平台会话镜像、模型范围已清除；它建的日历事件已置为取消。" +
			"此名字今后不可再注册（历史邮件的署名由此不会被冒用）。",
	})
}