chore: directory migration - gateway→server, web→client/electron

server/internal/repo/sessionrate_test.go

@@ -0,0 +1,127 @@
+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("窗口外的记录应被清掉，此次应放行")
+	}
+}