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「主动发信」的次数，不限制收信：
// 收信是被动的，卡住收信只会让邮件凭空消失；卡住发信才能真正阻止 Agent 无限自我循环。
//
// agents.max_rounds = 0 表示不限额。used_rounds 单调递增，由管理员显式重置。

// ErrQuotaExhausted 表示 Agent 的发信配额已用尽。
var ErrQuotaExhausted = errors.New("quota exhausted")

// Quota 是一个 Agent 的配额快照。
type Quota struct {
	AgentName string `json:"agent_name"`
	Max       int    `json:"max_rounds"`   // 0 = 不限
	Used      int    `json:"used_rounds"`
	Remaining int    `json:"remaining"`  // 不限时为 -1
	Unlimited bool   `json:"unlimited"`
}

func makeQuota(name string, max, used int) Quota {
	q := Quota{AgentName: name, Max: max, Used: used, Unlimited: max <= 0}
	if q.Unlimited {
		q.Remaining = -1
		return q
	}
	if r := max - used; r > 0 {
		q.Remaining = r
	}
	return q
}

// GetQuota 读取某 Agent 的配额状态。
func GetQuota(ctx context.Context, agentName string) (Quota, error) {
	var max, used int
	err := db.DB.QueryRowContext(ctx,
		`SELECT max_rounds, used_rounds FROM agents WHERE agent_name = $1`, agentName).Scan(&max, &used)
	if errors.Is(err, sql.ErrNoRows) {
		return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
	}
	if err != nil {
		return Quota{}, err
	}
	return makeQuota(agentName, max, used), nil
}

// ConsumeQuota 原子地占用一次发信配额，返回占用后的快照。
//
// 判断与自增必须在同一条 UPDATE 里完成（WHERE used_rounds < max_rounds），
// 否则并发发信会双双通过检查再各自 +1，把配额刷穿。
// 配额耗尽时返回 ErrQuotaExhausted，同时给出快照供调用方生成提示。
func ConsumeQuota(ctx context.Context, agentName string) (Quota, error) {
	tag, err := db.DB.ExecContext(ctx, `
		UPDATE agents SET used_rounds = used_rounds + 1
		WHERE agent_name = $1
		  AND (max_rounds <= 0 OR used_rounds < max_rounds)
	`, agentName)
	if err != nil {
		return Quota{}, err
	}
	n, _ := tag.RowsAffected()
	if n == 0 {
		// 要么 Agent 不存在，要么配额用尽——用快照区分
		q, qErr := GetQuota(ctx, agentName)
		if qErr != nil {
			return Quota{}, qErr
		}
		return q, ErrQuotaExhausted
	}
	return GetQuota(ctx, agentName)
}

// SetQuota 设置某 Agent 的配额上限（0 = 不限）。
func SetQuota(ctx context.Context, agentName string, max int) (Quota, error) {
	if max < 0 {
		max = 0
	}
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE agents SET max_rounds = $2 WHERE agent_name = $1`, agentName, max)
	if err != nil {
		return Quota{}, err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
	}
	return GetQuota(ctx, agentName)
}

// ResetQuota 把已用次数归零（配额上限不变）。
func ResetQuota(ctx context.Context, agentName string) (Quota, error) {
	tag, err := db.DB.ExecContext(ctx,
		`UPDATE agents SET used_rounds = 0 WHERE agent_name = $1`, agentName)
	if err != nil {
		return Quota{}, err
	}
	if n, _ := tag.RowsAffected(); n == 0 {
		return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
	}
	return GetQuota(ctx, agentName)
}

// ListQuotas 列出所有 Agent 的配额（管理员视图）。
func ListQuotas(ctx context.Context) ([]Quota, error) {
	rows, err := db.DB.QueryContext(ctx,
		`SELECT agent_name, max_rounds, used_rounds FROM agents ORDER BY agent_name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := []Quota{}
	for rows.Next() {
		var name string
		var max, used int
		if err := rows.Scan(&name, &max, &used); err != nil {
			return nil, err
		}
		out = append(out, makeQuota(name, max, used))
	}
	return out, rows.Err()
}

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

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