Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion cli/src/app.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -384,7 +384,8 @@ const AuthedSurface = ({
IS_FREEBUFF &&
(session === null ||
session.status === 'queued' ||
session.status === 'none')
session.status === 'none' ||
session.status === 'country_blocked')
) {
return <WaitingRoomScreen session={session} error={sessionError} />
}
Expand Down
17 changes: 17 additions & 0 deletions cli/src/components/waiting-room-screen.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,23 @@ export const WaitingRoomScreen: React.FC<WaitingRoomScreenProps> = ({
{session?.status === 'disabled' && (
<text style={{ fg: theme.muted }}>Waiting room disabled.</text>
)}

{/* Country outside the free-mode allowlist. Terminal — polling has
stopped. Tell the user up front rather than letting them wait in
the queue only to be rejected at the chat/completions gate. */}
{session?.status === 'country_blocked' && (
<>
<text style={{ fg: theme.secondary, marginBottom: 1 }}>
⚠ Free mode isn't available in your region
</text>
<text style={{ fg: theme.muted, wrapMode: 'word' }}>
We detected your location as{' '}
<span fg={theme.foreground}>{session.countryCode}</span>,
which is outside the countries where freebuff is currently
offered. Press Ctrl+C to exit.
</text>
</>
)}
</box>
</box>

Expand Down
15 changes: 15 additions & 0 deletions cli/src/hooks/use-freebuff-session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,20 @@ async function callSession(
if (resp.status === 404) {
return { status: 'disabled' }
}
// 403 with a country_blocked body is a terminal signal, not an error — the
// server rejects non-allowlist countries up front (see session _handlers.ts)
// so users don't wait through the queue only to be rejected at chat time.
// The 403 status (rather than 200) is deliberate: older CLIs that don't
// know this status treat it as a generic error and back off on the 10s
// error-retry cadence instead of tight-polling an unrecognized 200 body.
if (resp.status === 403) {
const body = (await resp.json().catch(() => null)) as
| FreebuffSessionResponse
| null
if (body && body.status === 'country_blocked') {
return body
}
}
if (!resp.ok) {
const text = await resp.text().catch(() => '')
throw new Error(
Expand Down Expand Up @@ -80,6 +94,7 @@ function nextDelayMs(next: FreebuffSessionResponse): number | null {
case 'none':
case 'disabled':
case 'superseded':
case 'country_blocked':
return null
}
}
Expand Down
9 changes: 9 additions & 0 deletions common/src/types/freebuff-session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,3 +59,12 @@ export type FreebuffSessionServerResponse =
* surfaces it as a 409 for fast in-flight feedback. */
status: 'superseded'
}
| {
/** Request originated from a country outside the free-mode allowlist.
* Returned before queue admission so users don't wait through the
* room only to be rejected on their first chat request. Terminal —
* CLI stops polling and shows a "not available in your country"
* screen. `countryCode` is the resolved country for display. */
status: 'country_blocked'
countryCode: string
}
33 changes: 5 additions & 28 deletions web/src/app/api/v1/chat/completions/_post.ts
Original file line number Diff line number Diff line change
Expand Up @@ -68,40 +68,17 @@ import {
OpenRouterError,
} from '@/llm-api/openrouter'
import { checkSessionAdmissible } from '@/server/free-session/public-api'
import {
FREE_MODE_ALLOWED_COUNTRIES,
extractClientIp,
getCountryCode,
} 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'

const FREE_MODE_ALLOWED_COUNTRIES = new Set([
'US', 'CA',
'GB', 'AU', 'NZ',
'NO', 'SE', 'NL', 'DK', 'DE', 'FI', 'BE', 'LU', 'CH', 'IE', 'IS',
])

function extractClientIp(req: NextRequest): string | undefined {
const forwardedFor = req.headers.get('x-forwarded-for')
if (forwardedFor) {
return forwardedFor.split(',')[0].trim()
}
return req.headers.get('x-real-ip') ?? undefined
}

function getCountryCode(req: NextRequest): string | null {
const cfCountry = req.headers.get('cf-ipcountry')
if (cfCountry && cfCountry !== 'XX' && cfCountry !== 'T1') {
return cfCountry.toUpperCase()
}

const clientIp = extractClientIp(req)
if (!clientIp) {
return null
}
const geo = geoip.lookup(clientIp)
return geo?.country ?? null
}

export const formatQuotaResetCountdown = (
nextQuotaReset: string | null | undefined,
): string => {
Expand Down
40 changes: 39 additions & 1 deletion web/src/app/api/v1/freebuff/session/__tests__/session.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,12 @@ import type { NextRequest } from 'next/server'

function makeReq(
apiKey: string | null,
opts: { instanceId?: string } = {},
opts: { instanceId?: string; cfCountry?: string } = {},
): NextRequest {
const headers = new Headers()
if (apiKey) headers.set('Authorization', `Bearer ${apiKey}`)
if (opts.instanceId) headers.set(FREEBUFF_INSTANCE_HEADER, opts.instanceId)
if (opts.cfCountry) headers.set('cf-ipcountry', opts.cfCountry)
return {
headers,
} as unknown as NextRequest
Expand Down Expand Up @@ -102,6 +103,31 @@ describe('POST /api/v1/freebuff/session', () => {
const body = await resp.json()
expect(body.status).toBe('disabled')
})

test('returns country_blocked without joining the queue for disallowed country', async () => {
const sessionDeps = makeSessionDeps()
const resp = await postFreebuffSession(
makeReq('ok', { cfCountry: 'FR' }),
makeDeps(sessionDeps, 'u1'),
)
// 403 (not 200) so older CLIs that don't know `country_blocked` fall into
// their error-retry backoff instead of tight-polling.
expect(resp.status).toBe(403)
const body = await resp.json()
expect(body.status).toBe('country_blocked')
expect(body.countryCode).toBe('FR')
expect(sessionDeps.rows.size).toBe(0)
})

test('allows queue entry for allowed country', async () => {
const sessionDeps = makeSessionDeps()
const resp = await postFreebuffSession(
makeReq('ok', { cfCountry: 'US' }),
makeDeps(sessionDeps, 'u1'),
)
const body = await resp.json()
expect(body.status).toBe('queued')
})
})
Comment on lines +107 to 131
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Missing test for the null-country fail-open path

The new tests cover blocked (FR) and allowed (US) countries, but the shared module's stated contract — "null country (VPN/localhost) fails open" — is not exercised. A request without cf-ipcountry and no resolvable IP should pass through to queue admission. Given this is the main escape hatch for VPN/localhost users, a test that verifies postFreebuffSession creates a queued session when no country headers are present would round out the coverage.

test('allows queue entry when country cannot be determined (null/VPN)', async () => {
  const sessionDeps = makeSessionDeps()
  // No cfCountry → geoip returns null → fail open
  const resp = await postFreebuffSession(makeReq('ok'), makeDeps(sessionDeps, 'u1'))
  const body = await resp.json()
  expect(body.status).toBe('queued')
})


describe('GET /api/v1/freebuff/session', () => {
Expand All @@ -113,6 +139,18 @@ describe('GET /api/v1/freebuff/session', () => {
expect(body.status).toBe('none')
})

test('returns country_blocked for disallowed country on GET', async () => {
const sessionDeps = makeSessionDeps()
const resp = await getFreebuffSession(
makeReq('ok', { cfCountry: 'FR' }),
makeDeps(sessionDeps, 'u1'),
)
expect(resp.status).toBe(403)
const body = await resp.json()
expect(body.status).toBe('country_blocked')
expect(body.countryCode).toBe('FR')
})

test('returns superseded when active row exists with mismatched instance id', async () => {
const sessionDeps = makeSessionDeps()
sessionDeps.rows.set('u1', {
Expand Down
30 changes: 30 additions & 0 deletions web/src/app/api/v1/freebuff/session/_handlers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,37 @@ import {
getSessionState,
requestSession,
} from '@/server/free-session/public-api'
import {
FREE_MODE_ALLOWED_COUNTRIES,
getCountryCode,
} from '@/server/free-mode-country'
import { extractApiKeyFromHeader } from '@/util/auth'

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'
import type { NextRequest } from 'next/server'

/** Early country gate. Mirrors the chat/completions check: if we can resolve
* the caller's country and it's not on the allowlist, short-circuit with a
* terminal `country_blocked` response so the CLI can show the warning
* screen without ever joining the queue. Null country (VPN / localhost)
* fails open — chat/completions will catch it later if it matters.
*
* Returns HTTP 403 (not 200) so older CLIs — which don't know the
* `country_blocked` status and would tight-poll on an unrecognized 200
* body — fall into their existing `!resp.ok` error path and back off on
* the 10s error retry cadence. The new CLI parses the 403 body directly. */
function countryBlockedResponse(req: NextRequest): NextResponse | null {
const countryCode = getCountryCode(req)
if (!countryCode) return null
if (FREE_MODE_ALLOWED_COUNTRIES.has(countryCode)) return null
return NextResponse.json(
{ status: 'country_blocked', countryCode },
{ status: 403 },
)
}

/** 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'
Expand Down Expand Up @@ -95,6 +119,9 @@ export async function postFreebuffSession(
const auth = await resolveUser(req, deps)
if ('error' in auth) return auth.error

const blocked = countryBlockedResponse(req)
if (blocked) return blocked

try {
const state = await requestSession({
userId: auth.userId,
Expand All @@ -117,6 +144,9 @@ export async function getFreebuffSession(
const auth = await resolveUser(req, deps)
if ('error' in auth) return auth.error

const blocked = countryBlockedResponse(req)
if (blocked) return blocked

try {
const claimedInstanceId = req.headers.get(FREEBUFF_INSTANCE_HEADER) ?? undefined
const state = await getSessionState({
Expand Down
43 changes: 43 additions & 0 deletions web/src/server/free-mode-country.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import geoip from 'geoip-lite'

import type { NextRequest } from 'next/server'

export const FREE_MODE_ALLOWED_COUNTRIES = new Set([
'US', 'CA',
'GB', 'AU', 'NZ',
'NO', 'SE', 'NL', 'DK', 'DE', 'FI', 'BE', 'LU', 'CH', 'IE', 'IS',
])

export function extractClientIp(req: NextRequest): string | undefined {
const forwardedFor = req.headers.get('x-forwarded-for')
if (forwardedFor) {
return forwardedFor.split(',')[0].trim()
}
return req.headers.get('x-real-ip') ?? undefined
}

export function getCountryCode(req: NextRequest): string | null {
const cfCountry = req.headers.get('cf-ipcountry')
if (cfCountry && cfCountry !== 'XX' && cfCountry !== 'T1') {
return cfCountry.toUpperCase()
}

const clientIp = extractClientIp(req)
if (!clientIp) {
return null
}
const geo = geoip.lookup(clientIp)
return geo?.country ?? null
}

/**
* Returns true if the request's resolved country is allowed to use free
* mode, false if it's explicitly disallowed. Returns null when country can't
* be determined (VPN / localhost / corporate proxy) — callers should fail
* open in that case to match the chat-completions gate.
*/
export function isCountryAllowedForFreeMode(req: NextRequest): boolean | null {
const countryCode = getCountryCode(req)
if (!countryCode) return null
return FREE_MODE_ALLOWED_COUNTRIES.has(countryCode)
}
Comment on lines +39 to +43
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 isCountryAllowedForFreeMode is exported but never imported

This helper was added to the shared module but nothing in the codebase actually imports it — both _handlers.ts and _post.ts call getCountryCode + FREE_MODE_ALLOWED_COUNTRIES.has() directly. Either wire up the callers to use this function (which would simplify countryBlockedResponse in _handlers.ts) or remove it to keep the public surface minimal.

Loading