Allow budgets to be set on the resources

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<void> {
+    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<string>();
+    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));
+        }
+    }
+}

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)

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

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<typeof TargetSchema>;
 export type Resource = z.infer<typeof PublicResourceSchema>;
 export type Config = z.infer<typeof ConfigSchema>;
 export type BlueprintResourcePolicy = z.infer<typeof ResourcePolicySchema>;
+export type BlueprintAiBudget = z.infer<typeof AiBudgetSchema>;