311 lines
11 KiB
Go
311 lines
11 KiB
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"`
|
||
// 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)
|
||
}
|