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