package ai import ( "errors" "fmt" "path/filepath" "sync" "testing" "time" "github.com/glebarez/sqlite" "github.com/stretchr/testify/require" "gorm.io/gorm" "gorm.io/gorm/logger" "senlinai-agent/backend/internal/models" ) func TestReserveRateLimitIsAtomicUnderConcurrentSQLiteRequests(t *testing.T) { database := newConcurrentAIRateTestDB(t) user := createAIRateTestUser(t, database, "sqlite-rate@example.com") gateway := NewGatewayWithSecret("system-key", "test-encryption-secret") const ( limit = 7 attempts = 40 ) start := make(chan struct{}) results := make(chan error, attempts) var wait sync.WaitGroup for range attempts { wait.Add(1) go func() { defer wait.Done() <-start results <- gateway.ReserveRateLimit(user.ID, "concurrent_session_create", limit, time.Hour) }() } close(start) wait.Wait() close(results) allowed := 0 limited := 0 for err := range results { switch { case err == nil: allowed++ case errors.Is(err, ErrAIRateLimited): limited++ default: require.NoError(t, err) } } require.Equal(t, limit, allowed) require.Equal(t, attempts-limit, limited) var buckets []models.SenlinAgentAIRateBucket require.NoError(t, database.Where("user_id = ? AND action = ?", user.ID, "concurrent_session_create").Find(&buckets).Error) require.Len(t, buckets, 1) require.Equal(t, limit, buckets[0].Count) } func TestReserveRateLimitUsesFixedWindowsAndCountsFailedAttempts(t *testing.T) { database := newConcurrentAIRateTestDB(t) user := createAIRateTestUser(t, database, "window-rate@example.com") gateway := NewGatewayWithSecret("system-key", "test-encryption-secret") current := time.Date(2026, 7, 21, 10, 15, 0, 0, time.UTC) gateway.now = func() time.Time { return current } require.NoError(t, gateway.ReserveRateLimit(user.ID, "windowed_session_create", 2, time.Hour)) require.NoError(t, gateway.ReserveRateLimit(user.ID, "windowed_session_create", 2, time.Hour)) require.ErrorIs(t, gateway.ReserveRateLimit(user.ID, "windowed_session_create", 2, time.Hour), ErrAIRateLimited) current = current.Add(time.Hour) require.NoError(t, gateway.ReserveRateLimit(user.ID, "windowed_session_create", 2, time.Hour)) var buckets []models.SenlinAgentAIRateBucket require.NoError(t, database.Where("user_id = ? AND action = ?", user.ID, "windowed_session_create").Order("window_start asc").Find(&buckets).Error) require.Len(t, buckets, 2) require.Equal(t, []int{2, 1}, []int{buckets[0].Count, buckets[1].Count}) } func newConcurrentAIRateTestDB(t *testing.T) *gorm.DB { t.Helper() databasePath := filepath.ToSlash(filepath.Join(t.TempDir(), "ai-rate.db")) dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)", databasePath) database, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) require.NoError(t, err) sqlDatabase, err := database.DB() require.NoError(t, err) t.Cleanup(func() { require.NoError(t, sqlDatabase.Close()) }) sqlDatabase.SetMaxOpenConns(20) require.NoError(t, models.AutoMigrate(database)) models.DBService = database return database } func createAIRateTestUser(t *testing.T, database *gorm.DB, email string) models.SenlinAgentUser { t.Helper() user := models.SenlinAgentUser{Email: email, DisplayName: email, PasswordHash: "hash"} require.NoError(t, database.Create(&user).Error) return user }