MailUI4Agents/gateway/internal/repo/models_scope.go

package repo

import (
	"context"
	"strings"

	"github.com/agentmail/gateway/internal/db"
)

// ---------- 邮件场景下的可用模型 ----------
//
// 两张表，两种真相：
//
//	agent_model_catalog  —— 平台**上报**它当前看得见哪些模型（注册时整表替换）
//	agent_allowed_models —— 管理员**选定**其中哪些能在邮件场景下用，rank 即优先级
//
// 为什么不合成一张带 allowed 标记的表：模型会从平台目录里消失（换了 provider 配置、
// 上游临时下线），那时整行被删掉就连带把管理员的选择也删了，模型回来还得重配一遍。
// 分开存之后，「选了什么」是持久的，目录只决定「这一项现在是否可用」。
//
// 为什么让平台上报而不是在 Gateway 里配一张静态表：模型清单是平台侧的事实 ——
// opencode 的 provider 配置、DSH 的 llm 适配器注册，都可能随时变。
// Gateway 猜不出来，猜错的后果是管理员在配置页选了一个平台其实调不到的模型。

// ModelRef 是一次「provider + model」路由。
type ModelRef struct {
	Provider string `json:"provider"`
	Model    string `json:"model"`
}

// CatalogModel 是平台上报的一个可选模型。
type CatalogModel struct {
	Provider    string `json:"provider"`
	Model       string `json:"model"`
	DisplayName string `json:"display_name,omitempty"`
	// Allowed 表示它已被管理员选入邮件场景。
	// 与目录合并后一起返回，前端才能画出「已勾选」的复选框。
	Allowed bool `json:"allowed"`
	// Rank 仅在 Allowed 为真时有意义，越小越先试。
	Rank int `json:"rank,omitempty"`
}

// maxCatalogModels 限制单个 Agent 上报的模型数。
//
// 有平台会把上游的全部模型都列出来（实测 opencode 的一个 provider 就有几十个），
// 无上限的话一次注册能写进几千行，而配置页面上几千个复选框对人毫无用处。
const maxCatalogModels = 300

// ReplaceModelCatalog 整表替换某 Agent 上报的模型目录。
//
// 整表替换而非增量合并：目录是平台当前状态的快照，
// 增量合并会让已经下线的模型永远留在列表里，而那正是「选了却调不到」的来源。
//
// 事务包住删+插：中途失败留下一个空目录，会让配置页显示「该平台没有可用模型」
// 而管理员根本没做任何操作。
func ReplaceModelCatalog(ctx context.Context, agentName string, models []CatalogModel) error {
	agentName = strings.TrimSpace(agentName)
	if agentName == "" {
		return nil
	}
	if len(models) > maxCatalogModels {
		models = models[:maxCatalogModels]
	}

	tx, err := db.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx,
		`DELETE FROM agent_model_catalog WHERE agent_name = $1`, agentName); err != nil {
		return err
	}

	seen := map[string]bool{}
	for _, m := range models {
		p := strings.TrimSpace(m.Provider)
		id := strings.TrimSpace(m.Model)
		if p == "" || id == "" {
			continue // 半条记录不如不要：它在配置页上是一个点不动的空复选框
		}
		key := p + "/" + id
		if seen[key] {
			continue
		}
		seen[key] = true
		if _, err := tx.ExecContext(ctx,
			`INSERT INTO agent_model_catalog (agent_name, provider, model, display_name, reported_at)
			 VALUES ($1, $2, $3, $4, NOW())`,
			agentName, p, id, strings.TrimSpace(m.DisplayName)); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// ListModelCatalog 返回某 Agent 的模型目录，并标出哪些已被选入邮件场景。
//
// LEFT JOIN 而不是两次查询：前端要的是一份「带勾选状态的清单」，
// 在 SQL 里合完比让前端自己对齐两个数组更难出错。
func ListModelCatalog(ctx context.Context, agentName string) ([]CatalogModel, error) {
	rows, err := db.DB.QueryContext(ctx, `
		SELECT c.provider, c.model, c.display_name,
		       CASE WHEN a.model IS NULL THEN 0 ELSE 1 END AS allowed,
		       COALESCE(a.rank, 0)
		FROM agent_model_catalog c
		LEFT JOIN agent_allowed_models a
		       ON a.agent_name = c.agent_name
		      AND a.provider   = c.provider
		      AND a.model      = c.model
		WHERE c.agent_name = $1
		ORDER BY c.provider, c.model
	`, agentName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []CatalogModel{}
	for rows.Next() {
		var m CatalogModel
		var allowed int
		if err := rows.Scan(&m.Provider, &m.Model, &m.DisplayName, &allowed, &m.Rank); err != nil {
			return nil, err
		}
		m.Allowed = allowed == 1
		out = append(out, m)
	}
	return out, rows.Err()
}

// ListAllowedModels 按 rank 返回该 Agent 在邮件场景下可用的模型。
//
// **不与目录做 JOIN**：目录是平台上次注册时的快照，插件重启前可能已经过期。
// 真正能不能调通只有插件试过才知道 —— 这也正是插件要按顺序降级的原因。
// 在这里用目录过滤，只会把「目录暂时没上报但其实可用」的模型挡掉。
func ListAllowedModels(ctx context.Context, agentName string) ([]ModelRef, error) {
	rows, err := db.DB.QueryContext(ctx, `
		SELECT provider, model FROM agent_allowed_models
		WHERE agent_name = $1
		ORDER BY rank ASC, provider ASC, model ASC
	`, agentName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []ModelRef{}
	for rows.Next() {
		var m ModelRef
		if err := rows.Scan(&m.Provider, &m.Model); err != nil {
			return nil, err
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// SetAllowedModels 整表替换某 Agent 的邮件场景可用模型，入参顺序即优先级。
//
// 允许传空列表：那表示「不限定」——插件回退到平台自己的默认模型。
// 这与「一个都不许用」不同，后者等于让 Agent 彻底哑掉，不该是一次误删的后果。
func SetAllowedModels(ctx context.Context, agentName string, picks []ModelRef) error {
	agentName = strings.TrimSpace(agentName)
	if agentName == "" {
		return nil
	}

	tx, err := db.DB.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.ExecContext(ctx,
		`DELETE FROM agent_allowed_models WHERE agent_name = $1`, agentName); err != nil {
		return err
	}

	rank := 0
	seen := map[string]bool{}
	for _, m := range picks {
		p := strings.TrimSpace(m.Provider)
		id := strings.TrimSpace(m.Model)
		if p == "" || id == "" {
			continue
		}
		key := p + "/" + id
		if seen[key] {
			// 重复项直接跳过而不是报错：它对最终顺序没有影响，
			// 为一次无害的重复让整次保存失败只会让人以为配置没生效。
			continue
		}
		seen[key] = true
		if _, err := tx.ExecContext(ctx,
			`INSERT INTO agent_allowed_models (agent_name, provider, model, rank)
			 VALUES ($1, $2, $3, $4)`, agentName, p, id, rank); err != nil {
			return err
		}
		rank++
	}
	return tx.Commit()
}