diff --git a/apps/web/src/app/api/openrouter/[...path]/route.test.ts b/apps/web/src/app/api/openrouter/[...path]/route.test.ts index c7e0285614..6dd830902f 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.test.ts +++ b/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', () => { diff --git a/apps/web/src/app/api/openrouter/[...path]/route.ts b/apps/web/src/app/api/openrouter/[...path]/route.ts index 2c091299c6..980b431ff1 100644 --- a/apps/web/src/app/api/openrouter/[...path]/route.ts +++ b/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 0) { + await sleepForRulesEngineAction(rulesEngineDecision.delayMs); + } + return customLlmCountryNotAllowedResponse(effectiveModelIdLowerCased); + } effectiveProviderContext = quarantineProviderResult; diff --git a/apps/web/src/lib/ai-gateway/custom-llm/country-access.test.ts b/apps/web/src/lib/ai-gateway/custom-llm/country-access.test.ts new file mode 100644 index 0000000000..ca2aa47501 --- /dev/null +++ b/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); + }); +}); diff --git a/apps/web/src/lib/ai-gateway/custom-llm/country-access.ts b/apps/web/src/lib/ai-gateway/custom-llm/country-access.ts new file mode 100644 index 0000000000..22134e7b6c --- /dev/null +++ b/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); +} diff --git a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts b/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts index 5c947fb1a3..81bb6ecd4a 100644 --- a/apps/web/src/lib/ai-gateway/llm-proxy-helpers.ts +++ b/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( diff --git a/apps/web/src/lib/ai-gateway/providers/get-provider.ts b/apps/web/src/lib/ai-gateway/providers/get-provider.ts index 22acbd4dcb..c647df9388 100644 --- a/apps/web/src/lib/ai-gateway/providers/get-provider.ts +++ b/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 { + organizationId: string, + clientCountry: string | null +): Promise { 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 { - 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; diff --git a/packages/db/src/schema-types.ts b/packages/db/src/schema-types.ts index 635352e271..7292aa3178 100644 --- a/packages/db/src/schema-types.ts +++ b/packages/db/src/schema-types.ts @@ -1516,11 +1516,20 @@ export const CustomLlmApiConfigSchema = z.object({ export type CustomLlmApiConfig = z.infer; +// 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; + 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)