Allow budgets to be set on the resources

2026-09-10 04:56:36 +02:00 · 2026-08-14 11:42:33 -04:00
parent 4989d1e31a
commit c8f170d197
4 changed files with 160 additions and 2 deletions
@@ -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));
+        }
+    }
+}