From c8f170d19771932f3eddb3438a7feea133d7def7 Mon Sep 17 00:00:00 2001 From: Owen Date: Fri, 14 Aug 2026 11:42:33 -0400 Subject: [PATCH] Allow budgets to be set on the resources --- server/lib/blueprints/aiBudgets.ts | 89 +++++++++++++++++++++++ server/lib/blueprints/privateResources.ts | 17 +++++ server/lib/blueprints/publicResources.ts | 17 +++++ server/lib/blueprints/types.ts | 39 +++++++++- 4 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 server/lib/blueprints/aiBudgets.ts diff --git a/server/lib/blueprints/aiBudgets.ts b/server/lib/blueprints/aiBudgets.ts new file mode 100644 index 000000000..0ae2fbdd4 --- /dev/null +++ b/server/lib/blueprints/aiBudgets.ts @@ -0,0 +1,89 @@ +import { eq } from "drizzle-orm"; +import { aiBudgets, Transaction } from "@server/db"; + +export type BlueprintAiBudgetInput = { + amount: number; + unit: "usd" | "tokens"; + period: + | "monthly" + | "yearly" + | "lifetime" + | "daily" + | "hourly" + | "weekly"; + enforcement: "hard" | "soft"; + enabled: boolean; +}; + +type SyncAiBudgetsInput = { + orgId: string; + trx: Transaction; + budgets: BlueprintAiBudgetInput[]; +} & ( + | { scope: "public"; resourceId: number } + | { scope: "site"; siteResourceId: number } +); + +/** + * Fully declarative: makes the resource's/site resource's AI budgets match + * exactly what the blueprint declares (omitted unit/period budgets are removed). + */ +export async function syncAiBudgets(input: SyncAiBudgetsInput): Promise { + const { orgId, trx, budgets } = input; + + const existing = await trx + .select() + .from(aiBudgets) + .where( + input.scope === "public" + ? eq(aiBudgets.resourceId, input.resourceId) + : eq(aiBudgets.siteResourceId, input.siteResourceId) + ); + + const existingByKey = new Map( + existing.map((b) => [`${b.unit}::${b.period}`, b]) + ); + + const seenKeys = new Set(); + const now = Date.now(); + + for (const budget of budgets) { + const key = `${budget.unit}::${budget.period}`; + seenKeys.add(key); + const existingBudget = existingByKey.get(key); + + if (existingBudget) { + await trx + .update(aiBudgets) + .set({ + amount: budget.amount, + enforcement: budget.enforcement, + enabled: budget.enabled, + updatedAt: now + }) + .where(eq(aiBudgets.budgetId, existingBudget.budgetId)); + } else { + await trx.insert(aiBudgets).values({ + orgId, + resourceId: input.scope === "public" ? input.resourceId : null, + siteResourceId: + input.scope === "site" ? input.siteResourceId : null, + amount: budget.amount, + unit: budget.unit, + period: budget.period, + enforcement: budget.enforcement, + enabled: budget.enabled, + createdAt: now, + updatedAt: now + }); + } + } + + for (const [key, existingBudget] of existingByKey) { + if (!seenKeys.has(key)) { + await trx + .delete(aiBudgets) + .where(eq(aiBudgets.budgetId, existingBudget.budgetId)); + } + } +} diff --git a/server/lib/blueprints/privateResources.ts b/server/lib/blueprints/privateResources.ts index 070ac68b4..bcbee7ee6 100644 --- a/server/lib/blueprints/privateResources.ts +++ b/server/lib/blueprints/privateResources.ts @@ -30,6 +30,7 @@ import { build } from "@server/build"; import { LimitId } from "../billing"; import { usageService } from "../billing/usageService"; import { syncInferenceAiConfig } from "./aiProviders"; +import { syncAiBudgets } from "./aiBudgets"; async function getDomainForSiteResource( siteResourceId: number | undefined, @@ -367,6 +368,14 @@ export async function updatePrivateResources( })) }); + await syncAiBudgets({ + orgId, + trx, + scope: "site", + siteResourceId, + budgets: resourceData["ai-budget"] + }); + await trx .delete(clientSiteResources) .where(eq(clientSiteResources.siteResourceId, siteResourceId)); @@ -668,6 +677,14 @@ export async function updatePrivateResources( })) }); + await syncAiBudgets({ + orgId, + trx, + scope: "site", + siteResourceId, + budgets: resourceData["ai-budget"] + }); + const [adminRole] = await trx .select() .from(roles) diff --git a/server/lib/blueprints/publicResources.ts b/server/lib/blueprints/publicResources.ts index 806d97e9c..4adc981ce 100644 --- a/server/lib/blueprints/publicResources.ts +++ b/server/lib/blueprints/publicResources.ts @@ -57,6 +57,7 @@ import next from "next"; import { LimitId } from "../billing"; import { usageService } from "../billing/usageService"; import { syncInferenceAiConfig } from "./aiProviders"; +import { syncAiBudgets } from "./aiBudgets"; export type PublicResourcesResults = { proxyResource: Resource; @@ -696,6 +697,14 @@ export async function updatePublicResources( }) ) }); + + await syncAiBudgets({ + orgId, + trx, + scope: "public", + resourceId: existingResource.resourceId, + budgets: resourceData["ai-budget"] || [] + }); } const existingResourceTargets = await trx @@ -1258,6 +1267,14 @@ export async function updatePublicResources( })) }); + await syncAiBudgets({ + orgId, + trx, + scope: "public", + resourceId: newResource.resourceId, + budgets: resourceData["ai-budget"] || [] + }); + await trx.insert(roleResources).values({ roleId: adminRole.roleId, resourceId: newResource.resourceId diff --git a/server/lib/blueprints/types.ts b/server/lib/blueprints/types.ts index 099656ddc..d9e9415d0 100644 --- a/server/lib/blueprints/types.ts +++ b/server/lib/blueprints/types.ts @@ -5,6 +5,11 @@ import { MaintenanceSchema } from "#dynamic/lib/blueprints/MaintenanceSchema"; import { isValidRegionId } from "@server/db/regions"; import { wildcardSubdomainSchema } from "@server/lib/schemas"; import config from "@server/lib/config"; +import { + aiBudgetEnforcementSchema, + aiBudgetPeriodSchema, + aiBudgetUnitSchema +} from "@server/routers/aiBudget/validation"; const maxmindDbPath = config.getRawConfig().server.maxmind_db_path; const maxmindAsnPath = config.getRawConfig().server.maxmind_asn_path; @@ -211,6 +216,33 @@ export const AiProviderAttachmentSchema = z } ); +export const AiBudgetSchema = z.object({ + amount: z.number().positive(), + unit: aiBudgetUnitSchema, + period: aiBudgetPeriodSchema.optional().default("monthly"), + enforcement: aiBudgetEnforcementSchema.optional().default("hard"), + enabled: z.boolean().optional().default(true) +}); + +const aiBudgetArraySchema = z.array(AiBudgetSchema).refine( + (budgets) => { + const keys = budgets.map((b) => `${b.unit}::${b.period}`); + return keys.length === new Set(keys).size; + }, + { + message: + "'ai-budget' entries must not overlap: only one budget per unit/period combination is allowed" + } +); + +// No default here: an object with only 'targets' set must remain +// recognized as a targets-only resource by isTargetsOnlyResource(). +export const AiBudgetListSchema = aiBudgetArraySchema.optional(); + +export const AiBudgetListSchemaWithDefault = aiBudgetArraySchema + .optional() + .default([]); + export const AuthDaemonSchema = z .object({ pam: z.enum(["passthrough", "push"]).optional().default("passthrough"), @@ -257,7 +289,8 @@ export const PublicResourceSchema = z "proxy-protocol": z.boolean().optional(), "proxy-protocol-version": z.int().min(1).optional(), labels: z.array(z.string().min(1)).optional(), - "ai-providers": z.array(AiProviderAttachmentSchema).optional() + "ai-providers": z.array(AiProviderAttachmentSchema).optional(), + "ai-budget": AiBudgetListSchema }) .refine( (resource) => { @@ -565,7 +598,8 @@ export const PrivateResourceSchema = z machines: z.array(z.string()).optional().default([]), labels: z.array(z.string().min(1)).optional().default([]), "auth-daemon": AuthDaemonSchema.optional(), - "ai-providers": z.array(AiProviderAttachmentSchema).optional().default([]) + "ai-providers": z.array(AiProviderAttachmentSchema).optional().default([]), + "ai-budget": AiBudgetListSchemaWithDefault }) .refine( (data) => { @@ -921,3 +955,4 @@ export type Target = z.infer; export type Resource = z.infer; export type Config = z.infer; export type BlueprintResourcePolicy = z.infer; +export type BlueprintAiBudget = z.infer;