MailUI4Agents/gateway/internal/repo/quota.go

package repo

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/models"
	"github.com/google/uuid"
)

// ---------- 配额 ----------
//
// **配额是任务的属性，不是 Agent 的属性。**
//
// 真正的约束在 `sessions.max_rounds`（见本文件末尾的「会话级往返预算」）：
// 每条会话独立计数，人在派活时给、在对话页里随时调。
//
// `agents` 表这边只剩两样东西：
//
//   default_rounds —— 派给这个 Agent 的**新任务**默认多少个来回。
//                     不同 Agent 能力不同（跑测试的小工具 vs 重构整个模块），
//                     默认值分开设才合理。
//
//   used_rounds    —— 纯统计，累计发信数。**不再拦任何请求。**
//                     它原本是「终身额度」：跑满就得管理员手工重置才能再干活，
//                     而 Agent 是长期在线的 —— 终身额度是错的工具。
//                     保留是因为「这个 Agent 一共发了多少信」本身有观测价值。
//
// 防止 Agent 用 `.new` 开一串新会话绕过预算，靠的是**新建会话速率限制**
// （见 sessionRateLimiter），而不是终身额度。
// AgentStats 是一个 Agent 的配额默认值与累计统计。
//
// 没有 Remaining / Unlimited 字段：这里不再有「剩余额度」的概念 ——
// 额度属于会话（SessionBudget），这里只有「新任务默认多少来回」与「一共发了多少信」。
type AgentStats struct {
	AgentName string `json:"agent_name"`
	// DefaultRounds 派给该 Agent 的新任务默认多少个来回（0 = 不限）
	DefaultRounds int `json:"default_rounds"`
	// SentTotal 累计发信数（纯统计，不拦请求）
	SentTotal int `json:"sent_total"`
	// ActiveSessions 该 Agent 参与的未归档会话数，配合默认值判断设多少合适
	ActiveSessions int `json:"active_sessions"`
}

// DefaultRoundsFor 读取该 Agent 的新任务默认预算。
//
// Agent 不存在时返回全局兜底值而非报错：派活的人不该因为「对方还没注册」
// 就拿不到一个合理的默认预算 —— 邮件本来就支持发给尚未上线的收件人。
func DefaultRoundsFor(ctx context.Context, agentName string) int {
	var n int
	err := db.DB.QueryRowContext(ctx,
		`SELECT COALESCE(default_rounds, 0) FROM agents WHERE agent_name = $1`,
		agentName).Scan(&n)
	if err != nil || n < 0 {
		return fallbackDefaultRounds
	}
	return n
}

// fallbackDefaultRounds 是 Agent 未注册时的兜底默认预算。
// 与建表默认值保持一致；改这里要同时改两份 schema。
const fallbackDefaultRounds = 20

// SetDefaultRounds 设置该 Agent 的新任务默认预算（0 = 不限）。
func SetDefaultRounds(ctx context.Context, agentName string, n int) (AgentStats, error) {
	if n < 0 {
		n = 0
	}
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE agents SET default_rounds = $2 WHERE agent_name = $1`, agentName, n)
	if err != nil {
		return AgentStats{}, err
	}
	if k, _ := tag.RowsAffected(); k == 0 {
		return AgentStats{}, fmt.Errorf("agent %q 不存在", agentName)
	}
	return GetAgentStats(ctx, agentName)
}

// GetAgentStats 读取某 Agent 的默认预算与累计统计。
func GetAgentStats(ctx context.Context, agentName string) (AgentStats, error) {
	var st AgentStats
	st.AgentName = agentName
	err := db.DB.QueryRowContext(ctx, `
		SELECT COALESCE(default_rounds, 0), COALESCE(used_rounds, 0)
		FROM agents WHERE agent_name = $1`, agentName,
	).Scan(&st.DefaultRounds, &st.SentTotal)
	if errors.Is(err, sql.ErrNoRows) {
		return AgentStats{}, fmt.Errorf("agent %q 不存在", agentName)
	}
	if err != nil {
		return AgentStats{}, err
	}
	st.ActiveSessions = countActiveSessionsFor(ctx, agentName)
	return st, nil
}

// countActiveSessionsFor 统计该 Agent 参与的未归档会话数。
// 查不出来返回 0：这只是个展示用的数字，不该让整个统计接口失败。
func countActiveSessionsFor(ctx context.Context, agentName string) int {
	var n int
	err := db.DB.QueryRowContext(ctx, `
		SELECT COUNT(DISTINCT s.session_id)
		FROM sessions s
		JOIN mails m ON m.session_id = s.session_id
		WHERE s.status <> 'archived'
		  AND (m.from_name = $1 OR m.to_name = $1 OR `+db.CCHas("m.cc_list", 1)+`)
	`, agentName).Scan(&n)
	if err != nil {
		return 0
	}
	return n
}

// BumpSentCount 累加发信统计。
//
// **绝不拦请求**：它是观测数据，不是额度。返回值只有 error，
// 而且调用方应当忽略它 —— 统计写失败不该让一封已经该发出的邮件失败。
func BumpSentCount(ctx context.Context, agentName string) {
	_, _ = db.DB.ExecContext(ctx,
		`UPDATE agents SET used_rounds = COALESCE(used_rounds, 0) + 1 WHERE agent_name = $1`,
		agentName)
}

// ListAgentStats 列出所有 Agent 的默认预算与统计（管理员视图）。
func ListAgentStats(ctx context.Context) ([]AgentStats, error) {
	rows, err := db.DB.QueryContext(ctx,
		`SELECT agent_name, COALESCE(default_rounds, 0), COALESCE(used_rounds, 0)
		 FROM agents ORDER BY agent_name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []AgentStats{}
	for rows.Next() {
		var st AgentStats
		if err := rows.Scan(&st.AgentName, &st.DefaultRounds, &st.SentTotal); err != nil {
			return nil, err
		}
		out = append(out, st)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	// 会话数逐个查：Agent 数量是个位数到几十，不值得为它写一个 GROUP BY 的联合查询
	for i := range out {
		out[i].ActiveSessions = countActiveSessionsFor(ctx, out[i].AgentName)
	}
	return out, nil
}

// ---------- 转发 ----------

// ForwardSource 是被转发邮件的必要信息。
type ForwardSource struct {
	Mail    *models.Mail
	Session uuid.UUID
}

// LoadForwardSource 读取待转发的邮件，并校验转发者确实参与过该邮件
//（收件人、发件人或被抄送方之一）。防止凭 mail_id 转发别人的邮件。
func LoadForwardSource(ctx context.Context, mailID uuid.UUID, actor string) (*models.Mail, error) {
	m, err := GetMailByID(ctx, mailID)
	if err != nil {
		return nil, ErrMailNotFound
	}

	if m.FromName == actor || m.ToName == actor {
		return m, nil
	}
	for _, cc := range m.CCList {
		if cc.Name == actor {
			return m, nil
		}
	}
	return nil, ErrForwardNotAllowed
}

var (
	// ErrMailNotFound 待转发的邮件不存在
	ErrMailNotFound = errors.New("mail not found")
	// ErrForwardNotAllowed 转发者未参与该邮件
	ErrForwardNotAllowed = errors.New("not a participant of that mail")
)

// ---------- 会话级往返预算 ----------
//
// 配额的真实语义是「这件事值得多少个来回」——那是**任务**的属性，不是 Agent 的属性。
// 只有 agents.max_rounds 一个全局计数器时有两个问题：
//   1. 两个并行任务互相抢额度：给紧急任务留的份被另一条线索吃掉
//   2. used_rounds 单调递增，跑满就得管理员手工重置才能再干活
// 因此预算下沉到会话，由人在写信时给、在对话页里随时调。
//
// **两层都要过**：会话预算 + Agent 全局配额。少了后者，Agent 自己 `.new` 开一串会话
// 每条都是全新预算，全局上限形同虚设；少了前者，就回到抢额度的老问题。

// SessionBudget 是一个会话的往返预算快照。
type SessionBudget struct {
	SessionID uuid.UUID `json:"session_id"`
	Max       int       `json:"max_rounds"` // 0 = 本会话不限
	Used      int       `json:"used_rounds"`
	Remaining int       `json:"remaining"` // 不限时为 -1
	Unlimited bool      `json:"unlimited"`
}

func makeSessionBudget(id uuid.UUID, max, used int) SessionBudget {
	b := SessionBudget{SessionID: id, Max: max, Used: used, Unlimited: max <= 0}
	if b.Unlimited {
		b.Remaining = -1
		return b
	}
	if r := max - used; r > 0 {
		b.Remaining = r
	}
	return b
}

// ErrSessionBudgetExhausted 表示该会话的往返预算已用尽。
var ErrSessionBudgetExhausted = errors.New("session budget exhausted")

// GetSessionBudget 读取会话预算。
func GetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
	var max, used int
	err := db.DB.QueryRowContext(ctx,
		`SELECT COALESCE(max_rounds,0), COALESCE(used_rounds,0) FROM sessions WHERE session_id = $1`,
		id).Scan(&max, &used)
	if errors.Is(err, sql.ErrNoRows) {
		return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
	}
	if err != nil {
		return SessionBudget{}, err
	}
	return makeSessionBudget(id, max, used), nil
}

// ConsumeSessionBudget 原子地占用会话的一次往返。
//
// 与 ConsumeQuota 同理：判断与自增必须在同一条 UPDATE 里（WHERE used_rounds < max_rounds），
// 否则并发发信会双双通过检查再各自 +1，把预算刷穿。
func ConsumeSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
	tag, err := db.DB.ExecContext(ctx, `
		UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) + 1
		WHERE session_id = $1
		  AND (COALESCE(max_rounds,0) <= 0 OR COALESCE(used_rounds,0) < max_rounds)
	`, id)
	if err != nil {
		return SessionBudget{}, err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		b, bErr := GetSessionBudget(ctx, id)
		if bErr != nil {
			return SessionBudget{}, bErr
		}
		return b, ErrSessionBudgetExhausted
	}
	return GetSessionBudget(ctx, id)
}

// SetSessionBudget 设置会话预算上限（0 = 不限）。
//
// 允许把上限调到低于已用次数：那表示「就到这里为止」，是人的合法意图，
// 不该因为算不出正的剩余量就拒绝。此时 Remaining 为 0，下次发信即被拦。
func SetSessionBudget(ctx context.Context, id uuid.UUID, max int) (SessionBudget, error) {
	if max < 0 {
		max = 0
	}
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE sessions SET max_rounds = $2, updated_at = NOW() WHERE session_id = $1`, id, max)
	if err != nil {
		return SessionBudget{}, err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
	}
	return GetSessionBudget(ctx, id)
}

// ResetSessionBudget 把该会话的已用次数归零（上限不变）。
func ResetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE sessions SET used_rounds = 0, updated_at = NOW() WHERE session_id = $1`, id)
	if err != nil {
		return SessionBudget{}, err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
	}
	return GetSessionBudget(ctx, id)
}

// RefundSessionBudget 退还一次往返。
//
// 会话预算先扣、Agent 全局配额后扣，全局那层拦下时必须把会话这次还回去，
// 否则会话预算白掉一格 —— 那次往返实际上没有发生。
func RefundSessionBudget(ctx context.Context, id uuid.UUID) {
	_, _ = db.DB.ExecContext(ctx,
		`UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) - 1
		 WHERE session_id = $1 AND COALESCE(used_rounds,0) > 0`, id)
}