Inject free mode rate limiter in completions tests

Author: jahooma

web/src/app/api/v1/chat/completions/__tests__/completions.test.ts

@@ -1104,9 +1104,18 @@ describe('/api/v1/chat/completions POST endpoint', () => {
     it(
       'counts child Gemini thinker requests toward the free-mode request limit',
       async () => {
-        expect(checkFreeModeRateLimit('user-gemini-rate-limit').limited).toBe(
-          false,
-        )
+        let rateLimitChecks = 0
+        const checkFreeModeRateLimitForTest = mock((userId: string) => {
+          expect(userId).toBe('user-gemini-rate-limit')
+          rateLimitChecks += 1
+          return rateLimitChecks === 1
+            ? { limited: false as const }
+            : {
+                limited: true as const,
+                windowName: '1 second',
+                retryAfterMs: 1_000,
+              }
+        })
 
         const createRequest = () =>
           new NextRequest('http://localhost:3000/api/v1/chat/completions', {
@@ -1135,6 +1144,7 @@ describe('/api/v1/chat/completions POST endpoint', () => {
           insertMessageBigquery: mockInsertMessageBigquery,
           loggerWithContext: mockLoggerWithContext,
           checkSessionAdmissible: mockCheckSessionAdmissibleAllow,
+          checkFreeModeRateLimit: checkFreeModeRateLimitForTest,
         })
 
         const firstResponse = await postChatCompletions(createPostParams())
@@ -1144,6 +1154,7 @@ describe('/api/v1/chat/completions POST endpoint', () => {
         expect(limitedResponse.status).toBe(429)
         const body = await limitedResponse.json()
         expect(body.error).toBe('free_mode_rate_limited')
+        expect(checkFreeModeRateLimitForTest).toHaveBeenCalledTimes(2)
       },
       FETCH_PATH_TEST_TIMEOUT_MS,
     )

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

@@ -78,7 +78,7 @@ import { getFreeModeCountryAccess } from '@/server/free-mode-country'
 import type { SessionGateResult } from '@/server/free-session/public-api'
 import { extractApiKeyFromHeader } from '@/util/auth'
 import { withDefaultProperties } from '@codebuff/common/analytics'
-import { checkFreeModeRateLimit } from './free-mode-rate-limiter'
+import { checkFreeModeRateLimit as defaultCheckFreeModeRateLimit } from './free-mode-rate-limiter'
 
 export const formatQuotaResetCountdown = (
   nextQuotaReset: string | null | undefined,
@@ -117,6 +117,7 @@ export const formatQuotaResetCountdown = (
 }
 
 export type CheckSessionAdmissibleFn = typeof checkSessionAdmissible
+export type CheckFreeModeRateLimitFn = typeof defaultCheckFreeModeRateLimit
 
 type GateRejectCode = Extract<SessionGateResult, { ok: false }>['code']
 
@@ -147,6 +148,9 @@ export async function postChatCompletions(params: {
   /** Optional override for the freebuff waiting-room gate. Defaults to the
    *  real check backed by Postgres; tests inject a no-op. */
   checkSessionAdmissible?: CheckSessionAdmissibleFn
+  /** Optional override for the free-mode rate limiter. Tests inject this to
+   *  avoid coupling to process-global limiter state. */
+  checkFreeModeRateLimit?: CheckFreeModeRateLimitFn
 }) {
   const {
     req,
@@ -159,6 +163,7 @@ export async function postChatCompletions(params: {
     ensureSubscriberBlockGrant,
     getUserPreferences,
     checkSessionAdmissible: checkSession = checkSessionAdmissible,
+    checkFreeModeRateLimit = defaultCheckFreeModeRateLimit,
   } = params
   let { logger } = params
   let { trackEvent } = params