In-memory rate limits for free mode

Author: jahooma

web/src/app/api/v1/chat/completions/__tests__/free-mode-rate-limiter.test.ts

@@ -0,0 +1,264 @@
+import { afterEach, beforeEach, describe, expect, it, spyOn } from 'bun:test'
+
+import {
+  checkFreeModeRateLimit,
+  FREE_MODE_RATE_LIMITS,
+  resetFreeModeRateLimits,
+} from '../free-mode-rate-limiter'
+
+const MINUTE_MS = 60 * 1000
+const HOUR_MS = 60 * MINUTE_MS
+
+describe('free-mode-rate-limiter', () => {
+  let nowSpy: ReturnType<typeof spyOn>
+  let fakeNow: number
+
+  beforeEach(() => {
+    resetFreeModeRateLimits()
+    fakeNow = 1_000_000_000_000
+    nowSpy = spyOn(Date, 'now').mockImplementation(() => fakeNow)
+  })
+
+  afterEach(() => {
+    nowSpy.mockRestore()
+  })
+
+  function advanceTime(ms: number) {
+    fakeNow += ms
+  }
+
+  function makeRequests(userId: string, count: number) {
+    for (let i = 0; i < count; i++) {
+      const result = checkFreeModeRateLimit(userId)
+      if (result.limited) {
+        throw new Error(`Unexpectedly rate limited on request ${i + 1}`)
+      }
+    }
+  }
+
+  describe('checkFreeModeRateLimit', () => {
+    it('allows the first request', () => {
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(false)
+    })
+
+    it('allows requests up to the per-minute limit', () => {
+      for (let i = 0; i < FREE_MODE_RATE_LIMITS.PER_MINUTE; i++) {
+        const result = checkFreeModeRateLimit('user-1')
+        expect(result.limited).toBe(false)
+      }
+    })
+
+    it('limits when per-minute limit is exceeded', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        expect(result.windowName).toBe('1 minute')
+      }
+    })
+
+    it('limits when per-30-minute limit is exceeded', () => {
+      const perMinute = FREE_MODE_RATE_LIMITS.PER_MINUTE
+      const per30Min = FREE_MODE_RATE_LIMITS.PER_30_MINUTES
+
+      // Spread requests across multiple 1-minute windows to avoid hitting the per-minute limit
+      let sent = 0
+      while (sent < per30Min) {
+        const batch = Math.min(perMinute, per30Min - sent)
+        makeRequests('user-1', batch)
+        sent += batch
+        if (sent < per30Min) {
+          // Advance past the 1-minute window so it resets
+          advanceTime(1 * MINUTE_MS + 1)
+        }
+      }
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        expect(result.windowName).toBe('30 minutes')
+      }
+    })
+
+    it('limits when per-5-hour limit is exceeded', () => {
+      const perMinute = FREE_MODE_RATE_LIMITS.PER_MINUTE
+      const per30Min = FREE_MODE_RATE_LIMITS.PER_30_MINUTES
+      const per5Hours = FREE_MODE_RATE_LIMITS.PER_5_HOURS
+
+      // Spread requests across multiple 30-minute windows
+      let sent = 0
+      while (sent < per5Hours) {
+        const batchFor30Min = Math.min(per30Min, per5Hours - sent)
+        // Within each 30-min window, spread across 1-min windows
+        let sentInWindow = 0
+        while (sentInWindow < batchFor30Min) {
+          const batch = Math.min(perMinute, batchFor30Min - sentInWindow)
+          makeRequests('user-1', batch)
+          sentInWindow += batch
+          if (sentInWindow < batchFor30Min) {
+            advanceTime(1 * MINUTE_MS + 1)
+          }
+        }
+        sent += sentInWindow
+        // Always advance past 30-min window to reset it for the next batch
+        // (stays well within the 5-hour window)
+        advanceTime(30 * MINUTE_MS + 1)
+      }
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        expect(result.windowName).toBe('5 hours')
+      }
+    })
+
+    it('limits when per-7-day limit is exceeded', () => {
+      const perMinute = FREE_MODE_RATE_LIMITS.PER_MINUTE
+      const per30Min = FREE_MODE_RATE_LIMITS.PER_30_MINUTES
+      const per5Hours = FREE_MODE_RATE_LIMITS.PER_5_HOURS
+      const per7Days = FREE_MODE_RATE_LIMITS.PER_7_DAYS
+
+      // Spread requests across multiple 5-hour windows
+      let sent = 0
+      while (sent < per7Days) {
+        const batchFor5Hours = Math.min(per5Hours, per7Days - sent)
+        let sentIn5Hr = 0
+        while (sentIn5Hr < batchFor5Hours) {
+          const batchFor30Min = Math.min(per30Min, batchFor5Hours - sentIn5Hr)
+          let sentIn30Min = 0
+          while (sentIn30Min < batchFor30Min) {
+            const batch = Math.min(perMinute, batchFor30Min - sentIn30Min)
+            makeRequests('user-1', batch)
+            sentIn30Min += batch
+            if (sentIn30Min < batchFor30Min) {
+              advanceTime(1 * MINUTE_MS + 1)
+            }
+          }
+          sentIn5Hr += sentIn30Min
+          advanceTime(30 * MINUTE_MS + 1)
+        }
+        sent += sentIn5Hr
+        // Advance past the 5-hour window (stays within 7-day window)
+        advanceTime(5 * HOUR_MS + 1)
+      }
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        expect(result.windowName).toBe('7 days')
+      }
+    })
+
+    it('does not increment counters when rate limited', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      // These should all be rejected without changing state
+      for (let i = 0; i < 5; i++) {
+        const result = checkFreeModeRateLimit('user-1')
+        expect(result.limited).toBe(true)
+      }
+
+      // After the 1-minute window expires, the user should only have used PER_MINUTE requests
+      // against the 30-minute window, not PER_MINUTE + 5
+      advanceTime(1 * MINUTE_MS + 1)
+
+      // Should be allowed again (1-min window reset)
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(false)
+    })
+
+    it('returns correct retryAfterMs for the violated window', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      // Advance 30 seconds into the 1-minute window
+      advanceTime(30_000)
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        // Should be approximately 30 seconds remaining in the 1-minute window
+        expect(result.retryAfterMs).toBe(1 * MINUTE_MS - 30_000)
+      }
+    })
+
+    it('resets per-minute window after expiry', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      const limited = checkFreeModeRateLimit('user-1')
+      expect(limited.limited).toBe(true)
+
+      // Advance past the 1-minute window
+      advanceTime(1 * MINUTE_MS + 1)
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(false)
+    })
+
+    it('isolates different users', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      // user-1 is rate limited
+      expect(checkFreeModeRateLimit('user-1').limited).toBe(true)
+
+      // user-2 should not be affected
+      const result = checkFreeModeRateLimit('user-2')
+      expect(result.limited).toBe(false)
+    })
+
+    it('retryAfterMs is never negative', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      // Advance to just before expiry
+      advanceTime(1 * MINUTE_MS - 1)
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(true)
+      if (result.limited) {
+        expect(result.retryAfterMs).toBeGreaterThanOrEqual(0)
+      }
+    })
+
+    it('tracks counts across all windows simultaneously', () => {
+      // Make some requests
+      makeRequests('user-1', 5)
+
+      // Advance past 1-minute window but within 30-minute window
+      advanceTime(1 * MINUTE_MS + 1)
+
+      // Make more requests — 1-min counter resets, but 30-min counter keeps accumulating
+      makeRequests('user-1', 5)
+
+      // Advance past 1-minute again
+      advanceTime(1 * MINUTE_MS + 1)
+
+      // The 30-min window should now have 10 requests counted
+      // and the 1-min window should be fresh
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(false)
+    })
+  })
+
+  describe('resetFreeModeRateLimits', () => {
+    it('clears all rate limit state', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+      expect(checkFreeModeRateLimit('user-1').limited).toBe(true)
+
+      resetFreeModeRateLimits()
+
+      const result = checkFreeModeRateLimit('user-1')
+      expect(result.limited).toBe(false)
+    })
+
+    it('clears state for all users', () => {
+      makeRequests('user-1', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+      makeRequests('user-2', FREE_MODE_RATE_LIMITS.PER_MINUTE)
+
+      resetFreeModeRateLimits()
+
+      expect(checkFreeModeRateLimit('user-1').limited).toBe(false)
+      expect(checkFreeModeRateLimit('user-2').limited).toBe(false)
+    })
+  })
+})

web/src/app/api/v1/chat/completions/_post.ts

@@ -65,6 +65,7 @@ import {
   OpenRouterError,
 } from '@/llm-api/openrouter'
 import { extractApiKeyFromHeader } from '@/util/auth'
+import { checkFreeModeRateLimit } from './free-mode-rate-limiter'
 
 const FREE_MODE_ALLOWED_COUNTRIES = new Set([
   'US', 'CA',
@@ -86,11 +87,6 @@ function getCountryCode(req: NextRequest): string | null {
     return cfCountry.toUpperCase()
   }
 
-  const vercelCountry = req.headers.get('x-vercel-ip-country')
-  if (vercelCountry && vercelCountry !== 'XX') {
-    return vercelCountry.toUpperCase()
-  }
-
   const clientIp = extractClientIp(req)
   if (!clientIp) {
     return null
@@ -263,10 +259,9 @@ export async function postChatCompletions(params: {
       const clientIp = extractClientIp(req)
 
       const cfHeader = req.headers.get('cf-ipcountry')
-      const vercelHeader = req.headers.get('x-vercel-ip-country')
       const geoipResult = clientIp ? geoip.lookup(clientIp)?.country ?? null : null
       logger.info(
-        { cfHeader, vercelHeader, geoipResult, resolvedCountry: countryCode, clientIp: clientIp ? '[redacted]' : undefined },
+        { cfHeader, geoipResult, resolvedCountry: countryCode, clientIp: clientIp ? '[redacted]' : undefined },
         'Free mode country detection',
       )
 
@@ -292,6 +287,36 @@ export async function postChatCompletions(params: {
           { status: 403 },
         )
       }
+
+      // Rate limit free mode requests
+      const rateLimitResult = checkFreeModeRateLimit(userId)
+      if (rateLimitResult.limited) {
+        const retryAfterSeconds = Math.ceil(rateLimitResult.retryAfterMs / 1000)
+        const resetTime = new Date(Date.now() + rateLimitResult.retryAfterMs).toISOString()
+        const resetCountdown = formatQuotaResetCountdown(resetTime)
+
+        trackEvent({
+          event: AnalyticsEvent.CHAT_COMPLETIONS_VALIDATION_ERROR,
+          userId,
+          properties: {
+            error: 'free_mode_rate_limited',
+            windowName: rateLimitResult.windowName,
+            retryAfterSeconds,
+          },
+          logger,
+        })
+
+        return NextResponse.json(
+          {
+            error: 'free_mode_rate_limited',
+            message: `Free mode rate limit exceeded (${rateLimitResult.windowName} limit). Try again ${resetCountdown}.`,
+          },
+          {
+            status: 429,
+            headers: { 'Retry-After': String(retryAfterSeconds) },
+          },
+        )
+      }
     }
 
     // Extract and validate agent run ID