feat(gateway): add custom BYOK upstream routing primitives#3544
feat(gateway): add custom BYOK upstream routing primitives#3544baanish wants to merge 3 commits into
Conversation
…ible-upstreams Codex/implement byok OpenAI compatible upstreams
![kilo-code-bot[bot]](https://avatars.githubusercontent.com/in/2193792?s=60&v=4){kind=small}
| input?.organizationId | ||
| ? eq(gateway_custom_upstreams.organization_id, input.organizationId) | ||
| : orgIds.length > 0 | ||
| ? inArray(gateway_custom_upstreams.organization_id, orgIds) |
There was a problem hiding this comment.
WARNING: Logic bug — user-owned upstreams are silently dropped when the user belongs to any organization.
When input?.organizationId is not provided and orgIds.length > 0, the inArray condition filters only on organization_id, meaning a user who is a member of one or more orgs will never see their personal (kilo_user_id-scoped) custom upstreams in the catalog.
The intent is likely to include both org-owned upstreams and user-owned upstreams when no organizationId is specified. Consider using an or(...) that combines both.
| .where(eq(organization_memberships.kilo_user_id, ctx.user.id)); | ||
| const orgIds = orgMemberships.map(m => m.organization_id); | ||
|
|
||
| const upstreams = await db |
There was a problem hiding this comment.
SUGGESTION: Over-fetching sensitive fields — the full select() pulls encrypted_api_key and encrypted_extra_headers which are not needed for the model catalog. Prefer selecting only the columns actually used (provider_id, model_metadata) to avoid unnecessarily loading encrypted key material into application memory on every catalog fetch.
| try { | ||
| const apiKey = decryptApiKey(row.encrypted_api_key, BYOK_ENCRYPTION_KEY); | ||
| const extraHeaders = row.encrypted_extra_headers | ||
| ? (JSON.parse(decryptApiKey(row.encrypted_extra_headers, BYOK_ENCRYPTION_KEY)) as Record< |
There was a problem hiding this comment.
WARNING: Unsafe as cast on decrypted JSON — JSON.parse(...) is cast directly to Record<string, string> without validation. If the stored encrypted_extra_headers payload was not a flat string-to-string object (e.g. nested values, numbers, or null), header injection could silently pass malformed data to the upstream. Consider adding a runtime shape check or using a Zod schema here.
|
|
||
| export type BYOKApiKey = typeof byok_api_keys.$inferSelect; | ||
|
|
||
| export const gateway_custom_upstreams = pgTable( |
There was a problem hiding this comment.
WARNING: Missing soft-delete cleanup — gateway_custom_upstreams stores kilo_user_id and encrypted API keys (encrypted_api_key, encrypted_extra_headers). The ON DELETE CASCADE FK only fires when the parent kilocode_users row is hard-deleted. softDeleteUser in apps/web/src/lib/user/index.ts anonymizes the user row (it does not delete it), so the cascade never fires and encrypted credentials remain in this table indefinitely after a user is soft-deleted. Per the repo's GDPR rule, softDeleteUser must also tx.delete(gateway_custom_upstreams).where(eq(gateway_custom_upstreams.kilo_user_id, userId)) and a corresponding test must be added to apps/web/src/lib/user/index.test.ts.
Code Review SummaryStatus: 4 Issues Found | Recommendation: Address before merge Executive SummaryEncrypted credentials stored in Overview
Issue Details (click to expand)WARNING
SUGGESTION
Other Observations (not in diff)
Files Reviewed (8 files)
Fix these issues in Kilo Cloud Reviewed by claude-sonnet-4.6 · 3,255,637 tokens Review guidance: REVIEW.md from base branch |
1 participant
Motivation
<provider_id>/<upstream_model_id>.Description
gateway_custom_upstreamswith owner scoping (user/org),provider_id,display_name,base_url,encrypted_api_key, optionalencrypted_extra_headers,model_metadata,is_enabled, timestamps, unique/index/check constraints, and a newGatewayCustomUpstreamtype. (packages/db/src/schema.ts)getCustomUpstreamForModelwhich parses<provider_id>/<upstream_model_id>ids, resolves an enabled owner-scoped upstream row, decrypts the stored API key and optional headers, and returns a resolved upstream configuration. (apps/web/src/lib/ai-gateway/byok/custom-upstreams.ts)get-provider.ts) to detect authenticated owner-scoped custom upstreams before OpenRouter/Vercel fallbacks, return aproviderentry that rewritesrequest.body.modelto the upstream model id, merges optional extra headers, and surface anisByokboolean on the provider result. (apps/web/src/lib/ai-gateway/providers/get-provider.ts)isByok || !!userByokfor abuse classification, usage context (user_byok), and OpenRouter 402 handling so custom upstreams are treated consistently with BYOK behavior. (apps/web/src/app/api/openrouter/[...path]/route.ts)/modelscatalog router to include enabled custom upstream models (frommodel_metadata) for the authenticated user/org so those models appear in the catalog. (apps/web/src/routers/models-router.ts)Testing
pnpm formatwhich completed successfully.pnpm run lintwhich completed successfully.pnpm run typecheck; the typecheck run completed (with expected workspace rollup/external warnings) after fixing local type/duplicate-property edits introduced during development.pnpm drizzle generatebut it failed in this environment becausePOSTGRES_URLis not configured, so no migration file was produced and DB migration verification could not be completed here.pnpm run test; the test suite failed due to many unrelated test/db environment failures (test DB/cleanup queries and fetch mocks) in the local environment rather than the changes introduced here, so dedicated unit/integration tests for custom upstreams were not added in this PR.Codex Task
Summary by CodeRabbit
New Features
Improvements
Chores