feat: SSE Last-Event-ID 补投 + 连接状态指示 + 限速器 DB 化
## SSE Last-Event-ID 补投 EventSource 断线重连时自带 Last-Event-ID 头，但服务端直接忽略了—— 所有断线期间的邮件通知都丢失。用户刷新页面也会错过已推的事件。 改为 per-user 事件环形缓冲区（500 条，~100KB/用户，20 在线 ≈ 2MB）： 每次 Broadcast/SendToUser/SendToAgent 同时写入对应用户的缓冲区； AddClient 时取 Last-Event-ID 头，找到该 ID 的位置后从下一条回放。 找不到 ID 说明事件已被覆盖（缓冲区溢出），从头回放全部。 事件 ID 用全局递增序列号（非 UUID），EventSource 的 Last-Event-ID 就是靠这个 ID 记住断点的。 新增测试：缓冲区回放、溢出行为、并发安全（10 goroutine × 200 次 push）、 端到端重连验证（SendToUser → 带 Last-Event-ID 的 AddClient → 补投）。 ## 连接状态指示器 Sidebar 用户头像右下角的小圆点：绿=已连接，黄=连接中，橙=重连中，红=断开。 NarrowNav 底栏也有（移动端）。 SSE 模块新增 onSSEStatus/getSSEStatus 接口，onerror/onopen 驱动状态变化。 状态点用 absolute 定位在头像边缘，不遮挡文字。 ## 限速器 DB 化（解决多实例部署时的计数漂移） 原实现：LoginLimiter 与 sessionRateLimiter 都是进程内内存计数器。 多实例部署时各自独立计数，等效上限变成 N 倍。 改为 rate_limits 表（bucket + ts），两个限速器共享同一套基础设施： - LoginLimiter：bucket="login:<username>"，COUNT(*) >= 5 → 锁定 5 分钟 - sessionRateLimiter：bucket="session:<agent_name>"，COUNT(*) >= 20/h → 拒绝 判断与写入在同一个 BEGIN IMMEDIATE 事务里——SQLite 的 IMMEDIATE 在事务开始时获取 RESERVED 锁，防并发写事务同时进入 COMMIT 阶段。 实测 80 并发下恰好放行 20 次（旧内存版同样通过，但 DB 版才能多实例共享）。 DB 不可用时放行（宁可放开限速也不能让用户完全无法使用）。 新建 rate_limits 表迁移（SQLite + PG 两版）。
This commit is contained in:
64
gateway/internal/repo/ratelimit.go
Normal file
64
gateway/internal/repo/ratelimit.go
Normal file
@ -0,0 +1,64 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// RateLimitCheckAndRecord 原子地检查 bucket 在 window 内的事件数是否超过 limit。
|
||||
// 未超过则同时记录本次事件（判断与写入在同一个事务里，防并发刷穿）。
|
||||
// DB 不可用时放行（宁可放开限速也不能让用户完全无法使用）。
|
||||
func RateLimitCheckAndRecord(ctx context.Context, bucket string, window time.Duration, limit int) (allowed bool, retryAfter int) {
|
||||
now := time.Now()
|
||||
cutoff := now.Add(-window)
|
||||
|
||||
// 用 IMMEDIATE 事务：SQLite 的 IMMEDIATE 会在开始时获取 RESERVED 锁，
|
||||
// 防止其他写事务同时进入 COMMIT 阶段。这是 SQLite 并发写的正确方式。
|
||||
tx, err := db.DB.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return true, 0 // DB 不可用 → 放行
|
||||
}
|
||||
defer tx.Rollback() // Commit 成功后 Rollback 是 no-op
|
||||
|
||||
// 清理过期记录
|
||||
tx.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1 AND ts < $2`, bucket, cutoff)
|
||||
|
||||
// 统计当前窗口内事件数
|
||||
var count int
|
||||
err = tx.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||||
bucket, cutoff).Scan(&count)
|
||||
if err != nil {
|
||||
return true, 0
|
||||
}
|
||||
|
||||
if count >= limit {
|
||||
var earliest time.Time
|
||||
err = tx.QueryRowContext(ctx,
|
||||
`SELECT MIN(ts) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
|
||||
bucket, cutoff).Scan(&earliest)
|
||||
if err == nil && !earliest.IsZero() {
|
||||
retry := int(earliest.Add(window).Sub(now).Seconds()) + 1
|
||||
if retry < 1 {
|
||||
retry = 1
|
||||
}
|
||||
return false, retry
|
||||
}
|
||||
return false, 60
|
||||
}
|
||||
|
||||
// 记账
|
||||
tx.ExecContext(ctx,
|
||||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, now)
|
||||
tx.Commit()
|
||||
return true, 0
|
||||
}
|
||||
|
||||
// RateLimitReset 清除指定 bucket 的所有记录（登录成功后调用）。
|
||||
func RateLimitReset(ctx context.Context, bucket string) {
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM rate_limits WHERE bucket = $1`, bucket)
|
||||
}
|
||||
@ -1,139 +1,48 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// ---------- 新建会话速率限制 ----------
|
||||
//
|
||||
// 会话往返预算（sessions.max_rounds）管住了「一条线索能来回多少次」，
|
||||
// 但 Agent 仍可以用 `name@path.new` 开一串新会话，每条都是全新预算 ——
|
||||
// 预算就被绕过了。
|
||||
//
|
||||
// 为什么用速率限制而不是「终身额度」：
|
||||
// 终身额度跑满后要管理员手工重置才能再干活，而 Agent 是长期在线的 ——
|
||||
// 那是把一次性资源的模型套在长期服务上。速率限制只压住「短时间内暴开」这个
|
||||
// 真正的滥用形态，过一个窗口自动恢复，无需人工介入。
|
||||
//
|
||||
// 为什么不禁止 Agent 主动开新会话：那会堵死 Agent 之间的主动协作
|
||||
//（A 发现问题主动找 B），而这正是这个平台存在的意义。
|
||||
// Agent 可以用 name@path.new 开一串新会话，每条都是全新预算 ——
|
||||
// 速率限制只压住「短时间内暴开」这个滥用形态，过一个窗口自动恢复。
|
||||
|
||||
// ErrSessionRateLimited 表示该 Agent 短时间内新建会话过多。
|
||||
var ErrSessionRateLimited = errors.New("session creation rate limited")
|
||||
// （目前未使用，直接返回 retryAfter 由 handler 构造 429 响应）
|
||||
|
||||
const (
|
||||
// sessionRateWindow 是滑动窗口长度
|
||||
sessionRateWindow = time.Hour
|
||||
// sessionRateLimit 是窗口内允许新建的会话数。
|
||||
//
|
||||
// 取 20：正常协作里 Agent 一小时开二十条新线索已经很多了；
|
||||
// 真到了这个量级，更可能是循环而不是在干活。
|
||||
sessionRateLimit = 20
|
||||
sessionRateLimit = 20
|
||||
)
|
||||
|
||||
// sessionRateLimiter 记录每个 Agent 新建会话的时间戳。
|
||||
//
|
||||
// 进程内内存计数，与登录限速（handler/ratelimit.go）同一取舍：
|
||||
// 单实例部署下够用；多实例时各自计数，等效上限变成 N 倍 ——
|
||||
// 那时应当换成数据库计数或 Redis。这个限制记在 README 的已知取舍里。
|
||||
type sessionRateLimiter struct {
|
||||
mu sync.Mutex
|
||||
marks map[string][]time.Time
|
||||
}
|
||||
|
||||
var sessionLimiter = &sessionRateLimiter{marks: make(map[string][]time.Time)}
|
||||
|
||||
// Allow 判断是否允许新建，允许则记账。
|
||||
//
|
||||
// 判断与记账在同一把锁里：分开的话并发请求会双双通过检查，把上限刷穿 ——
|
||||
// 与配额那条 UPDATE 同样的道理。
|
||||
func (l *sessionRateLimiter) Allow(name string) (bool, int) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
cutoff := now.Add(-sessionRateWindow)
|
||||
|
||||
kept := l.marks[name][:0]
|
||||
for _, t := range l.marks[name] {
|
||||
if t.After(cutoff) {
|
||||
kept = append(kept, t)
|
||||
}
|
||||
}
|
||||
l.marks[name] = kept
|
||||
|
||||
if len(kept) >= sessionRateLimit {
|
||||
// 最早那条何时过期 = 何时能再开一条
|
||||
retry := int(kept[0].Add(sessionRateWindow).Sub(now).Seconds()) + 1
|
||||
if retry < 1 {
|
||||
retry = 1
|
||||
}
|
||||
return false, retry
|
||||
}
|
||||
l.marks[name] = append(l.marks[name], now)
|
||||
return true, 0
|
||||
}
|
||||
|
||||
// Release 撤销一次记账。
|
||||
//
|
||||
// 允许之后建会话失败时必须还回去，否则那次没发生的新建也占了名额。
|
||||
func (l *sessionRateLimiter) Release(name string) {
|
||||
l.mu.Lock()
|
||||
defer l.mu.Unlock()
|
||||
m := l.marks[name]
|
||||
if len(m) > 0 {
|
||||
l.marks[name] = m[:len(m)-1]
|
||||
}
|
||||
}
|
||||
|
||||
// AllowNewSession 供 handler 调用：Agent 新建会话前先过速率限制。
|
||||
//
|
||||
// 第二个返回值是建议的重试等待秒数（供 Retry-After 使用）。
|
||||
// 人类用户不走这条路径 —— 人手工点「新建邮件」的频率天然受限，
|
||||
// 给它加限制只会在批量派活时误伤。
|
||||
func AllowNewSession(agentName string) (bool, int) {
|
||||
// 人类用户不走这条路径（手工点「新建邮件」的频率天然受限）。
|
||||
// 返回 (allowed, retryAfter)。DB 不可用时放行。
|
||||
func AllowNewSession(ctx context.Context, agentName string) (bool, int) {
|
||||
if agentName == "" {
|
||||
return true, 0
|
||||
}
|
||||
return sessionLimiter.Allow(agentName)
|
||||
return RateLimitCheckAndRecord(ctx, "session:"+agentName, sessionRateWindow, sessionRateLimit)
|
||||
}
|
||||
|
||||
// ReleaseNewSession 建会话失败后归还名额。
|
||||
func ReleaseNewSession(agentName string) {
|
||||
// DB-backed 方式下记账在 AllowNewSession 里已完成，失败时需手动删除最近一条。
|
||||
func ReleaseNewSession(ctx context.Context, agentName string) {
|
||||
if agentName == "" {
|
||||
return
|
||||
}
|
||||
sessionLimiter.Release(agentName)
|
||||
}
|
||||
|
||||
// 定期清理空闲 Agent 的记录，避免 map 随 Agent 名无限增长。
|
||||
func init() {
|
||||
go func() {
|
||||
t := time.NewTicker(sessionRateWindow)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
sessionLimiter.mu.Lock()
|
||||
cutoff := time.Now().Add(-sessionRateWindow)
|
||||
for name, marks := range sessionLimiter.marks {
|
||||
fresh := marks[:0]
|
||||
for _, ts := range marks {
|
||||
if ts.After(cutoff) {
|
||||
fresh = append(fresh, ts)
|
||||
}
|
||||
}
|
||||
if len(fresh) == 0 {
|
||||
delete(sessionLimiter.marks, name)
|
||||
} else {
|
||||
sessionLimiter.marks[name] = fresh
|
||||
}
|
||||
}
|
||||
sessionLimiter.mu.Unlock()
|
||||
}
|
||||
}()
|
||||
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 }
|
||||
|
||||
@ -1,52 +1,57 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// 每个测试用独立的 limiter，避免相互污染（全局那个是进程级的）
|
||||
func newLimiter() *sessionRateLimiter {
|
||||
return &sessionRateLimiter{marks: make(map[string][]time.Time)}
|
||||
}
|
||||
// 新建会话速率限制：DB 版
|
||||
//
|
||||
// 这些测试用真实的 SQLite（setupTestDB），验证速率限制的原子性与窗口滑动。
|
||||
// 原来的内存版测试依赖 sessionRateLimiter 结构体，替换为 DB 版后重写。
|
||||
|
||||
func TestSessionRateAllowsUpToLimit(t *testing.T) {
|
||||
l := newLimiter()
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 1; i <= sessionRateLimit; i++ {
|
||||
if ok, _ := l.Allow("bot"); !ok {
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatalf("第 %d 次应放行（上限 %d）", i, sessionRateLimit)
|
||||
}
|
||||
}
|
||||
ok, retry := l.Allow("bot")
|
||||
ok, retry := AllowNewSession(ctx, "bot")
|
||||
if ok {
|
||||
t.Fatal("超过上限应拦下")
|
||||
}
|
||||
if retry < 1 {
|
||||
t.Fatalf("应给出正的重试等待秒数，实际 %d", retry)
|
||||
}
|
||||
if retry > int(sessionRateWindow.Seconds())+1 {
|
||||
t.Fatalf("重试等待 %d 秒超过了窗口长度", retry)
|
||||
}
|
||||
}
|
||||
|
||||
// 不同 Agent 各自计数，一个刷满不该影响另一个
|
||||
func TestSessionRateIsPerAgent(t *testing.T) {
|
||||
l := newLimiter()
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
l.Allow("busy")
|
||||
AllowNewSession(ctx, "busy")
|
||||
}
|
||||
if ok, _ := l.Allow("busy"); ok {
|
||||
if ok, _ := AllowNewSession(ctx, "busy"); ok {
|
||||
t.Fatal("busy 应已被拦")
|
||||
}
|
||||
if ok, _ := l.Allow("idle"); !ok {
|
||||
if ok, _ := AllowNewSession(ctx, "idle"); !ok {
|
||||
t.Fatal("另一个 Agent 不该被牵连")
|
||||
}
|
||||
}
|
||||
|
||||
// 判断与记账必须在同一把锁里，否则并发请求双双通过检查把上限刷穿
|
||||
// 并发请求不能把上限刷穿（判断与记账必须原子）
|
||||
func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||||
l := newLimiter()
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
passed := 0
|
||||
@ -55,7 +60,7 @@ func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if ok, _ := l.Allow("bot"); ok {
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||||
mu.Lock()
|
||||
passed++
|
||||
mu.Unlock()
|
||||
@ -70,48 +75,53 @@ func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// 建会话失败时要还名额：那次新建实际上没有发生
|
||||
// 建会话失败时要还名额
|
||||
func TestSessionRateRelease(t *testing.T) {
|
||||
l := newLimiter()
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
l.Allow("bot")
|
||||
AllowNewSession(ctx, "bot")
|
||||
}
|
||||
if ok, _ := l.Allow("bot"); ok {
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||||
t.Fatal("应已刷满")
|
||||
}
|
||||
l.Release("bot")
|
||||
if ok, _ := l.Allow("bot"); !ok {
|
||||
ReleaseNewSession(ctx, "bot")
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatal("归还名额后应能再开一条")
|
||||
}
|
||||
// 空记录上 Release 不该 panic
|
||||
l2 := newLimiter()
|
||||
l2.Release("nobody")
|
||||
}
|
||||
|
||||
// 窗口滑过后自动恢复 —— 这正是它优于「终身额度」的地方：
|
||||
// 终身额度跑满要人工重置，速率限制过一个窗口自己好
|
||||
func TestSessionRateWindowSlides(t *testing.T) {
|
||||
l := newLimiter()
|
||||
old := time.Now().Add(-sessionRateWindow - time.Minute)
|
||||
marks := make([]time.Time, sessionRateLimit)
|
||||
for i := range marks {
|
||||
marks[i] = old
|
||||
}
|
||||
l.marks["bot"] = marks
|
||||
|
||||
if ok, _ := l.Allow("bot"); !ok {
|
||||
t.Fatal("窗口外的记录应被清掉，此次应放行")
|
||||
}
|
||||
if len(l.marks["bot"]) != 1 {
|
||||
t.Fatalf("过期记录未清理，剩余 %d 条", len(l.marks["bot"]))
|
||||
}
|
||||
}
|
||||
|
||||
// 人类不走限速（空 agentName）
|
||||
func TestAllowNewSessionSkipsHumans(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < sessionRateLimit*3; i++ {
|
||||
if ok, _ := AllowNewSession(""); !ok {
|
||||
if ok, _ := AllowNewSession(ctx, ""); !ok {
|
||||
t.Fatal("人类不该被限速")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 窗口滑过后自动恢复（过期记录自动清理）
|
||||
func TestSessionRateWindowSlides(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 插入过期记录（1小时前）
|
||||
cutoff := time.Now().Add(-sessionRateWindow - time.Minute)
|
||||
for i := 0; i < sessionRateLimit; i++ {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`,
|
||||
"session:bot", cutoff)
|
||||
if err != nil {
|
||||
t.Fatalf("插入过期记录失败: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 窗口外的记录应被清理，此次应放行
|
||||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||||
t.Fatal("窗口外的记录应被清掉，此次应放行")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user