feat(custom-models): per-key enterprise model catalog

apps/api/src/routes/index.ts

@@ -8,6 +8,7 @@ import { auditLogs } from "./audit-logs.js";
 import { chatPlans } from "./chat-plans.js";
 import { chat } from "./chat.js";
 import { chats } from "./chats.js";
+import { customModels } from "./custom-models.js";
 import { devPlanCancellationFeedback } from "./dev-plan-cancellation-feedback.js";
 import { devPlans } from "./dev-plans.js";
 import { guardrails } from "./guardrails.js";
@@ -73,4 +74,5 @@ routes.route("/audit-logs", auditLogs);
 routes.route("/model-ratings", modelRatings);
 routes.route("/guardrails", guardrails);
 routes.route("/routing-config", routingConfig);
+routes.route("/custom-models", customModels);
 routes.route("/video", video);

apps/api/src/routes/keys-provider.ts

@@ -7,7 +7,8 @@ import { getActiveUserOrganizationIds } from "@/utils/authorization.js";
 
 import { validateProviderKey } from "@llmgateway/actions";
 import { logAuditEvent } from "@llmgateway/audit";
-import { db, eq, tables } from "@llmgateway/db";
+import { invalidateSwrByTables } from "@llmgateway/cache";
+import { db, eq, getTableName, tables } from "@llmgateway/db";
 import { logger } from "@llmgateway/logger";
 import { providers } from "@llmgateway/models";
 
@@ -16,6 +17,8 @@ import type { ProviderId } from "@llmgateway/models";
 
 export const keysProvider = new OpenAPIHono<ServerTypes>();
 
+const providerKeyTableName = getTableName(tables.providerKey);
+
 // Create a schema for provider key responses
 // Using z.object directly instead of createSelectSchema due to compatibility issues
 const providerKeySchema = z.object({
@@ -61,6 +64,7 @@ const providerKeySchema = z.object({
 		})
 		.nullable(),
 	status: z.enum(["active", "inactive", "deleted"]).nullable(),
+	customModelsOnly: z.boolean(),
 	organizationId: z.string(),
 });
 
@@ -122,10 +126,16 @@ const createProviderKeySchema = z.object({
 	organizationId: z.string().min(1, "Organization ID is required"),
 });
 
-// Schema for updating a provider key status
-const updateProviderKeyStatusSchema = z.object({
-	status: z.enum(["active", "inactive"]),
-});
+// Schema for updating a provider key status / settings
+const updateProviderKeyStatusSchema = z
+	.object({
+		status: z.enum(["active", "inactive"]).optional(),
+		// Custom providers only: restrict requests to catalog-defined models.
+		customModelsOnly: z.boolean().optional(),
+	})
+	.refine((v) => v.status !== undefined || v.customModelsOnly !== undefined, {
+		message: "No updatable fields provided",
+	});
 
 // Create a new provider key
 const create = createRoute({
@@ -318,6 +328,10 @@ keysProvider.openapi(create, async (c) => {
 		},
 	});
 
+	// The gateway caches provider keys via SWR; invalidate so a newly added key
+	// is usable immediately.
+	await invalidateSwrByTables([providerKeyTableName]);
+
 	return c.json({
 		providerKey: {
 			...providerKey,
@@ -536,6 +550,10 @@ keysProvider.openapi(deleteKey, async (c) => {
 		},
 	});
 
+	// The gateway caches provider keys via SWR; invalidate so a deleted key
+	// stops being used immediately.
+	await invalidateSwrByTables([providerKeyTableName]);
+
 	return c.json({
 		message: "Provider key deleted successfully",
 	});
@@ -606,7 +624,7 @@ keysProvider.openapi(updateStatus, async (c) => {
 	}
 
 	const { id } = c.req.param();
-	const { status } = c.req.valid("json");
+	const { status, customModelsOnly } = c.req.valid("json");
 
 	// Get all active organization IDs the user has access to
 	const organizationIds = await getActiveUserOrganizationIds(user.id);
@@ -621,6 +639,9 @@ keysProvider.openapi(updateStatus, async (c) => {
 				in: organizationIds,
 			},
 		},
+		with: {
+			organization: true,
+		},
 	});
 
 	if (!providerKey) {
@@ -629,16 +650,53 @@ keysProvider.openapi(updateStatus, async (c) => {
 		});
 	}
 
-	// Update the provider key status
+	if (customModelsOnly !== undefined) {
+		if (providerKey.provider !== "custom") {
+			throw new HTTPException(400, {
+				message: "customModelsOnly can only be set on custom provider keys",
+			});
+		}
+		// Restricting to a custom catalog is an enterprise feature.
+		if (providerKey.organization?.plan !== "enterprise") {
+			throw new HTTPException(403, {
+				message: "Custom models require an enterprise plan",
+			});
+		}
+	}
+
+	const updates: {
+		status?: "active" | "inactive";
+		customModelsOnly?: boolean;
+	} = {};
+	if (status !== undefined) {
+		updates.status = status;
+	}
+	if (customModelsOnly !== undefined) {
+		updates.customModelsOnly = customModelsOnly;
+	}
+
+	// Update the provider key
 	const [updatedProviderKey] = await db
 		.update(tables.providerKey)
-		.set({
-			status,
-		})
+		.set(updates)
 		.where(eq(tables.providerKey.id, id))
 		.returning();
 
-	if (providerKey.status !== status) {
+	const changes: Record<string, { old: unknown; new: unknown }> = {};
+	if (status !== undefined && providerKey.status !== status) {
+		changes.status = { old: providerKey.status, new: status };
+	}
+	if (
+		customModelsOnly !== undefined &&
+		providerKey.customModelsOnly !== customModelsOnly
+	) {
+		changes.customModelsOnly = {
+			old: providerKey.customModelsOnly,
+			new: customModelsOnly,
+		};
+	}
+
+	if (Object.keys(changes).length > 0) {
 		await logAuditEvent({
 			organizationId: providerKey.organizationId,
 			userId: user.id,
@@ -647,15 +705,16 @@ keysProvider.openapi(updateStatus, async (c) => {
 			resourceId: id,
 			metadata: {
 				provider: providerKey.provider,
-				changes: {
-					status: { old: providerKey.status, new: status },
-				},
+				changes,
 			},
 		});
+		// The gateway caches provider keys (incl. customModelsOnly) via SWR;
+		// invalidate so status/restriction changes take effect promptly.
+		await invalidateSwrByTables([providerKeyTableName]);
 	}
 
 	return c.json({
-		message: `Provider key status updated to ${status}`,
+		message: "Provider key updated",
 		providerKey: {
 			...updatedProviderKey,
 			maskedToken: maskToken(updatedProviderKey.token),