feat(ai-gateway): restrict custom LLMs to allowed country codes#4288
feat(ai-gateway): restrict custom LLMs to allowed country codes#4288chrarnoldus wants to merge 1 commit into
Conversation
Add an optional `country_codes` allow-list (ISO 3166-1 alpha-2) to the custom LLM definition so admins can limit a custom LLM to requests from specific Vercel country codes. When the list is non-empty, requests whose resolved `x-vercel-ip-country` is not in the list are rejected at the LLM request level with a 403 `country_not_allowed`; an absent/empty list leaves the model unrestricted, and a configured list fails closed when the request country is unknown. No DB migration: the field lives in the existing JSONB `definition` column and is validated by `CustomLlmDefinitionSchema`. Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| if (!normalizedRequestCountry) { | ||
| return false; | ||
| } | ||
| const normalizedAllowed = new Set(allowedCountries.map(code => code.toUpperCase())); |
There was a problem hiding this comment.
WARNING: Allow-list entries are not trimmed before comparison
requestCountry is normalized with trim().toUpperCase(), but each country_codes entry is only uppercased here. A configured value like " US " or the new whitespace test added in this PR will still miss and incorrectly return forbidden, so the implementation doesn't match the documented behavior.
| const normalizedAllowed = new Set(allowedCountries.map(code => code.toUpperCase())); | |
| const normalizedAllowed = new Set(allowedCountries.map(code => code.trim().toUpperCase())); |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
| // `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)); |
There was a problem hiding this comment.
WARNING: country_codes validation accepts non-country values
The admin JSON editor validates through this schema, so entries like " U", "1*", or other non-letter pairs currently pass even though they are not valid Vercel country codes. Those values will later fail closed in isCountryAllowed and can block every request for a model by mistake. Restricting this to two letters keeps the case-insensitive behavior without accepting malformed codes.
| export const CustomLlmCountryCodesSchema = z.array(z.string().length(2)); | |
| export const CustomLlmCountryCodesSchema = z.array(z.string().regex(/^[A-Za-z]{2}$/)); |
Reply with @kilocode-bot fix it to have Kilo Code address this issue.
Code Review SummaryStatus: 2 Issues Found | Recommendation: Address before merge Executive SummaryThe new country-based custom LLM gate has two contract mismatches in validation/normalization that can either reject valid-looking allow-lists or let malformed country codes block traffic unexpectedly. Overview
Issue Details (click to expand)WARNING
Fix these issues in Kilo Cloud Files Reviewed (8 files)
Reviewed by gpt-5.4-20260305 · Input: 119.8K · Output: 13.4K · Cached: 941.8K Review guidance: REVIEW.md from base branch |
1 participant
Summary
Adds an optional
country_codesallow-list to the custom LLM definition (CustomLlmDefinitionSchema), letting admins limit a custom LLM to requests originating from specific Vercel country codes.country_codesis an optional array of ISO 3166-1 alpha-2 codes (e.g.["US", "GB"]), validated as 2-character strings. It lives in the existing JSONBdefinitioncolumn ofcustom_llm2, so no DB migration is required; the admin JSON editor validates it viadeepStrict(CustomLlmDefinitionSchema).checkCustomLlm(get-provider.ts), using the request'sx-vercel-ip-countryheader (newclientCountryinput threaded through bothgetProvidercall sites in the route handler). Matching is case-insensitive and trims whitespace.403 country_not_allowedresponse (newGetProviderResultvariant{ kind: 'forbidden' }andcustomLlmCountryNotAllowedResponsehelper +country_not_allowedproxy error type).organization_idscheck, so it only applies to orgs already permitted to use the model. Model-list output (listAvailableCustomLlms) is intentionally unchanged.Verification
pnpm testin an environment with the test DB.Visual Changes
N/A — admin custom-LLM editor already uses a raw JSON editor, so the new optional field is editable without UI changes.
Reviewer Notes
x-vercel-ip-countryis absent but a list is configured) is intentional to prevent header absence from bypassing the restriction. If an allow-list-all-when-unknown preference is desired instead, flag it for review.country_codesis access control metadata configured by admins, not user PII, and is not persisted as a new column/table, so the GDPR soft-delete flow is unaffected.country_not_allowederror type is additive to the zod enum; older clients that don't recognize it will fall back to a generic error message.