@hachej/boring-core 0.1.71 → 0.1.72

package/dist/{chunk-EFM5IWTK.js → chunk-JBEKMYST.js}

@@ -13,13 +13,13 @@ function createDatabase(config) {
   if (!config.databaseUrl) {
     throw new Error("databaseUrl is required to create a database connection");
   }
-  const sql8 = postgres(config.databaseUrl, {
+  const sql10 = postgres(config.databaseUrl, {
     max: 10,
     idle_timeout: 20,
     connect_timeout: 10
   });
-  const db = drizzle(sql8);
-  return { db, sql: sql8 };
+  const db = drizzle(sql10);
+  return { db, sql: sql10 };
 }
 
 // src/server/db/migrate.ts
@@ -553,6 +553,7 @@ var schema_exports = {};
 __export(schema_exports, {
   accounts: () => accounts,
   accountsRelations: () => accountsRelations,
+  budgetReservations: () => budgetReservations,
   creditGrants: () => creditGrants,
   creditPurchases: () => creditPurchases,
   idempotencyKeys: () => idempotencyKeys,
@@ -968,16 +969,17 @@ var usageReservations = pgTable2(
     )
   ]
 );
-var modelBudgetReservations = pgTable2(
-  "boring_model_budget_reservations",
+var budgetReservations = pgTable2(
+  "boring_budget_reservations",
   {
     id: uuid2("id").default(sql3`gen_random_uuid()`).primaryKey(),
     userId: text2("user_id").notNull(),
     workspaceId: text2("workspace_id"),
     sessionId: text2("session_id"),
     runId: text2("run_id").notNull(),
-    provider: text2("provider").notNull(),
-    model: text2("model").notNull(),
+    scope: text2("scope").notNull().default("model"),
+    provider: text2("provider"),
+    model: text2("model"),
     period: text2("period").notNull(),
     amountMicros: bigint("amount_micros", { mode: "number" }).notNull(),
     status: text2("status").notNull().default("active"),
@@ -985,16 +987,23 @@ var modelBudgetReservations = pgTable2(
     expiresAt: timestamp2("expires_at").notNull()
   },
   (table) => [
-    uniqueIndex("boring_model_budget_reservations_active_user_run_idx").on(table.userId, table.runId).where(sql3`${table.status} = 'active'`),
-    index2("boring_model_budget_reservations_budget_idx").on(table.userId, table.provider, table.model, table.period, table.status),
-    index2("boring_model_budget_reservations_stale_idx").on(table.status, table.expiresAt),
-    check("boring_model_budget_reservations_amount_check", sql3`${table.amountMicros} > 0`),
+    uniqueIndex("boring_budget_reservations_active_user_run_idx").on(table.scope, table.userId, table.runId).where(sql3`${table.status} = 'active'`),
+    index2("boring_budget_reservations_budget_idx").on(table.scope, table.userId, table.provider, table.model, table.period, table.status),
+    index2("boring_budget_reservations_user_budget_idx").on(table.scope, table.userId, table.period, table.status),
+    index2("boring_budget_reservations_stale_idx").on(table.status, table.expiresAt),
+    check("boring_budget_reservations_amount_check", sql3`${table.amountMicros} > 0`),
+    check("boring_budget_reservations_scope_check", sql3`${table.scope} IN ('model', 'user')`),
     check(
-      "boring_model_budget_reservations_status_check",
+      "boring_budget_reservations_scope_shape_check",
+      sql3`(${table.scope} = 'model' AND ${table.provider} IS NOT NULL AND length(btrim(${table.provider})) > 0 AND ${table.model} IS NOT NULL AND length(btrim(${table.model})) > 0) OR (${table.scope} = 'user' AND ${table.provider} IS NULL AND ${table.model} IS NULL)`
+    ),
+    check(
+      "boring_budget_reservations_status_check",
       sql3`${table.status} IN ('active', 'settled', 'released', 'expired')`
     )
   ]
 );
+var modelBudgetReservations = budgetReservations;
 var usageLedger = pgTable2(
   "boring_usage_ledger",
   {
@@ -2289,14 +2298,33 @@ var PostgresMeteringStore = class {
         providerCostMicros: usageLedger.providerCostMicros,
         billedCostMicros: usageLedger.billedCostMicros,
         stopReason: usageLedger.stopReason,
-        reservationId: sql6`${usageLedger.metadata}->>'reservationId'`
+        reservationId: sql6`${usageLedger.metadata}->>'reservationId'`,
+        modelBudgetReservationId: sql6`${usageLedger.metadata}->>'modelBudgetReservationId'`,
+        userBudgetReservationId: sql6`${usageLedger.metadata}->>'userBudgetReservationId'`
       }).from(usageLedger).where(eq3(usageLedger.id, input.usageId)).limit(1);
       const e = existing[0];
+      if (!e) throw new Error(`usage ledger id collision for ${input.usageId}: conflicting row disappeared before idempotency check`);
       const incomingReservationId = input.metadata?.reservationId ?? null;
-      const matches = e && e.userId === input.userId && e.runId === (input.runId ?? null) && e.messageId === (input.messageId ?? null) && e.source === (input.source ?? "") && e.provider === (input.provider ?? null) && e.model === (input.model ?? null) && e.inputTokens === (input.inputTokens ?? 0) && e.outputTokens === (input.outputTokens ?? 0) && e.cacheReadTokens === (input.cacheReadTokens ?? 0) && e.cacheWriteTokens === (input.cacheWriteTokens ?? 0) && e.providerCostMicros === (input.providerCostMicros ?? 0) && e.billedCostMicros === input.billedCostMicros && e.stopReason === (input.stopReason ?? null) && e.reservationId === incomingReservationId;
+      const incomingModelBudgetReservationId = input.metadata?.modelBudgetReservationId ?? null;
+      const incomingUserBudgetReservationId = input.metadata?.userBudgetReservationId ?? null;
+      const legacyMetadataCanBeBackfilled = (existing2, incoming) => existing2 === null && incoming !== null;
+      const retryMetadataIsAbsent = (existing2, incoming) => existing2 !== null && incoming === null;
+      const metadataFieldIsIdempotent = (existing2, incoming) => existing2 === incoming || legacyMetadataCanBeBackfilled(existing2, incoming) || retryMetadataIsAbsent(existing2, incoming);
+      const exactBudgetMetadataMatch = e.modelBudgetReservationId === incomingModelBudgetReservationId && e.userBudgetReservationId === incomingUserBudgetReservationId;
+      const backfillableBudgetMetadata = !exactBudgetMetadataMatch && (legacyMetadataCanBeBackfilled(e.modelBudgetReservationId, incomingModelBudgetReservationId) || legacyMetadataCanBeBackfilled(e.userBudgetReservationId, incomingUserBudgetReservationId)) && metadataFieldIsIdempotent(e.modelBudgetReservationId, incomingModelBudgetReservationId) && metadataFieldIsIdempotent(e.userBudgetReservationId, incomingUserBudgetReservationId);
+      const budgetMetadataIsIdempotent = exactBudgetMetadataMatch || backfillableBudgetMetadata || (retryMetadataIsAbsent(e.modelBudgetReservationId, incomingModelBudgetReservationId) || retryMetadataIsAbsent(e.userBudgetReservationId, incomingUserBudgetReservationId)) && metadataFieldIsIdempotent(e.modelBudgetReservationId, incomingModelBudgetReservationId) && metadataFieldIsIdempotent(e.userBudgetReservationId, incomingUserBudgetReservationId);
+      const matches = e && e.userId === input.userId && e.runId === (input.runId ?? null) && e.messageId === (input.messageId ?? null) && e.source === (input.source ?? "") && e.provider === (input.provider ?? null) && e.model === (input.model ?? null) && e.inputTokens === (input.inputTokens ?? 0) && e.outputTokens === (input.outputTokens ?? 0) && e.cacheReadTokens === (input.cacheReadTokens ?? 0) && e.cacheWriteTokens === (input.cacheWriteTokens ?? 0) && e.providerCostMicros === (input.providerCostMicros ?? 0) && e.billedCostMicros === input.billedCostMicros && e.stopReason === (input.stopReason ?? null) && e.reservationId === incomingReservationId && budgetMetadataIsIdempotent;
       if (!matches) {
         throw new Error(`usage ledger id collision for ${input.usageId}: existing row does not match this usage (refusing to silently drop the debit or corrupt the audit trail)`);
       }
+      if (!exactBudgetMetadataMatch && backfillableBudgetMetadata) {
+        const metadataPatch = {};
+        if (legacyMetadataCanBeBackfilled(e.modelBudgetReservationId, incomingModelBudgetReservationId)) metadataPatch.modelBudgetReservationId = incomingModelBudgetReservationId;
+        if (legacyMetadataCanBeBackfilled(e.userBudgetReservationId, incomingUserBudgetReservationId)) metadataPatch.userBudgetReservationId = incomingUserBudgetReservationId;
+        await tx.update(usageLedger).set({
+          metadata: sql6`COALESCE(${usageLedger.metadata}, '{}'::jsonb) || ${JSON.stringify(metadataPatch)}::jsonb`
+        }).where(eq3(usageLedger.id, input.usageId));
+      }
       return { inserted: false };
     });
   }
@@ -2452,8 +2480,157 @@ function describeUsage(source) {
   return { kind: "usage", description: "Agent usage" };
 }
 
-// src/server/db/stores/PostgresModelBudgetStore.ts
-import { and as and3, eq as eq4, inArray as inArray2, lt as lt2, or as or2, sql as sql7 } from "drizzle-orm";
+// src/server/db/stores/PostgresBudgetReservationStore.ts
+import { and as and3, eq as eq4, inArray as inArray3, lt as lt2, or as or2, sql as sql9 } from "drizzle-orm";
+
+// src/server/db/stores/BudgetReservationSupport.ts
+import { sql as sql7 } from "drizzle-orm";
+function assertPositiveSafeInteger(name, value) {
+  if (!Number.isSafeInteger(value) || value <= 0) throw new Error(`${name} must be a positive safe integer`);
+}
+function monthPeriodUtc(now) {
+  const year = now.getUTCFullYear();
+  const month = now.getUTCMonth();
+  return {
+    period: `${year}-${String(month + 1).padStart(2, "0")}`,
+    start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
+    end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0))
+  };
+}
+function chunkIds(ids, size = 1e3) {
+  const chunks = [];
+  for (let index3 = 0; index3 < ids.length; index3 += size) chunks.push(ids.slice(index3, index3 + size));
+  return chunks;
+}
+async function lockBudgetTargets(executor, targets, keyForTarget) {
+  const keys = Array.from(new Set(targets.map(keyForTarget))).sort();
+  for (const key of keys) await executor.execute(sql7`SELECT pg_advisory_xact_lock(hashtext(${key}))`);
+}
+
+// src/server/db/stores/BudgetSpendAttribution.ts
+import { inArray as inArray2, sql as sql8 } from "drizzle-orm";
+var budgetSpendAttributionStrategies = {
+  model: { usageMicros: modelUsageMicros },
+  user: { usageMicros: userUsageMicros }
+};
+async function computeBudgetSpend(tx, input, period, start, end, eligibleLegacySources, totals) {
+  const strategy = budgetSpendAttributionStrategies[input.scope];
+  return totals(input, await strategy.usageMicros(tx, input, period, start, end, eligibleLegacySources));
+}
+async function modelUsageMicros(tx, input, period, start, end) {
+  const periodStartIso = start.toISOString();
+  const periodEndIso = end.toISOString();
+  const metadataExpr = sql8`${usageLedger.metadata}->>'modelBudgetReservationId'`;
+  const usageRows = await tx.execute(sql8`
+    SELECT COALESCE(SUM(${usageLedger.billedCostMicros}), 0)::text AS total
+    FROM ${usageLedger}
+    WHERE ${usageLedger.userId} = ${input.userId}
+      AND ${usageLedger.provider} = ${input.provider}
+      AND ${usageLedger.model} = ${input.model}
+      AND (
+        EXISTS (
+          SELECT 1 FROM ${budgetReservations} r
+          WHERE r.scope = 'model'
+            AND r.user_id = ${usageLedger.userId}
+            AND r.provider = ${usageLedger.provider}
+            AND r.model = ${usageLedger.model}
+            AND r.run_id = ${usageLedger.runId}
+            AND r.period = ${period}
+            AND r.status NOT IN ('active', 'settled')
+            AND (
+              ${metadataExpr} = r.id::text
+              OR (
+                ${metadataExpr} IS NULL
+                AND r.created_at <= ${usageLedger.createdAt}
+                AND NOT EXISTS (
+                  SELECT 1 FROM ${budgetReservations} newer
+                  WHERE newer.scope = 'model'
+                    AND newer.user_id = r.user_id
+                    AND newer.provider = r.provider
+                    AND newer.model = r.model
+                    AND newer.run_id = r.run_id
+                    AND newer.created_at <= ${usageLedger.createdAt}
+                    AND (newer.created_at > r.created_at OR (newer.created_at = r.created_at AND newer.id > r.id))
+                )
+              )
+            )
+        )
+        OR (
+          ${metadataExpr} IS NULL
+          AND ${usageLedger.createdAt} >= ${periodStartIso}::timestamp
+          AND ${usageLedger.createdAt} < ${periodEndIso}::timestamp
+          AND NOT EXISTS (
+            SELECT 1 FROM ${budgetReservations} r
+            WHERE r.scope = 'model'
+              AND r.user_id = ${usageLedger.userId}
+              AND r.provider = ${usageLedger.provider}
+              AND r.model = ${usageLedger.model}
+              AND r.run_id = ${usageLedger.runId}
+              AND r.created_at <= ${usageLedger.createdAt}
+          )
+        )
+      )
+  `);
+  return Number(usageRows[0]?.total ?? 0);
+}
+async function userUsageMicros(tx, input, period, start, end, eligibleLegacySources) {
+  const periodStartIso = start.toISOString();
+  const periodEndIso = end.toISOString();
+  const metadataExpr = sql8`${usageLedger.metadata}->>'userBudgetReservationId'`;
+  const legacySourcePredicate = inArray2(usageLedger.source, [...eligibleLegacySources]);
+  const usageRows = await tx.execute(sql8`
+    SELECT COALESCE(SUM(${usageLedger.billedCostMicros}), 0)::text AS total
+    FROM ${usageLedger}
+    WHERE ${usageLedger.userId} = ${input.userId}
+      AND (
+        EXISTS (
+          SELECT 1 FROM ${budgetReservations} r
+          WHERE r.scope = 'user'
+            AND r.id::text = ${metadataExpr}
+            AND r.user_id = ${usageLedger.userId}
+            AND r.period = ${period}
+            AND r.status NOT IN ('active', 'settled')
+        )
+        OR (
+          ${legacySourcePredicate}
+          AND ${metadataExpr} IS NULL
+          AND EXISTS (
+            SELECT 1 FROM ${budgetReservations} r
+            WHERE r.scope = 'user'
+              AND r.user_id = ${usageLedger.userId}
+              AND r.run_id = ${usageLedger.runId}
+              AND r.period = ${period}
+              AND r.status NOT IN ('active', 'settled')
+              AND r.created_at <= ${usageLedger.createdAt}
+              AND NOT EXISTS (
+                SELECT 1 FROM ${budgetReservations} newer
+                WHERE newer.scope = 'user'
+                  AND newer.user_id = r.user_id
+                  AND newer.run_id = r.run_id
+                  AND newer.created_at <= ${usageLedger.createdAt}
+                  AND (newer.created_at > r.created_at OR (newer.created_at = r.created_at AND newer.id > r.id))
+              )
+          )
+        )
+        OR (
+          ${legacySourcePredicate}
+          AND ${metadataExpr} IS NULL
+          AND ${usageLedger.createdAt} >= ${periodStartIso}::timestamp
+          AND ${usageLedger.createdAt} < ${periodEndIso}::timestamp
+          AND NOT EXISTS (
+            SELECT 1 FROM ${budgetReservations} r
+            WHERE r.scope = 'user'
+              AND r.user_id = ${usageLedger.userId}
+              AND r.run_id = ${usageLedger.runId}
+              AND r.created_at <= ${usageLedger.createdAt}
+          )
+        )
+      )
+  `);
+  return Number(usageRows[0]?.total ?? 0);
+}
+
+// src/server/db/stores/PostgresBudgetReservationStore.ts
 var ModelBudgetExceededError = class extends Error {
   constructor(usedMicros, heldMicros, budgetMicros, requestedMicros) {
     super("Budget reached for this model.");
@@ -2470,221 +2647,210 @@ var ModelBudgetExceededError = class extends Error {
   statusCode = 402;
   code = "MODEL_BUDGET_EXCEEDED";
 };
-var UPDATE_ID_BATCH_SIZE = 1e3;
-function assertPositiveSafeInteger(name, value) {
-  if (!Number.isSafeInteger(value) || value <= 0) throw new Error(`${name} must be a positive safe integer`);
+var UserBudgetExceededError = class extends Error {
+  constructor(usedMicros, heldMicros, budgetMicros, requestedMicros) {
+    super("Budget reached for this user.");
+    this.usedMicros = usedMicros;
+    this.heldMicros = heldMicros;
+    this.budgetMicros = budgetMicros;
+    this.requestedMicros = requestedMicros;
+    this.name = "UserBudgetExceededError";
+  }
+  usedMicros;
+  heldMicros;
+  budgetMicros;
+  requestedMicros;
+  statusCode = 402;
+  code = "MODEL_BUDGET_EXCEEDED";
+};
+function toBudgetReservationRow(row) {
+  if (row.scope !== "model" && row.scope !== "user") throw new Error(`invalid budget reservation scope: ${row.scope}`);
+  return { id: row.id, scope: row.scope, userId: row.userId, provider: row.provider ?? null, model: row.model ?? null, period: row.period };
 }
-function monthPeriodUtc(now) {
-  const year = now.getUTCFullYear();
-  const month = now.getUTCMonth();
-  return {
-    period: `${year}-${String(month + 1).padStart(2, "0")}`,
-    start: new Date(Date.UTC(year, month, 1, 0, 0, 0, 0)),
-    end: new Date(Date.UTC(year, month + 1, 1, 0, 0, 0, 0))
-  };
+function requireModelBudgetInput(input) {
+  if (input.scope !== "model") throw new Error("expected model budget input");
+  return input;
+}
+var budgetScopePolicies = {
+  model: {
+    metadataKey: "modelBudgetReservationId",
+    lockKey: (target) => `model-budget:${target.userId}:${target.provider ?? ""}:${target.model ?? ""}:${target.period}`,
+    assertInput: (input) => {
+      const modelInput = requireModelBudgetInput(input);
+      if (!modelInput.provider.trim() || !modelInput.model.trim()) throw new Error("reserve requires provider/model");
+    },
+    sameReservation: (existing, input, period) => {
+      const modelInput = requireModelBudgetInput(input);
+      return existing.period === period && existing.provider === modelInput.provider && existing.model === modelInput.model;
+    },
+    exceededError: ModelBudgetExceededError,
+    totalFilter: (input, period) => {
+      const modelInput = requireModelBudgetInput(input);
+      return sql9`scope = 'model' AND user_id = ${modelInput.userId} AND provider = ${modelInput.provider} AND model = ${modelInput.model} AND period = ${period}`;
+    },
+    scopeBudgetWhere: (input) => {
+      const modelInput = requireModelBudgetInput(input);
+      return [eq4(budgetReservations.provider, modelInput.provider), eq4(budgetReservations.model, modelInput.model)];
+    }
+  },
+  user: {
+    metadataKey: "userBudgetReservationId",
+    lockKey: (target) => `user-budget:${target.userId}:${target.period}`,
+    assertInput: () => {
+    },
+    sameReservation: (existing, _input, period) => existing.period === period,
+    exceededError: UserBudgetExceededError,
+    totalFilter: (input, period) => sql9`scope = 'user' AND user_id = ${input.userId} AND period = ${period}`,
+    scopeBudgetWhere: () => []
+  }
+};
+function scopePolicy(scope) {
+  return budgetScopePolicies[scope];
+}
+function budgetLockKey(target) {
+  return scopePolicy(target.scope).lockKey(target);
 }
 function requireFinishKey(input) {
-  if (input.reservationId) return eq4(modelBudgetReservations.id, input.reservationId);
+  if (input.reservationId) return and3(eq4(budgetReservations.id, input.reservationId), eq4(budgetReservations.scope, input.scope));
   if (!input.runId || !input.userId) throw new Error("finish requires reservationId or runId+userId");
-  return and3(eq4(modelBudgetReservations.runId, input.runId), eq4(modelBudgetReservations.userId, input.userId));
+  return and3(eq4(budgetReservations.runId, input.runId), eq4(budgetReservations.userId, input.userId), eq4(budgetReservations.scope, input.scope));
 }
-function chunkIds(ids) {
-  const chunks = [];
-  for (let index3 = 0; index3 < ids.length; index3 += UPDATE_ID_BATCH_SIZE) {
-    chunks.push(ids.slice(index3, index3 + UPDATE_ID_BATCH_SIZE));
-  }
-  return chunks;
+function metadataKey(scope) {
+  return scopePolicy(scope).metadataKey;
 }
-function modelBudgetAdvisoryLockKey(target) {
-  return `model-budget:${target.userId}:${target.provider}:${target.model}:${target.period}`;
+function admissionFromResults(results) {
+  const user = results.find((result) => result.scope === "user");
+  const model = results.find((result) => result.scope === "model");
+  if (!model) throw new Error("budget admission requires a model reservation");
+  const handles = user ? [user, model] : [model];
+  return { user, model };
 }
-async function lockModelBudgetTargets(executor, targets) {
-  const keys = Array.from(new Set(targets.map(modelBudgetAdvisoryLockKey))).sort();
-  for (const key of keys) {
-    await executor.execute(sql7`SELECT pg_advisory_xact_lock(hashtext(${key}))`);
-  }
+function normalizeTarget(input, period) {
+  return { scope: input.scope, userId: input.userId, provider: input.provider ?? null, model: input.model ?? null, period };
 }
-var PostgresModelBudgetStore = class {
-  constructor(db) {
+function assertReserveInput(input) {
+  if (!input.userId) throw new Error("reserve requires userId");
+  scopePolicy(input.scope).assertInput(input);
+  assertPositiveSafeInteger("budgetMicros", input.budgetMicros);
+  assertPositiveSafeInteger("holdMicros", input.holdMicros);
+  if (!Number.isSafeInteger(input.ttlSeconds) || input.ttlSeconds <= 0) throw new Error("ttlSeconds must be a positive safe integer");
+}
+var PostgresBudgetReservationStore = class {
+  constructor(db, options = {}) {
     this.db = db;
+    this.options = options;
   }
   db;
+  options;
   static monthPeriodUtc(now = /* @__PURE__ */ new Date()) {
     return monthPeriodUtc(now).period;
   }
   async sweepExpired(now = /* @__PURE__ */ new Date()) {
     return this.db.transaction(async (tx) => {
       const expired = await tx.select({
-        id: modelBudgetReservations.id,
-        userId: modelBudgetReservations.userId,
-        provider: modelBudgetReservations.provider,
-        model: modelBudgetReservations.model,
-        period: modelBudgetReservations.period
-      }).from(modelBudgetReservations).where(and3(eq4(modelBudgetReservations.status, "active"), lt2(modelBudgetReservations.expiresAt, now)));
+        id: budgetReservations.id,
+        scope: budgetReservations.scope,
+        userId: budgetReservations.userId,
+        provider: budgetReservations.provider,
+        model: budgetReservations.model,
+        period: budgetReservations.period
+      }).from(budgetReservations).where(and3(eq4(budgetReservations.status, "active"), lt2(budgetReservations.expiresAt, now)));
       if (expired.length === 0) return 0;
-      await lockModelBudgetTargets(tx, expired);
+      await lockBudgetTargets(tx, expired.map(toBudgetReservationRow), budgetLockKey);
       let count = 0;
       for (const ids of chunkIds(expired.map((row) => row.id))) {
-        const rows = await tx.update(modelBudgetReservations).set({ status: "expired" }).where(and3(
-          inArray2(modelBudgetReservations.id, ids),
-          eq4(modelBudgetReservations.status, "active"),
-          lt2(modelBudgetReservations.expiresAt, now)
-        )).returning({ id: modelBudgetReservations.id });
+        const rows = await tx.update(budgetReservations).set({ status: "expired" }).where(and3(
+          inArray3(budgetReservations.id, ids),
+          eq4(budgetReservations.status, "active"),
+          lt2(budgetReservations.expiresAt, now)
+        )).returning({ id: budgetReservations.id });
         count += rows.length;
       }
       return count;
     });
   }
   async reserve(input) {
-    if (!input.userId) throw new Error("reserve requires userId");
-    if (!input.provider || !input.model) throw new Error("reserve requires provider/model");
-    assertPositiveSafeInteger("budgetMicros", input.budgetMicros);
-    assertPositiveSafeInteger("holdMicros", input.holdMicros);
-    if (!Number.isSafeInteger(input.ttlSeconds) || input.ttlSeconds <= 0) throw new Error("ttlSeconds must be a positive safe integer");
+    assertReserveInput(input);
     const now = input.now ?? /* @__PURE__ */ new Date();
+    return this.db.transaction((tx) => this.reserveInTransaction(tx, input, now));
+  }
+  async reserveAdmission(input) {
+    return admissionFromResults(await this.reserveMany(input.user ? [input.user, input.model] : [input.model]));
+  }
+  async reserveMany(inputs) {
+    const now = inputs[0]?.now ?? /* @__PURE__ */ new Date();
+    const period = monthPeriodUtc(now).period;
+    for (const input of inputs) {
+      assertReserveInput(input);
+      if (input.now && input.now.getTime() !== now.getTime()) throw new Error("reserveMany requires all item clocks to match");
+    }
+    return this.db.transaction(async (tx) => {
+      await lockBudgetTargets(tx, inputs.map((input) => normalizeTarget(input, period)), budgetLockKey);
+      const results = [];
+      for (const input of inputs) results.push(await this.reserveInTransaction(tx, { ...input, now }, now));
+      return results;
+    });
+  }
+  async reserveInTransaction(tx, input, now) {
     const period = monthPeriodUtc(now);
     const expiresAt = new Date(now.getTime() + input.ttlSeconds * 1e3);
-    return this.db.transaction(async (tx) => {
-      const expired = await tx.select({
-        id: modelBudgetReservations.id,
-        userId: modelBudgetReservations.userId,
-        provider: modelBudgetReservations.provider,
-        model: modelBudgetReservations.model,
-        period: modelBudgetReservations.period
-      }).from(modelBudgetReservations).where(and3(
-        eq4(modelBudgetReservations.status, "active"),
-        lt2(modelBudgetReservations.expiresAt, now),
-        or2(
-          and3(
-            eq4(modelBudgetReservations.userId, input.userId),
-            eq4(modelBudgetReservations.provider, input.provider),
-            eq4(modelBudgetReservations.model, input.model),
-            eq4(modelBudgetReservations.period, period.period)
-          ),
-          and3(
-            eq4(modelBudgetReservations.userId, input.userId),
-            eq4(modelBudgetReservations.runId, input.runId)
-          )
-        )
+    const expired = await tx.select({
+      id: budgetReservations.id,
+      scope: budgetReservations.scope,
+      userId: budgetReservations.userId,
+      provider: budgetReservations.provider,
+      model: budgetReservations.model,
+      period: budgetReservations.period
+    }).from(budgetReservations).where(and3(
+      eq4(budgetReservations.status, "active"),
+      lt2(budgetReservations.expiresAt, now),
+      or2(
+        this.scopeBudgetWhere(input, period.period),
+        and3(eq4(budgetReservations.scope, input.scope), eq4(budgetReservations.userId, input.userId), eq4(budgetReservations.runId, input.runId))
+      )
+    ));
+    await lockBudgetTargets(tx, [normalizeTarget(input, period.period), ...expired.map(toBudgetReservationRow)], budgetLockKey);
+    for (const ids of chunkIds(expired.map((row) => row.id))) {
+      await tx.update(budgetReservations).set({ status: "expired" }).where(and3(
+        inArray3(budgetReservations.id, ids),
+        eq4(budgetReservations.status, "active"),
+        lt2(budgetReservations.expiresAt, now)
       ));
-      await lockModelBudgetTargets(tx, [
-        { userId: input.userId, provider: input.provider, model: input.model, period: period.period },
-        ...expired
-      ]);
-      for (const ids of chunkIds(expired.map((row) => row.id))) {
-        await tx.update(modelBudgetReservations).set({ status: "expired" }).where(and3(
-          inArray2(modelBudgetReservations.id, ids),
-          eq4(modelBudgetReservations.status, "active"),
-          lt2(modelBudgetReservations.expiresAt, now)
-        ));
-      }
-      const existing = await tx.select({
-        id: modelBudgetReservations.id,
-        provider: modelBudgetReservations.provider,
-        model: modelBudgetReservations.model,
-        period: modelBudgetReservations.period
-      }).from(modelBudgetReservations).where(and3(
-        eq4(modelBudgetReservations.status, "active"),
-        eq4(modelBudgetReservations.userId, input.userId),
-        eq4(modelBudgetReservations.runId, input.runId)
-      )).limit(1);
-      const prior = existing[0];
-      if (prior) {
-        if (prior.provider !== input.provider || prior.model !== input.model || prior.period !== period.period) {
-          throw new Error(`active model budget reservation for run ${input.runId} has conflicting tuple`);
-        }
-        return { reservationId: prior.id, created: false, period: period.period };
-      }
-      const periodStartIso = period.start.toISOString();
-      const periodEndIso = period.end.toISOString();
-      const usageRows = await tx.execute(sql7`
-        SELECT COALESCE(SUM(${usageLedger.billedCostMicros}), 0)::text AS total
-        FROM ${usageLedger}
-        WHERE ${usageLedger.userId} = ${input.userId}
-          AND ${usageLedger.provider} = ${input.provider}
-          AND ${usageLedger.model} = ${input.model}
-          AND (
-            (
-              EXISTS (
-                SELECT 1 FROM ${modelBudgetReservations} r
-                WHERE r.user_id = ${usageLedger.userId}
-                  AND r.provider = ${usageLedger.provider}
-                  AND r.model = ${usageLedger.model}
-                  AND r.run_id = ${usageLedger.runId}
-                  AND r.period = ${period.period}
-                  AND r.status NOT IN ('active', 'settled')
-                  AND (
-                    ${usageLedger.metadata}->>'modelBudgetReservationId' = r.id::text
-                    OR (
-                      ${usageLedger.metadata}->>'modelBudgetReservationId' IS NULL
-                      AND r.created_at <= ${usageLedger.createdAt}
-                      AND NOT EXISTS (
-                        SELECT 1 FROM ${modelBudgetReservations} newer
-                        WHERE newer.user_id = r.user_id
-                          AND newer.provider = r.provider
-                          AND newer.model = r.model
-                          AND newer.run_id = r.run_id
-                          AND newer.created_at <= ${usageLedger.createdAt}
-                          AND (
-                            newer.created_at > r.created_at
-                            OR (newer.created_at = r.created_at AND newer.id > r.id)
-                          )
-                      )
-                    )
-                  )
-              )
-            )
-            OR (
-              ${usageLedger.createdAt} >= ${periodStartIso}::timestamp
-              AND ${usageLedger.createdAt} < ${periodEndIso}::timestamp
-              AND ${usageLedger.metadata}->>'modelBudgetReservationId' IS NULL
-              AND NOT EXISTS (
-                SELECT 1 FROM ${modelBudgetReservations} r
-                WHERE r.user_id = ${usageLedger.userId}
-                  AND r.provider = ${usageLedger.provider}
-                  AND r.model = ${usageLedger.model}
-                  AND r.run_id = ${usageLedger.runId}
-                  AND r.created_at <= ${usageLedger.createdAt}
-              )
-            )
-          )
-      `);
-      const heldRows = await tx.execute(sql7`
-        SELECT COALESCE(SUM(amount_micros), 0)::text AS total
-        FROM ${modelBudgetReservations}
-        WHERE status = 'active'
-          AND user_id = ${input.userId}
-          AND provider = ${input.provider}
-          AND model = ${input.model}
-          AND period = ${period.period}
-      `);
-      const settledFallbackRows = await tx.execute(sql7`
-        SELECT COALESCE(SUM(amount_micros), 0)::text AS total
-        FROM ${modelBudgetReservations}
-        WHERE status = 'settled'
-          AND user_id = ${input.userId}
-          AND provider = ${input.provider}
-          AND model = ${input.model}
-          AND period = ${period.period}
-      `);
-      const usedMicros = Number(usageRows[0]?.total ?? 0) + Number(settledFallbackRows[0]?.total ?? 0);
-      const heldMicros = Number(heldRows[0]?.total ?? 0);
-      if (usedMicros + heldMicros + input.holdMicros > input.budgetMicros) {
-        throw new ModelBudgetExceededError(usedMicros, heldMicros, input.budgetMicros, input.holdMicros);
-      }
-      const inserted = await tx.insert(modelBudgetReservations).values({
-        userId: input.userId,
-        workspaceId: input.workspaceId ?? null,
-        sessionId: input.sessionId ?? null,
-        runId: input.runId,
-        provider: input.provider,
-        model: input.model,
-        period: period.period,
-        amountMicros: input.holdMicros,
-        expiresAt
-      }).returning({ id: modelBudgetReservations.id });
-      return { reservationId: inserted[0].id, created: true, period: period.period };
-    });
+    }
+    const existing = await tx.select({
+      id: budgetReservations.id,
+      provider: budgetReservations.provider,
+      model: budgetReservations.model,
+      period: budgetReservations.period
+    }).from(budgetReservations).where(and3(
+      eq4(budgetReservations.status, "active"),
+      eq4(budgetReservations.scope, input.scope),
+      eq4(budgetReservations.userId, input.userId),
+      eq4(budgetReservations.runId, input.runId)
+    )).limit(1);
+    const prior = existing[0];
+    if (prior) {
+      if (!scopePolicy(input.scope).sameReservation(prior, input, period.period)) throw new Error(`active ${input.scope} budget reservation for run ${input.runId} has conflicting target`);
+      return { scope: input.scope, reservationId: prior.id, created: false, period: period.period };
+    }
+    const { usedMicros, heldMicros } = await this.spend(tx, input, period.period, period.start, period.end);
+    if (usedMicros + heldMicros + input.holdMicros > input.budgetMicros) {
+      const ErrorCtor = scopePolicy(input.scope).exceededError;
+      throw new ErrorCtor(usedMicros, heldMicros, input.budgetMicros, input.holdMicros);
+    }
+    const inserted = await tx.insert(budgetReservations).values({
+      userId: input.userId,
+      workspaceId: input.workspaceId ?? null,
+      sessionId: input.sessionId ?? null,
+      runId: input.runId,
+      scope: input.scope,
+      provider: input.provider ?? null,
+      model: input.model ?? null,
+      period: period.period,
+      amountMicros: input.holdMicros,
+      expiresAt
+    }).returning({ id: budgetReservations.id });
+    return { scope: input.scope, reservationId: inserted[0].id, created: true, period: period.period };
   }
   async settle(input) {
     await this.finish(input, "settled");
@@ -2692,23 +2858,126 @@ var PostgresModelBudgetStore = class {
   async release(input) {
     await this.finish(input, "released");
   }
-  async finish(input, status) {
+  metadataForAdmission(admission) {
+    return Object.fromEntries(this.handlesForAdmission(admission).map((result) => [metadataKey(result.scope), result.reservationId]));
+  }
+  async releaseCreated(admission) {
+    const created = this.handlesForAdmission(admission).filter((handle) => handle.created).map((handle) => ({ scope: handle.scope, reservationId: handle.reservationId }));
+    if (created.length > 0) await this.finishMany(created, "released");
+  }
+  async finishAdmission(admission, status) {
+    await this.finishMany(this.handlesForAdmission(admission).map((handle) => ({ scope: handle.scope, reservationId: handle.reservationId })), status);
+  }
+  handlesForAdmission(admission) {
+    return admission.user ? [admission.user, admission.model] : [admission.model];
+  }
+  async finishRun(input, status) {
+    await this.finishMany(["model", "user"].map((scope) => ({ scope, runId: input.runId, userId: input.userId })), status);
+  }
+  async finishMany(inputs, status) {
     await this.db.transaction(async (tx) => {
-      const active = await tx.select({
-        id: modelBudgetReservations.id,
-        userId: modelBudgetReservations.userId,
-        provider: modelBudgetReservations.provider,
-        model: modelBudgetReservations.model,
-        period: modelBudgetReservations.period
-      }).from(modelBudgetReservations).where(and3(requireFinishKey(input), eq4(modelBudgetReservations.status, "active")));
+      const active = [];
+      for (const input of inputs) {
+        const scope = input.scope;
+        const rows = await tx.select({
+          id: budgetReservations.id,
+          scope: budgetReservations.scope,
+          userId: budgetReservations.userId,
+          provider: budgetReservations.provider,
+          model: budgetReservations.model,
+          period: budgetReservations.period
+        }).from(budgetReservations).where(and3(requireFinishKey({ ...input, scope }), eq4(budgetReservations.status, "active")));
+        active.push(...rows.map(toBudgetReservationRow));
+      }
       if (active.length === 0) return;
-      await lockModelBudgetTargets(tx, active);
-      await tx.update(modelBudgetReservations).set({ status }).where(and3(
-        inArray2(modelBudgetReservations.id, active.map((row) => row.id)),
-        eq4(modelBudgetReservations.status, "active")
-      ));
+      await lockBudgetTargets(tx, active, budgetLockKey);
+      for (const ids of chunkIds(active.map((row) => row.id))) {
+        await tx.update(budgetReservations).set({ status }).where(and3(
+          inArray3(budgetReservations.id, ids),
+          eq4(budgetReservations.status, "active")
+        ));
+      }
     });
   }
+  scopeBudgetWhere(input, period) {
+    const base = [
+      eq4(budgetReservations.scope, input.scope),
+      eq4(budgetReservations.status, "active"),
+      eq4(budgetReservations.userId, input.userId),
+      eq4(budgetReservations.period, period)
+    ];
+    base.push(...scopePolicy(input.scope).scopeBudgetWhere(input));
+    return and3(...base);
+  }
+  async reservationTotal(tx, input, period, status) {
+    const scopeFilter = scopePolicy(input.scope).totalFilter(input, period);
+    const rows = await tx.execute(sql9`
+      SELECT COALESCE(SUM(amount_micros), 0)::text AS total
+      FROM ${budgetReservations}
+      WHERE status = ${status} AND ${scopeFilter}
+    `);
+    return Number(rows[0]?.total ?? 0);
+  }
+  async spend(tx, input, period, start, end) {
+    return computeBudgetSpend(
+      tx,
+      input,
+      period,
+      start,
+      end,
+      this.options.eligibleLegacySources ?? [],
+      (budgetInput, usageMicros) => this.spendTotals(tx, budgetInput, period, usageMicros)
+    );
+  }
+  async spendTotals(tx, input, period, usageMicros) {
+    const heldMicros = await this.reservationTotal(tx, input, period, "active");
+    const settledFallbackMicros = await this.reservationTotal(tx, input, period, "settled");
+    return { usedMicros: usageMicros + settledFallbackMicros, heldMicros };
+  }
+  async finish(input, status) {
+    await this.db.transaction((tx) => this.finishInTransaction(tx, input, status));
+  }
+  async finishInTransaction(tx, input, status) {
+    const scope = input.scope;
+    const active = await tx.select({
+      id: budgetReservations.id,
+      scope: budgetReservations.scope,
+      userId: budgetReservations.userId,
+      provider: budgetReservations.provider,
+      model: budgetReservations.model,
+      period: budgetReservations.period
+    }).from(budgetReservations).where(and3(requireFinishKey({ ...input, scope }), eq4(budgetReservations.status, "active")));
+    if (active.length === 0) return;
+    await lockBudgetTargets(tx, active.map(toBudgetReservationRow), budgetLockKey);
+    await tx.update(budgetReservations).set({ status }).where(and3(
+      inArray3(budgetReservations.id, active.map((row) => row.id)),
+      eq4(budgetReservations.status, "active")
+    ));
+  }
+};
+
+// src/server/db/stores/PostgresModelBudgetStore.ts
+var PostgresModelBudgetStore = class {
+  store;
+  constructor(db) {
+    this.store = new PostgresBudgetReservationStore(db);
+  }
+  static monthPeriodUtc(now = /* @__PURE__ */ new Date()) {
+    return PostgresBudgetReservationStore.monthPeriodUtc(now);
+  }
+  sweepExpired(now = /* @__PURE__ */ new Date()) {
+    return this.store.sweepExpired(now);
+  }
+  async reserve(input) {
+    const { scope: _scope, ...result } = await this.store.reserve({ ...input, scope: "model" });
+    return result;
+  }
+  settle(input) {
+    return this.store.settle({ ...input, scope: "model" });
+  }
+  release(input) {
+    return this.store.release({ ...input, scope: "model" });
+  }
 };
 
 export {
@@ -2723,6 +2992,7 @@ export {
   creditGrants,
   creditPurchases,
   usageReservations,
+  modelBudgetReservations,
   usageLedger,
   telemetryEvents,
   schema_exports,
@@ -2735,5 +3005,7 @@ export {
   InsufficientCreditError,
   PostgresMeteringStore,
   ModelBudgetExceededError,
+  UserBudgetExceededError,
+  PostgresBudgetReservationStore,
   PostgresModelBudgetStore
 };