feat(gateway): add custom BYOK upstream routing primitives

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

@@ -409,6 +409,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     skipProviderPin,
     skipKiloExclusiveModelSettings,
     experiment,
+    isByok,
   } = providerResult;
 
   // Request-level data-collection opt-out: a caller can set
@@ -439,7 +440,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     organizationId,
     projectId,
     provider: provider.id,
-    isByok: !!userByok,
+    isByok: isByok || !!userByok,
     feature,
   });
 
@@ -481,7 +482,7 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
     status_code: null,
     editor_name: extractHeaderAndLimitLength(request, 'x-kilocode-editorname'),
     machine_id: machineIdHeader,
-    user_byok: !!userByok,
+    user_byok: isByok || !!userByok,
     has_tools: (requestBodyParsed.body.tools?.length ?? 0) > 0,
     botId,
     tokenSource,
@@ -637,8 +638,8 @@ export async function POST(request: NextRequest): Promise<NextResponseType<unkno
   usageContext.status_code = response.status;
 
   // Handle OpenRouter 402 errors - don't pass them through to the client. We need to pay, not them.
-  // Skip this conversion when user BYOK is used - the 402 is about their account, not ours.
-  if (response.status === 402 && !userByok) {
+  // Skip conversion when provider-level (custom upstream) BYOK or user-level BYOK is used (`response.status === 402 && !(isByok || !!userByok)`).
+  if (response.status === 402 && !(isByok || !!userByok)) {
     await captureProxyError({
       user,
       request: requestBodyParsed.body,

apps/web/src/lib/ai-gateway/byok/custom-upstreams.ts

@@ -0,0 +1,54 @@
+import { and, eq } from 'drizzle-orm';
+import { gateway_custom_upstreams } from '@kilocode/db/schema';
+import { decryptApiKey } from '@/lib/ai-gateway/byok/encryption';
+import { BYOK_ENCRYPTION_KEY } from '@/lib/config.server';
+import type { db as drizzleDb } from '@/lib/drizzle';
+
+export type CustomUpstreamResolved = {
+  providerId: string;
+  upstreamModelId: string;
+  apiKey: string;
+  baseUrl: string;
+  extraHeaders: Record<string, string>;
+};
+
+export async function getCustomUpstreamForModel(
+  db: typeof drizzleDb,
+  owner: { organizationId?: string; userId: string },
+  requestedModel: string
+): Promise<CustomUpstreamResolved | null> {
+  const [providerId, ...rest] = requestedModel.split('/');
+  const upstreamModelId = rest.join('/');
+  if (!providerId || !upstreamModelId) return null;
+
+  const [row] = await db
+    .select()
+    .from(gateway_custom_upstreams)
+    .where(
+      and(
+        eq(gateway_custom_upstreams.provider_id, providerId),
+        eq(gateway_custom_upstreams.is_enabled, true),
+        owner.organizationId
+          ? eq(gateway_custom_upstreams.organization_id, owner.organizationId)
+          : eq(gateway_custom_upstreams.kilo_user_id, owner.userId)
+      )
+    );
+  if (!row) return null;
+  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<
+          string,
+          string
+        >)
+      : {};
+    return { providerId, upstreamModelId, apiKey, baseUrl: row.base_url, extraHeaders };
+  } catch (error) {
+    console.error('[getCustomUpstreamForModel] failed to decrypt custom upstream', {
+      providerId,
+      upstreamModelId,
+      error,
+    });
+    return null;
+  }
+}

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

@@ -11,6 +11,7 @@ import { db } from '@/lib/drizzle';
 import { eq } from 'drizzle-orm';
 import type { AnonymousUserContext } from '@/lib/anonymous';
 import { isAnonymousContext } from '@/lib/anonymous';
+import { getCustomUpstreamForModel } from '@/lib/ai-gateway/byok/custom-upstreams';
 import type { BYOKResult, Provider } from '@/lib/ai-gateway/providers/types';
 import PROVIDERS from '@/lib/ai-gateway/providers/provider-definitions';
 import { getDirectByokModel } from '@/lib/ai-gateway/providers/direct-byok';
@@ -45,6 +46,7 @@ export type GetProviderProviderResult = {
    *  by direct-byok and custom_llm2 because both already require explicit
    *  admin opt-in. */
   bypassAccessCheck: boolean;
+  isByok: boolean;
   /** Skip pinning `body.provider` from the organization-determined config.
    *  Set for direct experiment upstreams where the partner endpoint is
    *  selected by the variant, not by gateway routing. */
@@ -96,6 +98,7 @@ async function checkDirectBYOK(
     } satisfies Provider,
     userByok,
     bypassAccessCheck: true,
+    isByok: false,
   };
 }
 
@@ -132,6 +135,7 @@ async function checkCustomLlm(
     }),
     userByok: null,
     bypassAccessCheck: true,
+    isByok: true,
   };
 }
 
@@ -173,6 +177,32 @@ export type GetProviderInput = {
 export async function getProvider(input: GetProviderInput): Promise<GetProviderResult> {
   const { requestedModel, request, user, organizationId, taskId, clientIp, machineId } = input;
 
+  if (!isAnonymousContext(user)) {
+    const customUpstream = await getCustomUpstreamForModel(
+      db,
+      { userId: user.id, organizationId },
+      requestedModel
+    );
+    if (customUpstream) {
+      return {
+        kind: 'provider',
+        provider: {
+          id: 'custom',
+          apiUrl: customUpstream.baseUrl,
+          apiKey: customUpstream.apiKey,
+          supportedChatApis: ['chat_completions'],
+          transformRequest(context) {
+            context.request.body.model = customUpstream.upstreamModelId;
+            Object.assign(context.extraHeaders, customUpstream.extraHeaders);
+          },
+        },
+        userByok: null,
+        bypassAccessCheck: true,
+        isByok: true,
+      };
+    }
+  }
+
   const directByokByok = await checkDirectBYOK(user, requestedModel, organizationId);
   if (directByokByok) {
     return directByokByok;
@@ -185,6 +215,7 @@ export async function getProvider(input: GetProviderInput): Promise<GetProviderR
       provider: PROVIDERS.VERCEL_AI_GATEWAY,
       userByok: vercelByok,
       bypassAccessCheck: false,
+      isByok: false,
     };
   }
 
@@ -212,6 +243,7 @@ export async function getProvider(input: GetProviderInput): Promise<GetProviderR
         provider: buildDirectProvider(selection.upstream),
         userByok: null,
         bypassAccessCheck: false,
+        isByok: false,
         skipProviderPin: true,
         skipKiloExclusiveModelSettings: true,
         experiment: {
@@ -246,6 +278,7 @@ export async function getProvider(input: GetProviderInput): Promise<GetProviderR
       provider: PROVIDERS.VERCEL_AI_GATEWAY,
       userByok: null,
       bypassAccessCheck: false,
+      isByok: false,
     };
   }
 
@@ -256,6 +289,7 @@ export async function getProvider(input: GetProviderInput): Promise<GetProviderR
       PROVIDERS.OPENROUTER,
     userByok: null,
     bypassAccessCheck: false,
+    isByok: false,
   };
 }
 

apps/web/src/routers/models-router.ts

@@ -1,18 +1,63 @@
 import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init';
 import { preferredModels } from '@/lib/ai-gateway/models';
 import { getEnhancedOpenRouterModels } from '@/lib/ai-gateway/providers/openrouter';
+import { db } from '@/lib/drizzle';
+import { gateway_custom_upstreams, organization_memberships } from '@kilocode/db/schema';
+import { and, eq, inArray } from 'drizzle-orm';
+import { ensureOrganizationAccess } from '@/routers/organizations/utils';
+import * as z from 'zod';
 
 const preferredSet = new Set(preferredModels);
 
 export const modelsRouter = createTRPCRouter({
-  list: baseProcedure.query(async () => {
-    const response = await getEnhancedOpenRouterModels();
-
-    return (response.data ?? []).map(model => ({
-      id: model.id,
-      name: model.name,
-      supportsVision: model.architecture.input_modalities.includes('image'),
-      isPreferred: preferredSet.has(model.id),
-    }));
-  }),
+  list: baseProcedure
+    .input(z.object({ organizationId: z.string().uuid().optional() }).optional())
+    .query(async ({ ctx, input }) => {
+      const response = await getEnhancedOpenRouterModels();
+
+      const gatewayModels = (response.data ?? []).map(model => ({
+        id: model.id,
+        name: model.name,
+        supportsVision: model.architecture.input_modalities.includes('image'),
+        isPreferred: preferredSet.has(model.id),
+      }));
+
+      if (input?.organizationId) {
+        await ensureOrganizationAccess(ctx, input.organizationId);
+      }
+
+      const orgMemberships = await db
+        .select({ organization_id: organization_memberships.organization_id })
+        .from(organization_memberships)
+        .where(eq(organization_memberships.kilo_user_id, ctx.user.id));
+      const orgIds = orgMemberships.map(m => m.organization_id);
+
+      const upstreams = await db
+        .select()
+        .from(gateway_custom_upstreams)
+        .where(
+          and(
+            eq(gateway_custom_upstreams.is_enabled, true),
+            input?.organizationId
+              ? eq(gateway_custom_upstreams.organization_id, input.organizationId)
+              : orgIds.length > 0
+                ? inArray(gateway_custom_upstreams.organization_id, orgIds)
+                : eq(gateway_custom_upstreams.kilo_user_id, ctx.user.id)
+          )
+        );
+
+      const customModels = upstreams.flatMap(upstream => {
+        const mm = upstream.model_metadata as {
+          models?: Array<{ id: string; name?: string; supportsVision?: boolean }>;
+        };
+        return (mm.models ?? []).map(model => ({
+          id: `${upstream.provider_id}/${model.id}`,
+          name: model.name ?? model.id,
+          supportsVision: !!model.supportsVision,
+          isPreferred: false,
+        }));
+      });
+
+      return [...gatewayModels, ...customModels];
+    }),
 });

packages/db/src/migrations/0145_nappy_iron_man.sql

@@ -0,0 +1,24 @@
+CREATE TABLE "gateway_custom_upstreams" (
+	"id" uuid PRIMARY KEY DEFAULT pg_catalog.gen_random_uuid() NOT NULL,
+	"organization_id" uuid,
+	"kilo_user_id" text,
+	"provider_id" text NOT NULL,
+	"display_name" text NOT NULL,
+	"base_url" text NOT NULL,
+	"encrypted_api_key" jsonb NOT NULL,
+	"encrypted_extra_headers" jsonb,
+	"model_metadata" jsonb DEFAULT '{}'::jsonb NOT NULL,
+	"is_enabled" boolean DEFAULT true NOT NULL,
+	"created_at" timestamp with time zone DEFAULT now() NOT NULL,
+	"updated_at" timestamp with time zone DEFAULT now() NOT NULL,
+	"created_by" text NOT NULL,
+	CONSTRAINT "UQ_gateway_custom_upstreams_org_provider" UNIQUE("organization_id","provider_id"),
+	CONSTRAINT "UQ_gateway_custom_upstreams_user_provider" UNIQUE("kilo_user_id","provider_id"),
+	CONSTRAINT "gateway_custom_upstreams_owner_check" CHECK ((( "gateway_custom_upstreams"."kilo_user_id" IS NOT NULL AND "gateway_custom_upstreams"."organization_id" IS NULL) OR ( "gateway_custom_upstreams"."kilo_user_id" IS NULL AND "gateway_custom_upstreams"."organization_id" IS NOT NULL)))
+);
+--> statement-breakpoint
+ALTER TABLE "gateway_custom_upstreams" ADD CONSTRAINT "gateway_custom_upstreams_organization_id_organizations_id_fk" FOREIGN KEY ("organization_id") REFERENCES "public"."organizations"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "gateway_custom_upstreams" ADD CONSTRAINT "gateway_custom_upstreams_kilo_user_id_kilocode_users_id_fk" FOREIGN KEY ("kilo_user_id") REFERENCES "public"."kilocode_users"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+CREATE INDEX "IDX_gateway_custom_upstreams_organization_id" ON "gateway_custom_upstreams" USING btree ("organization_id");--> statement-breakpoint
+CREATE INDEX "IDX_gateway_custom_upstreams_kilo_user_id" ON "gateway_custom_upstreams" USING btree ("kilo_user_id");--> statement-breakpoint
+CREATE INDEX "IDX_gateway_custom_upstreams_provider_id" ON "gateway_custom_upstreams" USING btree ("provider_id");