Address Freebuff gating review comments

Author: jahooma

web/src/app/api/v1/freebuff/session/__tests__/session.test.ts

@@ -163,11 +163,15 @@ const LOGGER = {
 function makeDeps(
   sessionDeps: SessionDeps,
   userId: string | null,
-  opts: { banned?: boolean } = {},
+  opts: {
+    banned?: boolean
+    getCountryAccess?: FreebuffSessionDeps['getCountryAccess']
+  } = {},
 ): FreebuffSessionDeps {
   return {
     logger: LOGGER as unknown as FreebuffSessionDeps['logger'],
-    getCountryAccess: async (req) => testCountryAccess(req),
+    getCountryAccess:
+      opts.getCountryAccess ?? (async (req) => testCountryAccess(req)),
     getUserInfoFromApiKey: (async () =>
       userId
         ? { id: userId, banned: opts.banned ?? false }
@@ -332,6 +336,42 @@ describe('GET /api/v1/freebuff/session', () => {
     expect(body.countryBlockReason).toBe('country_not_allowed')
   })
 
+  test('skips country recheck on GET when the stored check is recent', async () => {
+    const sessionDeps = makeSessionDeps()
+    sessionDeps.rows.set('u1', {
+      user_id: 'u1',
+      status: 'queued',
+      active_instance_id: 'inst-1',
+      model: DEFAULT_MODEL,
+      country_code: 'US',
+      cf_country: 'US',
+      geoip_country: null,
+      country_block_reason: null,
+      ip_privacy_signals: [],
+      client_ip_hash: 'test-ip-hash',
+      country_checked_at: new Date('2026-04-17T11:45:00Z'),
+      queued_at: new Date('2026-04-17T11:45:00Z'),
+      admitted_at: null,
+      expires_at: null,
+      created_at: new Date('2026-04-17T11:45:00Z'),
+      updated_at: new Date('2026-04-17T11:45:00Z'),
+    })
+    let countryChecks = 0
+    const resp = await getFreebuffSession(
+      makeReq('ok', { cfCountry: 'FR' }),
+      makeDeps(sessionDeps, 'u1', {
+        getCountryAccess: async (req) => {
+          countryChecks++
+          return testCountryAccess(req)
+        },
+      }),
+    )
+    const body = await resp.json()
+    expect(resp.status).toBe(200)
+    expect(body.status).toBe('queued')
+    expect(countryChecks).toBe(0)
+  })
+
   test('returns banned 403 on GET for banned user', async () => {
     const sessionDeps = makeSessionDeps()
     const resp = await getFreebuffSession(

web/src/app/api/v1/freebuff/session/_handlers.ts

@@ -6,11 +6,19 @@ import {
   getSessionState,
   requestSession,
 } from '@/server/free-session/public-api'
-import { getFreeModeCountryAccess } from '@/server/free-mode-country'
+import { getSessionRow as getStoredSessionRow } from '@/server/free-session/store'
+import {
+  FREE_MODE_ALLOWED_COUNTRIES,
+  getFreeModeCountryAccess,
+  IPINFO_PRIVACY_CACHE_TTL_MS,
+} from '@/server/free-mode-country'
 import { extractApiKeyFromHeader } from '@/util/auth'
 
 import type { FreeModeCountryAccess } from '@/server/free-mode-country'
-import type { FreeSessionCountryAccessMetadata } from '@/server/free-session/types'
+import type {
+  FreeSessionCountryAccessMetadata,
+  InternalSessionRow,
+} from '@/server/free-session/types'
 import type { SessionDeps } from '@/server/free-session/public-api'
 import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
@@ -57,10 +65,10 @@ function toSessionCountryAccess(
 async function countryBlockedResponse(
   req: NextRequest,
   deps: FreebuffSessionDeps,
-): Promise<
-  | { response: NextResponse; countryAccess: FreeModeCountryAccess }
-  | { response: null; countryAccess: FreeModeCountryAccess }
-> {
+): Promise<{
+  response: NextResponse | null
+  countryAccess: FreeModeCountryAccess
+}> {
   const countryAccess = await getCountryAccess(req, deps)
   if (countryAccess.allowed) {
     return { response: null, countryAccess }
@@ -79,6 +87,32 @@ async function countryBlockedResponse(
   }
 }
 
+function hasRecentAllowedCountryCheck(
+  row: InternalSessionRow | null,
+  now: Date,
+): boolean {
+  if (!row?.country_checked_at || row.country_block_reason !== null) {
+    return false
+  }
+  if (!row.country_code || !FREE_MODE_ALLOWED_COUNTRIES.has(row.country_code)) {
+    return false
+  }
+  return (
+    now.getTime() - row.country_checked_at.getTime() <
+    IPINFO_PRIVACY_CACHE_TTL_MS
+  )
+}
+
+async function shouldSkipGetCountryCheck(
+  userId: string,
+  deps: FreebuffSessionDeps,
+): Promise<boolean> {
+  const getSessionRow = deps.sessionDeps?.getSessionRow ?? getStoredSessionRow
+  const row = await getSessionRow(userId)
+  const now = deps.sessionDeps?.now?.() ?? new Date()
+  return hasRecentAllowedCountryCheck(row, now)
+}
+
 /** Header the CLI uses to identify which instance is polling. Used by GET to
  *  detect when another CLI on the same account has rotated the id. */
 export const FREEBUFF_INSTANCE_HEADER = 'x-freebuff-instance-id'
@@ -220,10 +254,12 @@ export async function getFreebuffSession(
   const auth = await resolveUser(req, deps)
   if ('error' in auth) return auth.error
 
-  const { response: blocked } = await countryBlockedResponse(req, deps)
-  if (blocked) return blocked
-
   try {
+    if (!(await shouldSkipGetCountryCheck(auth.userId, deps))) {
+      const { response: blocked } = await countryBlockedResponse(req, deps)
+      if (blocked) return blocked
+    }
+
     const claimedInstanceId =
       req.headers.get(FREEBUFF_INSTANCE_HEADER) ?? undefined
     const state = await getSessionState({

web/src/server/free-mode-country.ts

@@ -65,7 +65,7 @@ type ResolvedCountryAccess = Omit<
   countryCode: string
 }
 
-const IPINFO_PRIVACY_CACHE_TTL_MS = 30 * 60 * 1000
+export const IPINFO_PRIVACY_CACHE_TTL_MS = 30 * 60 * 1000
 const IPINFO_PRIVACY_CACHE_MAX_ENTRIES = 5000
 const ipinfoPrivacyCache = new Map<
   string,
@@ -109,21 +109,14 @@ function setIpinfoPrivacyCache(
   ip: string,
   privacy: FreeModeIpPrivacy | null,
 ): void {
-  const now = Date.now()
-  for (const [cachedIp, cached] of ipinfoPrivacyCache) {
-    if (cached.expiresAt <= now) {
-      ipinfoPrivacyCache.delete(cachedIp)
-    }
-  }
-
   while (ipinfoPrivacyCache.size >= IPINFO_PRIVACY_CACHE_MAX_ENTRIES) {
     const oldestIp = ipinfoPrivacyCache.keys().next().value
     if (!oldestIp) break
     ipinfoPrivacyCache.delete(oldestIp)
   }
 
   ipinfoPrivacyCache.set(ip, {
-    expiresAt: now + IPINFO_PRIVACY_CACHE_TTL_MS,
+    expiresAt: Date.now() + IPINFO_PRIVACY_CACHE_TTL_MS,
     privacy,
   })
 }
@@ -182,25 +175,6 @@ export async function lookupIpinfoPrivacy(params: {
   return privacy
 }
 
-async function getIpPrivacy(
-  clientIp: string | undefined,
-  options: FreeModeCountryAccessOptions,
-): Promise<FreeModeIpPrivacy | null> {
-  if (!clientIp) return null
-  try {
-    if (options.lookupIpPrivacy) {
-      return await options.lookupIpPrivacy(clientIp)
-    }
-    return await lookupIpinfoPrivacy({
-      ip: clientIp,
-      token: options.ipinfoToken,
-      fetch: options.fetch ?? globalThis.fetch,
-    })
-  } catch {
-    return null
-  }
-}
-
 export async function getFreeModeCountryAccess(
   req: NextRequest,
   options: FreeModeCountryAccessOptions,
@@ -290,7 +264,19 @@ export async function getFreeModeCountryAccess(
     }
   }
 
-  const ipPrivacy = await getIpPrivacy(clientIp, options)
+  let ipPrivacy: FreeModeIpPrivacy | null
+  try {
+    ipPrivacy = options.lookupIpPrivacy
+      ? await options.lookupIpPrivacy(clientIp)
+      : await lookupIpinfoPrivacy({
+          ip: clientIp,
+          token: options.ipinfoToken,
+          fetch: options.fetch ?? globalThis.fetch,
+        })
+  } catch {
+    ipPrivacy = null
+  }
+
   if (!ipPrivacy) {
     return {
       ...baseAccess,