add virtual api keys to budgets

This commit is contained in:
Owen
2026-08-12 13:21:46 -04:00
parent c33fa8782b
commit 47ae017f94
14 changed files with 415 additions and 32 deletions
+2
View File
@@ -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}",
+10 -1
View File
@@ -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
)
]
);
+10 -1
View File
@@ -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
)
]
);
+19 -1
View File
@@ -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;
}
+32 -9
View File
@@ -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(
eq(aiBudgets.orgId, orgId),
isNull(aiBudgets.providerId),
isNull(aiBudgets.modelId),
isNull(aiBudgets.resourceId),
isNull(aiBudgets.siteResourceId),
isNull(aiBudgets.roleId)
);
: virtualApiKeyId !== undefined
? eq(aiBudgets.virtualApiKeyId, virtualApiKeyId)
: and(
eq(aiBudgets.orgId, orgId),
isNull(aiBudgets.providerId),
isNull(aiBudgets.modelId),
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,
+1
View File
@@ -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";
@@ -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")
);
}
}
+52 -11
View File
@@ -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(
eq(aiBudgets.orgId, orgId),
isNull(aiBudgets.providerId),
isNull(aiBudgets.modelId),
isNull(aiBudgets.resourceId),
isNull(aiBudgets.siteResourceId),
isNull(aiBudgets.roleId)
);
: nextVirtualApiKeyId !== null
? eq(
aiBudgets.virtualApiKeyId,
nextVirtualApiKeyId
)
: and(
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;
}
+4 -2
View File
@@ -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"]
});
}
+2 -1
View File
@@ -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;
+7
View File
@@ -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,
+91 -1
View File
@@ -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>
)}
+103 -1
View File
@@ -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>
+13 -4
View File
@@ -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"
}
};