MailUI4Agents/gateway/internal/handler/models_scope.go

package handler

import (
	"net/http"
	"strings"

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

// ---------- 邮件场景下的可用模型 ----------
//
//	GET /agent/models/allowed         读取被允许的模型（Agent 凭证）
//	GET /admin/agents/{name}/models   管理员读目录 + 已选
//	PUT /admin/agents/{name}/models   管理员保存选择与优先级
//
// **目录上报走心跳**（见 agents.go 的 heartbeatRequest.Models），不另设端点：
// 模型清单会在运行中变（换 provider 配置、上游上下线、换 API key），
// 心跳本来就是 30 秒一次的现成通道。另设一个 POST 等于给「目录是谁写的」
// 这个问题留两个答案，排查时要同时看两处。
//
// 生效的模型范围同样随心跳响应回传（allowed_models），因此插件通常不需要调
// 下面这个 GET —— 它是给非插件的第三方客户端（没有心跳循环）与排查用的。

// GET /api/v1/agent/models/allowed —— 插件读取被允许的模型
//
// 返回按优先级排序的列表。空列表表示**不限定**，插件应回退到平台自己的默认模型
// —— 与「一个都不许用」不同，后者等于让 Agent 彻底哑掉，不该是一次误配的后果。
func GetAllowedModels(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}
	models, err := repo.ListAllowedModels(r.Context(), agentName)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list allowed models")
		return
	}
	JSON(w, http.StatusOK, map[string]any{
		"models": models,
		// unrestricted 明确表达「没配 = 不限」，省得插件自己去判断空数组的含义
		"unrestricted": len(models) == 0,
	})
}

// GET /api/v1/admin/agents/{name}/models —— 管理员读目录（带已选标记）
func AdminListAgentModels(w http.ResponseWriter, r *http.Request) {
	name := strings.TrimSpace(chi.URLParam(r, "name"))
	if name == "" {
		Error(w, http.StatusBadRequest, "Missing agent name")
		return
	}
	catalog, err := repo.ListModelCatalog(r.Context(), name)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list model catalog")
		return
	}
	// 已选但已不在目录里的模型要单独给出来：平台可能临时下线了某个模型，
	// 界面上不显示的话管理员会以为自己没选过它，而它其实还在被插件尝试。
	stale, err := repo.ListStaleAllowedModels(r.Context(), name)
	if err != nil {
		stale = []repo.ModelRef{}
	}
	JSON(w, http.StatusOK, map[string]any{
		"agent_name": name,
		"catalog":    catalog,
		"stale":      stale,
	})
}

// PUT /api/v1/admin/agents/{name}/models —— 管理员保存选择
//
// 入参顺序即优先级（rank）。插件按这个顺序逐个尝试，全部失败才回一封失败邮件。
func AdminSetAgentModels(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 struct {
		Models []repo.ModelRef `json:"models"`
	}
	if err := Decode(r, &req); err != nil {
		Error(w, http.StatusBadRequest, "Invalid JSON")
		return
	}
	if len(req.Models) > maxAllowedModels {
		Error(w, http.StatusBadRequest,
			"选定的模型过多（上限 "+itoa(maxAllowedModels)+" 个）")
		return
	}

	if err := repo.SetAllowedModels(r.Context(), name, req.Models); err != nil {
		Error(w, http.StatusInternalServerError, "Failed to save allowed models")
		return
	}
	// 回传保存后的实际结果而不是回显入参：repo 层会跳过重复项与空字段，
	// 回显入参会让前端以为那些也存下来了。
	saved, err := repo.ListAllowedModels(r.Context(), name)
	if err != nil {
		saved = []repo.ModelRef{}
	}
	JSON(w, http.StatusOK, map[string]any{
		"status": "saved",
		"models": saved,
	})
}

// maxAllowedModels 限制管理员能选多少个模型。
//
// 降级尝试是串行的：选 50 个意味着最坏情况下一封邮件要等 50 次模型调用超时。
// 十个已经足够表达「主力 + 几个备选」。
const maxAllowedModels = 10

// itoa 避免为一个数字引入 strconv 导入（本文件只此一处用到）。
func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}