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"` // Status 是 agents.status:online / offline / disabled。 // 管理页靠它决定显示「停用」还是「恢复」。 Status string `json:"status"` } // 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) { // 带上 status:管理页靠它区分「在线 / 离线 / 已停用」并决定显示 // 「停用」还是「恢复」按钮。不过滤 disabled —— 这里是唯一能把已停用的 // Agent 恢复回来的地方,过滤掉就再也找不到它了。 rows, err := db.DB.QueryContext(ctx, `SELECT agent_name, COALESCE(default_rounds, 0), COALESCE(used_rounds, 0), COALESCE(status, 'offline') 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, &st.Status); 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) }