feat(ai-gateway): restrict custom LLMs to allowed country codes

apps/web/src/app/api/openrouter/[...path]/route.test.ts

@@ -408,6 +408,19 @@ describe('POST /api/openrouter/v1/chat/completions rules-engine actions', () =>
     expect(mockedGetProvider).toHaveBeenCalledTimes(1);
     expect(mockedUpstreamRequest.mock.calls[0]?.[0].body.model).toBe('openai/gpt-4o');
   });
+
+  it('rejects with 403 country_not_allowed when provider resolution is forbidden', async () => {
+    mockedGetProvider.mockResolvedValueOnce({ kind: 'forbidden' });
+
+    const { POST } = await import('./route');
+    const response = await POST(makeRequest(makeBody('kilo-internal/custom-model')) as never);
+
+    expect(response.status).toBe(403);
+    expect(await response.json()).toMatchObject({
+      error_type: 'country_not_allowed',
+    });
+    expect(mockedUpstreamRequest).not.toHaveBeenCalled();
+  });
 });
 
 describe('kilo-auto/efficient classifier billing', () => {

apps/web/src/app/api/openrouter/[...path]/route.ts

@@ -52,6 +52,7 @@ import {
   usageLimitExceededResponse,
   wrapInSafeNextResponse,
   forbiddenFreeModelResponse,
+  customLlmCountryNotAllowedResponse,
   storeAndPreviousResponseIdIsNotSupported,
   apiKindNotSupportedResponse,
 } from '@/lib/ai-gateway/llm-proxy-helpers';
@@ -239,6 +240,9 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
 
   // Extract IP early (needed for free model routing fallback and rate limiting)
   const ipAddress = request.headers.get('x-forwarded-for')?.split(',')[0]?.trim();
+  // Vercel sets `x-vercel-ip-country` to the ISO 3166-1 alpha-2 country code
+  // of the client. Used to enforce a custom LLM's `country_codes` allow-list.
+  const clientCountry = request.headers.get('x-vercel-ip-country')?.trim() || null;
 
   const modeHeader = extractHeaderAndLimitLength(request, 'x-kilocode-mode');
   const taskId = extractHeaderAndLimitLength(request, 'x-kilocode-taskid') ?? undefined;
@@ -544,6 +548,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     organizationId,
     taskId,
     clientIp: ipAddress ?? null,
+    clientCountry,
     machineId: machineIdHeader,
   });
   if (initialProviderResultForAbuseService.kind === 'not-found') {
@@ -554,6 +559,9 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
   if (initialProviderResultForAbuseService.kind === 'unavailable') {
     return temporarilyUnavailableResponse();
   }
+  if (initialProviderResultForAbuseService.kind === 'forbidden') {
+    return customLlmCountryNotAllowedResponse(effectiveModelIdLowerCased);
+  }
   let effectiveProviderContext = initialProviderResultForAbuseService;
 
   if (
@@ -661,6 +669,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
       organizationId,
       taskId,
       clientIp: ipAddress ?? null,
+      clientCountry,
       machineId: machineIdHeader,
     });
     if (quarantineProviderResult.kind === 'not-found') {
@@ -675,6 +684,12 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
       }
       return temporarilyUnavailableResponse();
     }
+    if (quarantineProviderResult.kind === 'forbidden') {
+      if (rulesEngineDecision.delayMs > 0) {
+        await sleepForRulesEngineAction(rulesEngineDecision.delayMs);
+      }
+      return customLlmCountryNotAllowedResponse(effectiveModelIdLowerCased);
+    }
 
     effectiveProviderContext = quarantineProviderResult;
 

apps/web/src/lib/ai-gateway/custom-llm/country-access.test.ts

@@ -0,0 +1,39 @@
+import { describe, expect, it } from '@jest/globals';
+import { isCountryAllowed } from './country-access';
+
+describe('isCountryAllowed', () => {
+  it('allows all countries when the allow-list is absent', () => {
+    expect(isCountryAllowed(undefined, 'US')).toBe(true);
+    expect(isCountryAllowed(undefined, null)).toBe(true);
+  });
+
+  it('allows all countries when the allow-list is empty', () => {
+    expect(isCountryAllowed([], 'US')).toBe(true);
+    expect(isCountryAllowed([], null)).toBe(true);
+  });
+
+  it('allows a request whose country is in the allow-list', () => {
+    expect(isCountryAllowed(['US', 'CA'], 'US')).toBe(true);
+    expect(isCountryAllowed(['US', 'CA'], 'CA')).toBe(true);
+  });
+
+  it('denies a request whose country is not in the allow-list', () => {
+    expect(isCountryAllowed(['US', 'CA'], 'GB')).toBe(false);
+  });
+
+  it('compares case-insensitively', () => {
+    expect(isCountryAllowed(['us', 'ca'], 'US')).toBe(true);
+    expect(isCountryAllowed(['US', 'CA'], 'ca')).toBe(true);
+  });
+
+  it('trims surrounding whitespace before comparing', () => {
+    expect(isCountryAllowed([' US ', 'CA '], 'US')).toBe(true);
+    expect(isCountryAllowed(['US'], '  us  ')).toBe(true);
+  });
+
+  it('fails closed when the request country cannot be determined', () => {
+    expect(isCountryAllowed(['US'], null)).toBe(false);
+    expect(isCountryAllowed(['US'], '')).toBe(false);
+    expect(isCountryAllowed(['US'], '   ')).toBe(false);
+  });
+});

apps/web/src/lib/ai-gateway/custom-llm/country-access.ts

@@ -0,0 +1,30 @@
+/**
+ * Determines whether a request is allowed to access a custom LLM based on its
+ * `country_codes` allow-list.
+ *
+ * `country_codes` are ISO 3166-1 alpha-2 codes (e.g. "US", "GB") sourced from
+ * Vercel's `x-vercel-ip-country` request header. Comparison is
+ * case-insensitive and trims surrounding whitespace.
+ *
+ * Semantics:
+ * - An absent or empty `allowedCountries` list disables the restriction
+ *   (every country is allowed).
+ * - When a non-empty list is configured and the request country cannot be
+ *   determined (`requestCountry` is null or empty), access is denied
+ *   (fail-closed) so a misconfigured or spoofed header cannot bypass the
+ *   intended geographic restriction.
+ */
+export function isCountryAllowed(
+  allowedCountries: readonly string[] | undefined,
+  requestCountry: string | null
+): boolean {
+  if (!allowedCountries || allowedCountries.length === 0) {
+    return true;
+  }
+  const normalizedRequestCountry = requestCountry?.trim().toUpperCase() ?? null;
+  if (!normalizedRequestCountry) {
+    return false;
+  }
+  const normalizedAllowed = new Set(allowedCountries.map(code => code.toUpperCase()));
+  return normalizedAllowed.has(normalizedRequestCountry);
+}

apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts

@@ -325,6 +325,14 @@ export function featureExclusiveModelResponse(modelId: string) {
   );
 }
 
+export function customLlmCountryNotAllowedResponse(modelId: string) {
+  const error = `${modelId} is not available in the requester's country.`;
+  return NextResponse.json(
+    { error, error_type: ProxyErrorType.country_not_allowed, message: error },
+    { status: 403 }
+  );
+}
+
 export function storeAndPreviousResponseIdIsNotSupported() {
   const error = 'The store and previous_response_id fields are not supported.';
   return NextResponse.json(

apps/web/src/lib/ai-gateway/providers/get-provider.ts

@@ -17,6 +17,7 @@ import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions';
 import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok';
 import { CustomLlmDefinitionSchema } from '@kilocode/db';
 import { buildDirectProvider } from '@/lib/ai-gateway/experiments/build-direct-provider';
+import { isCountryAllowed } from '@/lib/ai-gateway/custom-llm/country-access';
 import { isPublicIdExperimented } from '@/lib/ai-gateway/experiments/membership';
 import {
   pickModelExperimentVariant,
@@ -50,12 +51,15 @@ export type GetProviderProviderResult = {
 /**
  * Discriminated routing result. `not-found` maps to the local
  * model-unavailable response (used by paused experiments); `unavailable`
- * maps to a 503 temporarily-unavailable response (cache/DB/config failure).
+ * maps to a 503 temporarily-unavailable response (cache/DB/config failure);
+ * `forbidden` maps to a 403 response (e.g. a custom LLM whose `country_codes`
+ * allow-list does not include the requester's country).
  */
 export type GetProviderResult =
   | GetProviderProviderResult
   | { kind: 'not-found' }
-  | { kind: 'unavailable' };
+  | { kind: 'unavailable' }
+  | { kind: 'forbidden' };
 
 async function checkDirectBYOK(
   user: User | AnonymousUserContext,
@@ -91,8 +95,9 @@ async function checkDirectBYOK(
 
 async function checkCustomLlm(
   requestedModel: string,
-  organizationId: string
-): Promise<GetProviderProviderResult | null> {
+  organizationId: string,
+  clientCountry: string | null
+): Promise<GetProviderResult | null> {
   const [row] = await db
     .select()
     .from(custom_llm2)
@@ -105,6 +110,13 @@ async function checkCustomLlm(
   if (!customLlm || !customLlm.organization_ids.includes(organizationId)) {
     return null;
   }
+  // `country_codes` is an allow-list of ISO 3166-1 alpha-2 codes sourced from
+  // Vercel's `x-vercel-ip-country` header. An empty/absent list allows all
+  // countries; a configured list fails closed when the request country is
+  // unknown or not in the list.
+  if (!isCountryAllowed(customLlm.country_codes, clientCountry)) {
+    return { kind: 'forbidden' };
+  }
   return {
     kind: 'provider',
     provider: buildDirectProvider(
@@ -153,13 +165,26 @@ export type GetProviderInput = {
    *  allocation subject for experiment routing when no userId/machineId
    *  is available. */
   clientIp: string | null;
+  /** Resolved client country (ISO 3166-1 alpha-2) from Vercel's
+   *  `x-vercel-ip-country` header. Used to enforce a custom LLM's
+   *  `country_codes` allow-list. */
+  clientCountry: string | null;
   /** Machine identifier from `x-kilocode-machineid`. Used as the machine-
    *  cohort allocation subject for experiment routing. */
   machineId: string | null;
 };
 
 export async function getProvider(input: GetProviderInput): Promise<GetProviderResult> {
-  const { requestedModel, request, user, organizationId, taskId, clientIp, machineId } = input;
+  const {
+    requestedModel,
+    request,
+    user,
+    organizationId,
+    taskId,
+    clientIp,
+    clientCountry,
+    machineId,
+  } = input;
 
   const directByokByok = await checkDirectBYOK(user, requestedModel, organizationId);
   if (directByokByok) {
@@ -220,7 +245,7 @@ export async function getProvider(input: GetProviderInput): Promise<GetProviderR
   }
 
   if (requestedModel.startsWith(CUSTOM_LLM_PREFIX) && organizationId) {
-    const customLlmResult = await checkCustomLlm(requestedModel, organizationId);
+    const customLlmResult = await checkCustomLlm(requestedModel, organizationId, clientCountry);
     if (customLlmResult) {
       return customLlmResult;
     }

apps/web/src/lib/proxy-error-types.ts

@@ -28,6 +28,7 @@ export const proxyErrorTypeSchema = z.enum([
   'upstream_error',
   'no_free_models_available',
   'abuse_blocked',
+  'country_not_allowed',
 ]);
 
 export type ProxyErrorType = z.infer<typeof proxyErrorTypeSchema>;

packages/db/src/schema-types.ts

@@ -1516,11 +1516,20 @@ export const CustomLlmApiConfigSchema = z.object({
 
 export type CustomLlmApiConfig = z.infer<typeof CustomLlmApiConfigSchema>;
 
+// ISO 3166-1 alpha-2 country codes (e.g. "US", "GB") sourced from Vercel's
+// `x-vercel-ip-country` request header. When non-empty, requests whose
+// resolved country is not in the list are rejected at the LLM request level.
+// Case-insensitive at enforcement; entries are validated as 2 characters.
+export const CustomLlmCountryCodesSchema = z.array(z.string().length(2));
+
+export type CustomLlmCountryCodes = z.infer<typeof CustomLlmCountryCodesSchema>;
+
 export const CustomLlmDefinitionSchema = z
   .object({
     display_name: z.string(),
     api_key: z.string(),
     organization_ids: z.array(z.string()),
+    country_codes: CustomLlmCountryCodesSchema.optional(),
     pricing: CustomLlmPricingSchema.optional(),
   })
   .and(CustomLlmMetadataSchema)