360 lines
13 KiB
Go
360 lines
13 KiB
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"`
|
||
|
||
// 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
|
||
}
|
||
|
||
// 可选的平台会话快照。解不开就当作没带：心跳的主职责是「我还活着」，
|
||
// 不该因为上报体格式不对就把 Agent 判成离线。
|
||
//
|
||
// 但**未知字段必须回报**（resp["unknown_fields"]）：这是全站唯一一处宽容
|
||
// 解码的端点，若还静默忽略，插件把 `models` 拼成 `modles` 就永远没人知道 ——
|
||
// 而那与 `attachments` vs `attachment_ids` 是同一种事故形状。
|
||
var req heartbeatRequest
|
||
var unknownFields []string
|
||
if r.ContentLength > 0 {
|
||
unknownFields, _ = DecodeLenient(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)
|
||
}
|
||
}
|
||
|
||
// 档位强制能力：省略时不动（保留现有值）。
|
||
// 写失败不影响心跳本身 —— 心跳的主职责是「我还活着」。
|
||
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{}{
|
||
"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": "已删除。邮件与会话保留（审计凭据）；" +
|
||
"密钥、平台会话镜像、模型范围已清除；它建的日历事件已置为取消。" +
|
||
"此名字今后不可再注册（历史邮件的署名由此不会被冒用）。",
|
||
})
|
||
}
|