MailUI4Agents/gateway/internal/handler/agents.go

package handler

import (
	"net/http"

	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
)

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

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

// 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 err := Decode(r, &req); err != nil {
		Error(w, http.StatusBadRequest, "Invalid JSON")
		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 {
		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 上下文，
	// 让它在配额耗尽前主动发总结，而不是撞到 403 才发现。
	quota, qErr := repo.GetQuota(r.Context(), agentName)
	if qErr != nil {
		// 配额读不到不影响心跳本身，降级为不限额
		quota = repo.Quota{AgentName: agentName, Unlimited: true, Remaining: -1}
	}

	JSON(w, http.StatusOK, map[string]interface{}{
		"status":        "ok",
		"pending_mails": pending,
		"quota":         quota,
	})
}

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