add virtual api keys to budgets

messages/en-US.json

@@ -1697,6 +1697,8 @@
     "virtualApiKeysErrorFetchSecret": "Error loading secret",
     "virtualApiKeysErrorFetchSecretDescription": "Failed to load the virtual API key secret",
     "virtualApiKeysFilterUnassigned": "Unassigned",
+    "virtualApiKeysInferenceBudget": "Inference Budget",
+    "virtualApiKeysInferenceBudgetDescription": "Configure how this key restricts AI usage based on spending or token limits",
     "myVirtualApiKeysTitle": "Your API Keys",
     "myVirtualApiKeysDescription": "View your identity key and any virtual API keys attributed to you in this organization",
     "myVirtualApiKeysResourceTitle": "Your API Keys for {resourceName}",

server/db/pg/schema/schema.ts

@@ -1776,6 +1776,10 @@ export const aiBudgets = pgTable(
         roleId: integer("roleId").references(() => roles.roleId, {
             onDelete: "cascade"
         }),
+        virtualApiKeyId: varchar("virtualApiKeyId").references(
+            () => virtualApiKeys.virtualApiKeyId,
+            { onDelete: "cascade" }
+        ),
         amount: real("amount").notNull(),
         unit: varchar("unit").$type<"usd" | "tokens">().notNull(),
         period: varchar("period")
@@ -1806,7 +1810,12 @@ export const aiBudgets = pgTable(
             t.unit,
             t.period
         ),
-        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
+        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period),
+        unique("ai_budget_virtual_api_key_uniq").on(
+            t.virtualApiKeyId,
+            t.unit,
+            t.period
+        )
     ]
 );
 

server/db/sqlite/schema/schema.ts

@@ -1762,6 +1762,10 @@ export const aiBudgets = sqliteTable(
         roleId: integer("roleId").references(() => roles.roleId, {
             onDelete: "cascade"
         }),
+        virtualApiKeyId: text("virtualApiKeyId").references(
+            () => virtualApiKeys.virtualApiKeyId,
+            { onDelete: "cascade" }
+        ),
         amount: real("amount").notNull(),
         unit: text("unit").$type<"usd" | "tokens">().notNull(),
         period: text("period")
@@ -1794,7 +1798,12 @@ export const aiBudgets = sqliteTable(
             t.unit,
             t.period
         ),
-        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period)
+        unique("ai_budget_role_uniq").on(t.roleId, t.unit, t.period),
+        unique("ai_budget_virtual_api_key_uniq").on(
+            t.virtualApiKeyId,
+            t.unit,
+            t.period
+        )
     ]
 );
 

server/lib/aiBudgetEnforcement.ts

@@ -53,7 +53,8 @@ function applicableBudgetsCacheKey(ctx: BudgetScopeContext): string {
         ctx.requestedModel,
         ctx.resourceId ?? "",
         ctx.siteResourceId ?? "",
-        roleKey
+        roleKey,
+        ctx.virtualApiKeyId ?? ""
     ].join(":");
 }
 
@@ -83,6 +84,7 @@ export type BudgetScopeContext = {
     siteResourceId: number | null;
     roleIds: number[];
     requestUserId: string | null;
+    virtualApiKeyId: string | null;
 };
 
 /**
@@ -142,6 +144,11 @@ async function fetchApplicableBudgets(
     if (ctx.roleIds.length > 0) {
         scopeConditions.push(inArray(aiBudgets.roleId, ctx.roleIds));
     }
+    if (ctx.virtualApiKeyId != null) {
+        scopeConditions.push(
+            eq(aiBudgets.virtualApiKeyId, ctx.virtualApiKeyId)
+        );
+    }
 
     return db
         .select()
@@ -274,6 +281,17 @@ export async function sumUsageForBudget(
         );
     }
 
+    if (budget.virtualApiKeyId != null) {
+        return sumUsageAmount(
+            and(
+                eq(aiUsageRecords.orgId, ctx.orgId),
+                eq(aiUsageRecords.virtualApiKeyId, budget.virtualApiKeyId),
+                gte(aiUsageRecords.createdAt, start)
+            )!,
+            budget.unit
+        );
+    }
+
     return 0;
 }
 

server/routers/aiBudget/createAiBudget.ts

@@ -7,7 +7,8 @@ import {
     db,
     resources,
     roles,
-    siteResources
+    siteResources,
+    virtualApiKeys
 } from "@server/db";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -35,6 +36,7 @@ const bodySchema = z
         resourceId: z.coerce.number().int().positive().optional(),
         siteResourceId: z.coerce.number().int().positive().optional(),
         roleId: z.coerce.number().int().positive().optional(),
+        virtualApiKeyId: z.string().nonempty().optional(),
         amount: z.number().positive(),
         unit: aiBudgetUnitSchema,
         period: aiBudgetPeriodSchema.optional().default("monthly"),
@@ -98,6 +100,7 @@ export async function createAiBudget(
             resourceId,
             siteResourceId,
             roleId,
+            virtualApiKeyId,
             amount,
             unit,
             period,
@@ -189,6 +192,22 @@ export async function createAiBudget(
             }
         }
 
+        if (virtualApiKeyId !== undefined) {
+            const [key] = await db
+                .select({ orgId: virtualApiKeys.orgId })
+                .from(virtualApiKeys)
+                .where(eq(virtualApiKeys.virtualApiKeyId, virtualApiKeyId))
+                .limit(1);
+            if (!key || key.orgId !== orgId) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        `Virtual API key with ID ${virtualApiKeyId} not found in this organization`
+                    )
+                );
+            }
+        }
+
         const scopeCondition =
             providerId !== undefined
                 ? eq(aiBudgets.providerId, providerId)
@@ -200,14 +219,17 @@ export async function createAiBudget(
                       ? eq(aiBudgets.siteResourceId, siteResourceId)
                       : roleId !== undefined
                         ? eq(aiBudgets.roleId, roleId)
-                        : and(
+                        : virtualApiKeyId !== undefined
-                              eq(aiBudgets.orgId, orgId),
+                          ? eq(aiBudgets.virtualApiKeyId, virtualApiKeyId)
-                              isNull(aiBudgets.providerId),
+                          : and(
-                              isNull(aiBudgets.modelId),
+                                eq(aiBudgets.orgId, orgId),
-                              isNull(aiBudgets.resourceId),
+                                isNull(aiBudgets.providerId),
-                              isNull(aiBudgets.siteResourceId),
+                                isNull(aiBudgets.modelId),
-                              isNull(aiBudgets.roleId)
+                                isNull(aiBudgets.resourceId),
-                          );
+                                isNull(aiBudgets.siteResourceId),
+                                isNull(aiBudgets.roleId),
+                                isNull(aiBudgets.virtualApiKeyId)
+                            );
 
         const [existing] = await db
             .select({ budgetId: aiBudgets.budgetId })
@@ -239,6 +261,7 @@ export async function createAiBudget(
                 resourceId: resourceId ?? null,
                 siteResourceId: siteResourceId ?? null,
                 roleId: roleId ?? null,
+                virtualApiKeyId: virtualApiKeyId ?? null,
                 amount,
                 unit,
                 period,

server/routers/aiBudget/index.ts

@@ -5,6 +5,7 @@ export * from "./listAiBudgetsForModel";
 export * from "./listAiBudgetsForResource";
 export * from "./listAiBudgetsForSiteResource";
 export * from "./listAiBudgetsForRole";
+export * from "./listAiBudgetsForVirtualApiKey";
 export * from "./getAiBudget";
 export * from "./updateAiBudget";
 export * from "./deleteAiBudget";

server/routers/aiBudget/listAiBudgetsForVirtualApiKey.ts

@@ -0,0 +1,69 @@
+import { Request, Response, NextFunction } from "express";
+import { z } from "zod";
+import { aiBudgets, db } from "@server/db";
+import response from "@server/lib/response";
+import HttpCode from "@server/types/HttpCode";
+import createHttpError from "http-errors";
+import logger from "@server/logger";
+import { fromError } from "zod-validation-error";
+import { OpenAPITags, registry } from "@server/openApi";
+import { asc, eq } from "drizzle-orm";
+import type { ListAiBudgetsByScopeResponse } from "@server/routers/aiBudget/types";
+
+const paramsSchema = z.strictObject({
+    virtualApiKeyId: z.string().nonempty()
+});
+
+registry.registerPath({
+    method: "get",
+    path: "/virtual-api-key/{virtualApiKeyId}/ai-budgets",
+    description: "List AI budgets scoped to a virtual API key.",
+    tags: [OpenAPITags.AiBudget],
+    request: {
+        params: paramsSchema
+    },
+    responses: {
+        200: {
+            description: "Successful response"
+        }
+    }
+});
+
+export async function listAiBudgetsForVirtualApiKey(
+    req: Request,
+    res: Response,
+    next: NextFunction
+): Promise<any> {
+    try {
+        const parsedParams = paramsSchema.safeParse(req.params);
+        if (!parsedParams.success) {
+            return next(
+                createHttpError(
+                    HttpCode.BAD_REQUEST,
+                    fromError(parsedParams.error).toString()
+                )
+            );
+        }
+
+        const { virtualApiKeyId } = parsedParams.data;
+
+        const budgets = await db
+            .select()
+            .from(aiBudgets)
+            .where(eq(aiBudgets.virtualApiKeyId, virtualApiKeyId))
+            .orderBy(asc(aiBudgets.budgetId));
+
+        return response<ListAiBudgetsByScopeResponse>(res, {
+            data: { budgets },
+            success: true,
+            error: false,
+            message: "AI budgets retrieved successfully",
+            status: HttpCode.OK
+        });
+    } catch (error) {
+        logger.error(error);
+        return next(
+            createHttpError(HttpCode.INTERNAL_SERVER_ERROR, "An error occurred")
+        );
+    }
+}

server/routers/aiBudget/updateAiBudget.ts

@@ -7,7 +7,8 @@ import {
     db,
     resources,
     roles,
-    siteResources
+    siteResources,
+    virtualApiKeys
 } from "@server/db";
 import response from "@server/lib/response";
 import HttpCode from "@server/types/HttpCode";
@@ -34,6 +35,7 @@ const bodySchema = z.strictObject({
     resourceId: z.coerce.number().int().positive().nullable().optional(),
     siteResourceId: z.coerce.number().int().positive().nullable().optional(),
     roleId: z.coerce.number().int().positive().nullable().optional(),
+    virtualApiKeyId: z.string().nonempty().nullable().optional(),
     amount: z.number().positive().optional(),
     unit: aiBudgetUnitSchema.optional(),
     period: aiBudgetPeriodSchema.optional(),
@@ -128,6 +130,10 @@ export async function updateAiBudget(
                 : existing.siteResourceId;
         const nextRoleId =
             body.roleId !== undefined ? body.roleId : existing.roleId;
+        const nextVirtualApiKeyId =
+            body.virtualApiKeyId !== undefined
+                ? body.virtualApiKeyId
+                : existing.virtualApiKeyId;
         const nextUnit = body.unit !== undefined ? body.unit : existing.unit;
         const nextPeriod =
             body.period !== undefined ? body.period : existing.period;
@@ -138,7 +144,8 @@ export async function updateAiBudget(
                 modelId: z.number().nullable().optional(),
                 resourceId: z.number().nullable().optional(),
                 siteResourceId: z.number().nullable().optional(),
-                roleId: z.number().nullable().optional()
+                roleId: z.number().nullable().optional(),
+                virtualApiKeyId: z.string().nullable().optional()
             })
             .superRefine((data, ctx) => refineBudgetScopeFields(data, ctx))
             .safeParse({
@@ -146,7 +153,8 @@ export async function updateAiBudget(
                 modelId: nextModelId,
                 resourceId: nextResourceId,
                 siteResourceId: nextSiteResourceId,
-                roleId: nextRoleId
+                roleId: nextRoleId,
+                virtualApiKeyId: nextVirtualApiKeyId
             });
 
         if (!scopeValidation.success) {
@@ -245,6 +253,30 @@ export async function updateAiBudget(
             }
         }
 
+        if (
+            body.virtualApiKeyId !== undefined &&
+            body.virtualApiKeyId !== null
+        ) {
+            const [key] = await db
+                .select({ orgId: virtualApiKeys.orgId })
+                .from(virtualApiKeys)
+                .where(
+                    eq(
+                        virtualApiKeys.virtualApiKeyId,
+                        body.virtualApiKeyId
+                    )
+                )
+                .limit(1);
+            if (!key || key.orgId !== orgId) {
+                return next(
+                    createHttpError(
+                        HttpCode.NOT_FOUND,
+                        `Virtual API key with ID ${body.virtualApiKeyId} not found in this organization`
+                    )
+                );
+            }
+        }
+
         const scopeCondition =
             nextProviderId !== null
                 ? eq(aiBudgets.providerId, nextProviderId)
@@ -256,14 +288,20 @@ export async function updateAiBudget(
                       ? eq(aiBudgets.siteResourceId, nextSiteResourceId)
                       : nextRoleId !== null
                         ? eq(aiBudgets.roleId, nextRoleId)
-                        : and(
+                        : nextVirtualApiKeyId !== null
-                              eq(aiBudgets.orgId, orgId),
+                          ? eq(
-                              isNull(aiBudgets.providerId),
+                                aiBudgets.virtualApiKeyId,
-                              isNull(aiBudgets.modelId),
+                                nextVirtualApiKeyId
-                              isNull(aiBudgets.resourceId),
+                            )
-                              isNull(aiBudgets.siteResourceId),
+                          : and(
-                              isNull(aiBudgets.roleId)
+                                eq(aiBudgets.orgId, orgId),
-                          );
+                                isNull(aiBudgets.providerId),
+                                isNull(aiBudgets.modelId),
+                                isNull(aiBudgets.resourceId),
+                                isNull(aiBudgets.siteResourceId),
+                                isNull(aiBudgets.roleId),
+                                isNull(aiBudgets.virtualApiKeyId)
+                            );
 
         const [conflict] = await db
             .select({ budgetId: aiBudgets.budgetId })
@@ -305,6 +343,9 @@ export async function updateAiBudget(
         if (body.roleId !== undefined) {
             updateData.roleId = body.roleId;
         }
+        if (body.virtualApiKeyId !== undefined) {
+            updateData.virtualApiKeyId = body.virtualApiKeyId;
+        }
         if (body.amount !== undefined) {
             updateData.amount = body.amount;
         }

server/routers/aiBudget/validation.ts

@@ -20,6 +20,7 @@ export function refineBudgetScopeFields(
         resourceId?: number | null;
         siteResourceId?: number | null;
         roleId?: number | null;
+        virtualApiKeyId?: string | null;
     },
     ctx: z.RefinementCtx
 ) {
@@ -28,7 +29,8 @@ export function refineBudgetScopeFields(
         data.modelId,
         data.resourceId,
         data.siteResourceId,
-        data.roleId
+        data.roleId,
+        data.virtualApiKeyId
     ];
 
     const setCount = scopeFields.filter(
@@ -39,7 +41,7 @@ export function refineBudgetScopeFields(
         ctx.addIssue({
             code: "custom",
             message:
-                "Only one of providerId, modelId, resourceId, siteResourceId, or roleId may be set on a budget",
+                "Only one of providerId, modelId, resourceId, siteResourceId, roleId, or virtualApiKeyId may be set on a budget",
             path: ["providerId"]
         });
     }

server/routers/aiGateway/pipeline.ts

@@ -830,7 +830,8 @@ export async function handleAiGatewayProxy(
                 resourceId,
                 siteResourceId,
                 roleIds: requestUser?.roleIds ?? [],
-                requestUserId: requestUser?.userId ?? null
+                requestUserId: requestUser?.userId ?? null,
+                virtualApiKeyId: identity.virtualApiKeyId
             });
             appliedBudgets = budgetCheck.budgets;
 

server/routers/external.ts

@@ -1791,6 +1791,13 @@ authenticated.get(
     aiBudget.listAiBudgetsForRole
 );
 
+authenticated.get(
+    "/virtual-api-key/:virtualApiKeyId/ai-budgets",
+    verifyVirtualApiKeyAccess,
+    verifyUserHasAction(ActionsEnum.listAiBudgets),
+    aiBudget.listAiBudgetsForVirtualApiKey
+);
+
 authenticated.get(
     "/org/:orgId/labels",
     verifyOrgAccess,

src/components/CreateVirtualApiKeyForm.tsx

@@ -48,6 +48,12 @@ import {
     formatMultiResourcesSelectorLabel
 } from "@app/components/multi-resource-selector";
 import type { SelectedResource } from "@app/components/resource-selector";
+import { HorizontalTabs } from "@app/components/HorizontalTabs";
+import {
+    BudgetRowsFields,
+    getBudgetRowsErrors,
+    type BudgetRow
+} from "@app/components/BudgetsEditor";
 
 export type CreatedVirtualApiKey = {
     virtualApiKeyId: string;
@@ -93,6 +99,10 @@ export default function CreateVirtualApiKeyForm({
     const [selectedResources, setSelectedResources] = useState<
         SelectedResource[]
     >([]);
+    const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
+        []
+    );
+    const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
 
     const formSchema = z.object({
         name: z.string().min(1),
@@ -113,9 +123,29 @@ export default function CreateVirtualApiKeyForm({
         setAllResources(false);
         setSelectedUser(null);
         setSelectedResources([]);
+        setPendingBudgetRows([]);
+        setAttemptedBudgetsSave(false);
         form.reset();
     }
 
+    function handleFormSubmit(values: z.infer<typeof formSchema>) {
+        const { conflictingKeys, invalidAmountKeys } =
+            getBudgetRowsErrors(pendingBudgetRows);
+        if (conflictingKeys.size > 0 || invalidAmountKeys.size > 0) {
+            setAttemptedBudgetsSave(true);
+            toast({
+                variant: "destructive",
+                title: t("aiBudgetErrorSave"),
+                description: conflictingKeys.size
+                    ? t("aiBudgetConflictError")
+                    : t("aiBudgetInvalidAmountError")
+            });
+            return;
+        }
+
+        return onSubmit(values);
+    }
+
     async function onSubmit(values: z.infer<typeof formSchema>) {
         setLoading(true);
 
@@ -155,6 +185,33 @@ export default function CreateVirtualApiKeyForm({
                 );
             }
 
+            const pendingBudgets = pendingBudgetRows.filter(
+                (budget) => budget.amount.trim() !== ""
+            );
+            if (pendingBudgets.length > 0) {
+                try {
+                    await Promise.all(
+                        pendingBudgets.map((budget) =>
+                            api.put(`/org/${org.org.orgId}/ai-budget`, {
+                                virtualApiKeyId: key.virtualApiKeyId,
+                                amount: Number(budget.amount),
+                                unit: budget.unit,
+                                period: budget.period
+                            })
+                        )
+                    );
+                } catch (e) {
+                    toast({
+                        variant: "destructive",
+                        title: t("aiBudgetErrorSave"),
+                        description: formatAxiosError(
+                            e,
+                            t("aiBudgetErrorSave")
+                        )
+                    });
+                }
+            }
+
             const resourceLookup = new Map(
                 selectedResources.map((r) => [
                     r.resourceId,
@@ -219,10 +276,26 @@ export default function CreateVirtualApiKeyForm({
                         {!credential && (
                             <Form {...form}>
                                 <form
-                                    onSubmit={form.handleSubmit(onSubmit)}
+                                    onSubmit={form.handleSubmit(
+                                        handleFormSubmit
+                                    )}
                                     className="space-y-4"
                                     id="virtual-api-key-form"
                                 >
+                                    <HorizontalTabs
+                                        clientSide={true}
+                                        defaultTab={0}
+                                        items={[
+                                            { title: t("general"), href: "#" },
+                                            {
+                                                title: t(
+                                                    "virtualApiKeysInferenceBudget"
+                                                ),
+                                                href: "#"
+                                            }
+                                        ]}
+                                    >
+                                    <div className="space-y-4 mt-4">
                                     <FormField
                                         control={form.control}
                                         name="name"
@@ -392,6 +465,23 @@ export default function CreateVirtualApiKeyForm({
                                             </div>
                                         )}
                                     </div>
+                                    </div>
+
+                                    <div className="space-y-4 mt-4">
+                                        <p className="text-sm text-muted-foreground">
+                                            {t(
+                                                "virtualApiKeysInferenceBudgetDescription"
+                                            )}
+                                        </p>
+                                        <BudgetRowsFields
+                                            rows={pendingBudgetRows}
+                                            onChange={setPendingBudgetRows}
+                                            attemptedSave={
+                                                attemptedBudgetsSave
+                                            }
+                                        />
+                                    </div>
+                                    </HorizontalTabs>
                                 </form>
                             </Form>
                         )}

src/components/EditVirtualApiKeyForm.tsx

@@ -50,6 +50,16 @@ import { getUserDisplayName } from "@app/lib/getUserDisplayName";
 import CopyTextBox from "@app/components/CopyTextBox";
 import type { CreatedVirtualApiKey } from "@app/components/CreateVirtualApiKeyForm";
 import type { GetVirtualApiKeyResponse } from "@server/routers/virtualApiKey/types";
+import { HorizontalTabs } from "@app/components/HorizontalTabs";
+import {
+    BudgetRowsFields,
+    getBudgetRowsErrors,
+    rowsFromBudgets,
+    saveBudgetRows,
+    type BudgetRow
+} from "@app/components/BudgetsEditor";
+import { aiBudgetQueries } from "@app/lib/queries";
+import { useQuery, useQueryClient } from "@tanstack/react-query";
 
 type FormProps = {
     open: boolean;
@@ -93,6 +103,7 @@ export default function EditVirtualApiKeyForm({
     const { env } = useEnvContext();
     const api = createApiClient({ env });
     const t = useTranslations();
+    const queryClient = useQueryClient();
 
     const [loading, setLoading] = useState(false);
     const [selectedUser, setSelectedUser] = useState<SelectedUser | null>(null);
@@ -101,6 +112,19 @@ export default function EditVirtualApiKeyForm({
     >([]);
     const [credential, setCredential] = useState<string | null>(null);
     const [credentialLoading, setCredentialLoading] = useState(false);
+    const [pendingBudgetRows, setPendingBudgetRows] = useState<BudgetRow[]>(
+        []
+    );
+    const [attemptedBudgetsSave, setAttemptedBudgetsSave] = useState(false);
+
+    const budgetScope = {
+        type: "virtualApiKey" as const,
+        id: virtualApiKey?.virtualApiKeyId ?? ""
+    };
+    const budgetsQuery = useQuery({
+        ...aiBudgetQueries.scoped({ scope: budgetScope }),
+        enabled: open && !!virtualApiKey
+    });
 
     const formSchema = z
         .object({
@@ -191,6 +215,32 @@ export default function EditVirtualApiKeyForm({
         };
     }, [open, virtualApiKey, form]);
 
+    useEffect(() => {
+        if (!open || !budgetsQuery.data) {
+            return;
+        }
+        setPendingBudgetRows(rowsFromBudgets(budgetsQuery.data));
+        setAttemptedBudgetsSave(false);
+    }, [open, budgetsQuery.data]);
+
+    function handleFormSubmit(values: z.infer<typeof formSchema>) {
+        const { conflictingKeys, invalidAmountKeys } =
+            getBudgetRowsErrors(pendingBudgetRows);
+        if (conflictingKeys.size > 0 || invalidAmountKeys.size > 0) {
+            setAttemptedBudgetsSave(true);
+            toast({
+                variant: "destructive",
+                title: t("aiBudgetErrorSave"),
+                description: conflictingKeys.size
+                    ? t("aiBudgetConflictError")
+                    : t("aiBudgetInvalidAmountError")
+            });
+            return;
+        }
+
+        return onSubmit(values);
+    }
+
     async function onSubmit(values: z.infer<typeof formSchema>) {
         if (!virtualApiKey) {
             return;
@@ -223,6 +273,26 @@ export default function EditVirtualApiKeyForm({
 
         if (res?.data.data.virtualApiKey) {
             const key = res.data.data.virtualApiKey;
+
+            try {
+                await saveBudgetRows({
+                    api,
+                    orgId: virtualApiKey.orgId,
+                    scope: budgetScope,
+                    existingBudgets: budgetsQuery.data ?? [],
+                    rows: pendingBudgetRows
+                });
+                await queryClient.invalidateQueries(
+                    aiBudgetQueries.scoped({ scope: budgetScope })
+                );
+            } catch (e) {
+                toast({
+                    variant: "destructive",
+                    title: t("aiBudgetErrorSave"),
+                    description: formatAxiosError(e, t("aiBudgetErrorSave"))
+                });
+            }
+
             const resourceLookup = new Map(
                 selectedResources.map((r) => [
                     r.resourceId,
@@ -280,10 +350,26 @@ export default function EditVirtualApiKeyForm({
                     <div className="flex flex-col gap-y-4 px-1">
                         <Form {...form}>
                             <form
-                                onSubmit={form.handleSubmit(onSubmit)}
+                                onSubmit={form.handleSubmit(
+                                    handleFormSubmit
+                                )}
                                 className="space-y-4"
                                 id="edit-virtual-api-key-form"
                             >
+                                <HorizontalTabs
+                                    clientSide={true}
+                                    defaultTab={0}
+                                    items={[
+                                        { title: t("general"), href: "#" },
+                                        {
+                                            title: t(
+                                                "virtualApiKeysInferenceBudget"
+                                            ),
+                                            href: "#"
+                                        }
+                                    ]}
+                                >
+                                <div className="space-y-4 mt-4">
                                 <div className="space-y-2">
                                     <Label>
                                         {t(
@@ -430,6 +516,22 @@ export default function EditVirtualApiKeyForm({
                                         </div>
                                     )}
                                 </div>
+                                </div>
+
+                                <div className="space-y-4 mt-4">
+                                    <p className="text-sm text-muted-foreground">
+                                        {t(
+                                            "virtualApiKeysInferenceBudgetDescription"
+                                        )}
+                                    </p>
+                                    <BudgetRowsFields
+                                        rows={pendingBudgetRows}
+                                        onChange={setPendingBudgetRows}
+                                        disabled={budgetsQuery.isLoading}
+                                        attemptedSave={attemptedBudgetsSave}
+                                    />
+                                </div>
+                                </HorizontalTabs>
                             </form>
                         </Form>
                     </div>

src/lib/aiBudgetScope.ts

@@ -5,11 +5,12 @@ export type AiBudgetScopeType =
     | "model"
     | "resource"
     | "siteResource"
-    | "role";
+    | "role"
+    | "virtualApiKey";
 
 export type AiBudgetScope = {
     type: AiBudgetScopeType;
-    id: number;
+    id: number | string;
 };
 
 export type AiBudgetScopeBodyField =
@@ -17,11 +18,15 @@ export type AiBudgetScopeBodyField =
     | "modelId"
     | "resourceId"
     | "siteResourceId"
-    | "roleId";
+    | "roleId"
+    | "virtualApiKeyId";
 
 const scopeConfig: Record<
     AiBudgetScopeType,
-    { listPath: (id: number) => string; bodyField: AiBudgetScopeBodyField }
+    {
+        listPath: (id: number | string) => string;
+        bodyField: AiBudgetScopeBodyField;
+    }
 > = {
     provider: {
         listPath: (id) => `/ai-provider/${id}/ai-budgets`,
@@ -42,6 +47,10 @@ const scopeConfig: Record<
     role: {
         listPath: (id) => `/role/${id}/ai-budgets`,
         bodyField: "roleId"
+    },
+    virtualApiKey: {
+        listPath: (id) => `/virtual-api-key/${id}/ai-budgets`,
+        bodyField: "virtualApiKeyId"
     }
 };