JianFeeeee
9d4718a412
## 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 两版）。
128 lines
3.0 KiB
Go
128 lines
3.0 KiB
Go
package repo
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/agentmail/gateway/internal/db"
|
||
)
|
||
|
||
// 新建会话速率限制：DB 版
|
||
//
|
||
// 这些测试用真实的 SQLite（setupTestDB），验证速率限制的原子性与窗口滑动。
|
||
// 原来的内存版测试依赖 sessionRateLimiter 结构体，替换为 DB 版后重写。
|
||
|
||
func TestSessionRateAllowsUpToLimit(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 1; i <= sessionRateLimit; i++ {
|
||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||
t.Fatalf("第 %d 次应放行（上限 %d）", i, sessionRateLimit)
|
||
}
|
||
}
|
||
ok, retry := AllowNewSession(ctx, "bot")
|
||
if ok {
|
||
t.Fatal("超过上限应拦下")
|
||
}
|
||
if retry < 1 {
|
||
t.Fatalf("应给出正的重试等待秒数，实际 %d", retry)
|
||
}
|
||
}
|
||
|
||
func TestSessionRateIsPerAgent(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit; i++ {
|
||
AllowNewSession(ctx, "busy")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "busy"); ok {
|
||
t.Fatal("busy 应已被拦")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "idle"); !ok {
|
||
t.Fatal("另一个 Agent 不该被牵连")
|
||
}
|
||
}
|
||
|
||
// 并发请求不能把上限刷穿（判断与记账必须原子）
|
||
func TestSessionRateConcurrentDoesNotOverrun(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
var wg sync.WaitGroup
|
||
var mu sync.Mutex
|
||
passed := 0
|
||
|
||
for i := 0; i < sessionRateLimit*4; i++ {
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||
mu.Lock()
|
||
passed++
|
||
mu.Unlock()
|
||
}
|
||
}()
|
||
}
|
||
wg.Wait()
|
||
|
||
if passed != sessionRateLimit {
|
||
t.Fatalf("%d 并发下放行 %d 次，期望恰好 %d 次",
|
||
sessionRateLimit*4, passed, sessionRateLimit)
|
||
}
|
||
}
|
||
|
||
// 建会话失败时要还名额
|
||
func TestSessionRateRelease(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit; i++ {
|
||
AllowNewSession(ctx, "bot")
|
||
}
|
||
if ok, _ := AllowNewSession(ctx, "bot"); ok {
|
||
t.Fatal("应已刷满")
|
||
}
|
||
ReleaseNewSession(ctx, "bot")
|
||
if ok, _ := AllowNewSession(ctx, "bot"); !ok {
|
||
t.Fatal("归还名额后应能再开一条")
|
||
}
|
||
}
|
||
|
||
// 人类不走限速（空 agentName）
|
||
func TestAllowNewSessionSkipsHumans(t *testing.T) {
|
||
setupTestDB(t)
|
||
ctx := context.Background()
|
||
|
||
for i := 0; i < sessionRateLimit*3; i++ {
|
||
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("窗口外的记录应被清掉，此次应放行")
|
||
}
|
||
}
|