package repo import ( "context" "time" "github.com/agentmail/gateway/internal/db" ) // ---------- 新建会话速率限制 ---------- // // Agent 可以用 name@path.new 开一串新会话,每条都是全新预算 —— // 速率限制只压住「短时间内暴开」这个滥用形态,过一个窗口自动恢复。 // ErrSessionRateLimited 表示该 Agent 短时间内新建会话过多。 // (目前未使用,直接返回 retryAfter 由 handler 构造 429 响应) const ( sessionRateWindow = time.Hour sessionRateLimit = 20 ) // AllowNewSession 供 handler 调用:Agent 新建会话前先过速率限制。 // 人类用户不走这条路径(手工点「新建邮件」的频率天然受限)。 // 返回 (allowed, retryAfter)。DB 不可用时放行。 func AllowNewSession(ctx context.Context, agentName string) (bool, int) { if agentName == "" { return true, 0 } return RateLimitCheckAndRecord(ctx, "session:"+agentName, sessionRateWindow, sessionRateLimit) } // ReleaseNewSession 建会话失败后归还名额。 // DB-backed 方式下记账在 AllowNewSession 里已完成,失败时需手动删除最近一条。 func ReleaseNewSession(ctx context.Context, agentName string) { if agentName == "" { return } bucket := "session:" + agentName // 删掉最近一条(建会话失败,那次不该占名额) _, _ = db.DB.ExecContext(ctx, `DELETE FROM rate_limits WHERE bucket = $1 AND ts = ( SELECT MAX(ts) FROM rate_limits WHERE bucket = $1 )`, bucket) } // SessionRateLimit 暴露窗口内的新建上限,供错误文案使用。 func SessionRateLimit() int { return sessionRateLimit } // ---------- Agent 建日历事件的速率限制 ---------- // // 与新建会话同一套机制、独立的桶。为什么必须限: // // 日历事件是**长效**的 —— 一条每日重复的提醒会一直发下去,直到有人去删。 // 模型在循环里每轮建一个「10 分钟后提醒我检查」,攒出几十条定时任务后, // 即使那条会话早已归档,提醒仍会按时发出。这比 `.new` 洪泛更难收拾: // 后者只是多几条空会话,前者是持续产生新邮件的源头。 // // 上限与新建会话一致(20 次/小时):正常用法下 Agent 一次任务里建 // 一两条日程,20 条足够宽松;而循环失控时一小时内就会撞上限。 const ( calendarRateWindow = time.Hour calendarRateLimit = 20 ) // AllowAgentCalendarEvent 供 handler 调用:Agent 建日历事件前先过速率限制。 // // 人类不走这条路径(在界面上手工填表的频率天然受限), // 因此桶名带 agent: 前缀,与人类操作完全隔离。 // 返回 (allowed, retryAfter)。DB 不可用时放行 —— 限速不该成为可用性的单点。 func AllowAgentCalendarEvent(ctx context.Context, agentName string) (bool, int) { if agentName == "" { return true, 0 } return RateLimitCheckAndRecord(ctx, "calendar:"+agentName, calendarRateWindow, calendarRateLimit) } // ReleaseAgentCalendarEvent 建事件失败后归还名额。 // // 与 ReleaseNewSession 同理:记账发生在检查那一刻, // 后续的写库失败意味着「那次创建实际没有发生」,不该占名额。 func ReleaseAgentCalendarEvent(ctx context.Context, agentName string) { if agentName == "" { return } bucket := "calendar:" + agentName _, _ = db.DB.ExecContext(ctx, `DELETE FROM rate_limits WHERE bucket = $1 AND ts = ( SELECT MAX(ts) FROM rate_limits WHERE bucket = $1 )`, bucket) } // CalendarRateLimit 暴露窗口内的上限,供错误文案使用。 func CalendarRateLimit() int { return calendarRateLimit }