zen: budget

packages/console/app/src/routes/zen/util/handler.ts

@@ -137,7 +137,7 @@ export async function handler(
     const providerBudgetTracker = createProviderBudgetTracker(
       modelInfo.providers.map((provider) => ({ ...zenData.providers[provider.id], ...provider })),
     )
-    const providerBudgetUsage = await providerBudgetTracker?.check()
+    const providerBudget = await providerBudgetTracker?.check()
 
     const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
       const providerInfo = selectProvider(
@@ -151,7 +151,7 @@ export async function handler(
         stickyProvider,
         modelTpmLimits,
         modelTpsLimits,
-        providerBudgetUsage,
+        providerBudget,
       )
       validateModelSettings(billingSource, authInfo)
       updateProviderKey(authInfo, providerInfo)
@@ -284,7 +284,8 @@ export async function handler(
         const costInfo = calculateCost(modelInfo, usageInfo)
         await trialLimiter?.track(usageInfo)
         await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
-        await providerBudgetTracker?.track(providerInfo.id, costInfo.totalCostInCent)
+        if (providerInfo.budgetPriority !== undefined)
+          await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent)
         await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
         await reload(billingSource, authInfo, costInfo)
         json.cost = calculateOccurredCost(billingSource, costInfo)
@@ -345,7 +346,12 @@ export async function handler(
                     timestampLastByte,
                     usageInfo,
                   )
-                  await providerBudgetTracker?.track(providerInfo.id, costInfo.totalCostInCent)
+                  if (providerInfo.budgetPriority !== undefined)
+                    await providerBudgetTracker?.track(
+                      providerInfo.id,
+                      providerInfo.budgetPriority,
+                      costInfo.totalCostInCent,
+                    )
                   await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
                   await reload(billingSource, authInfo, costInfo)
                   const cost = calculateOccurredCost(billingSource, costInfo)
@@ -512,7 +518,7 @@ export async function handler(
     stickyProviderId: string | undefined,
     modelTpmLimits: Record<string, number> | undefined,
     modelTpsLimits: Record<string, { qualify: number; unqualify: number }> | undefined,
-    providerBudgetUsage: Record<string, number> | undefined,
+    providerBudget: { qualify: (providerId: string, priority: number) => boolean } | undefined,
   ) {
     const modelProvider = (() => {
       // Byok is top priority b/c if user set their own API key, we should use it
@@ -536,10 +542,9 @@ export async function handler(
           .filter((provider) => provider.weight !== 0)
           .filter((provider) => !retry.excludeProviders.includes(provider.id))
           .filter((provider) => {
-            if (provider.budgetMode !== "fill") return true
-            const budget = zenData.providers[provider.id]?.budget
-            if (budget === undefined) return false
-            return (providerBudgetUsage?.[provider.id] ?? 0) < centsToMicroCents(budget * 100)
+            if (provider.budgetPriority === undefined) return true
+            if (!providerBudget) return true
+            return providerBudget.qualify(provider.id, provider.budgetPriority)
           })
           .filter((provider) => {
             if (!provider.tpmLimit) return true

packages/console/app/src/routes/zen/util/modelTpsLimiter.ts

@@ -20,7 +20,7 @@ export function createModelTpsLimiter(providers: { id: string; model: string; tp
     )
   const now = Date.now()
   const currInterval = toInterval(new Date(now))
-  const prevInterval = toInterval(new Date(now - 60 * 1000))
+  const prevInterval = toInterval(new Date(now - 60_000))
 
   return {
     check: async () => {

packages/console/app/src/routes/zen/util/providerBudgetTracker.ts

@@ -2,50 +2,135 @@ import { centsToMicroCents } from "@opencode-ai/console-core/util/price.js"
 import { buildRateLimitKey, getRedis } from "./redis"
 import { logger } from "./logger"
 
+// Per-provider, per-minute budget with priorities. The budget belongs to a
+// provider and is shared across every model that routes to it. Each model's
+// provider entry carries a `budgetPriority`: priority 1 ("always") routes
+// unconditionally, while higher priorities ("fill") only route while the provider's
+// current-minute spend through that priority is still under budget.
+//
+// Spend is tracked per (provider, priority, minute) so a fill priority can yield its
+// leftover headroom to the next priority down. The previous minute is also read so
+// higher priorities can reserve the next minute's budget first.
 export function createProviderBudgetTracker(
   providers: {
     id: string
     budget?: number
     budgetContribution?: number
-    budgetMode?: "always" | "fill"
+    budgetPriority?: number
   }[],
 ) {
   const tracked = providers.filter(
-    (provider) => provider.budget !== undefined && provider.budgetContribution !== undefined,
+    (provider) =>
+      provider.budget !== undefined &&
+      provider.budgetContribution !== undefined &&
+      provider.budgetPriority !== undefined,
   )
   if (tracked.length === 0) return undefined
 
-  const interval = new Date()
-    .toISOString()
-    .replace(/[^0-9]/g, "")
-    .substring(0, 12)
+  const intervalAt = (date: Date) =>
+    date
+      .toISOString()
+      .replace(/[^0-9]/g, "")
+      .substring(0, 12)
+  const now = new Date()
+  const currInterval = intervalAt(now)
+  const prevInterval = intervalAt(new Date(now.getTime() - 60_000))
+
   const redis = getRedis()
-  const keys = Object.fromEntries(
-    tracked.map((provider) => [provider.id, buildRateLimitKey("provider-budget", provider.id, interval)]),
-  )
-  let budgetUsage: Record<string, number> = {}
+  const key = (providerId: string, priority: number, withInterval: string) =>
+    buildRateLimitKey("provider-budget", `${providerId}:${priority}`, withInterval)
+
+  const budgetByProvider = tracked.reduce<Record<string, number>>((acc, provider) => {
+    acc[provider.id] = provider.budget!
+    return acc
+  }, {})
+
+  const maxPriorityByProvider = tracked.reduce<Record<string, number>>((acc, provider) => {
+    acc[provider.id] = Math.max(acc[provider.id] ?? 0, provider.budgetPriority!)
+    return acc
+  }, {})
+
+  // Effective budget in micro-cents per provider/priority, computed in check()
+  // from the configured budget minus previous-minute usage from higher priorities.
+  let effectiveBudget: Record<string, Record<number, number>> = {}
+  // Cumulative current-minute spend through each priority, per provider.
+  let spentThroughPriority: Record<string, Record<number, number>> = {}
 
   return {
+    // Returns whether a provider at a given priority still has budget headroom.
+    // Priority 1 always qualifies; higher priorities qualify only while everything through
+    // the current priority hasn't already filled the previous-minute adjusted
+    // budget.
     check: async () => {
-      const ids = tracked.map((provider) => provider.id)
-      if (ids.length === 0) return {}
-      const values = await redis.mget<(string | number | null)[]>(ids.map((id) => keys[id]))
-      budgetUsage = Object.fromEntries(ids.map((id, index) => [id, Number(values[index] ?? 0)]))
-      return budgetUsage
+      const reads = Object.entries(maxPriorityByProvider).flatMap(([providerId, maxPriority]) =>
+        Array.from({ length: maxPriority }, (_, index) => index + 1).flatMap((priority) => [
+          { providerId, priority, interval: currInterval, prev: false },
+          { providerId, priority, interval: prevInterval, prev: true },
+        ]),
+      )
+      const values = await redis.mget<(string | number | null)[]>(
+        reads.map((r) => key(r.providerId, r.priority, r.interval)),
+      )
+
+      const current: Record<string, Record<number, number>> = {}
+      const previous: Record<string, Record<number, number>> = {}
+      reads.forEach((r, index) => {
+        const amount = Number(values[index] ?? 0)
+        if (r.prev) {
+          previous[r.providerId] ??= {}
+          previous[r.providerId][r.priority] = amount
+          return
+        }
+        current[r.providerId] ??= {}
+        current[r.providerId][r.priority] = amount
+      })
+
+      effectiveBudget = {}
+      spentThroughPriority = {}
+      Object.entries(maxPriorityByProvider).forEach(([providerId, maxPriority]) => {
+        const providerBudget = budgetByProvider[providerId]
+        if (providerBudget === undefined) return
+        const budget = centsToMicroCents(providerBudget * 100)
+
+        let currentRunning = 0
+        let previousRunning = 0
+        effectiveBudget[providerId] = {}
+        spentThroughPriority[providerId] = {}
+        Array.from({ length: maxPriority }, (_, index) => index + 1).forEach((priority) => {
+          currentRunning += current[providerId]?.[priority] ?? 0
+          effectiveBudget[providerId][priority] = Math.max(0, budget - previousRunning)
+          previousRunning += previous[providerId]?.[priority] ?? 0
+          spentThroughPriority[providerId][priority] = currentRunning
+        })
+      })
+
+      return {
+        // Priority 1 is unconditional. Higher priorities gate on the spend through
+        // the current priority against the effective budget.
+        qualify: (providerId: string, priority: number) => {
+          if (priority <= 1) return true
+          const budget = effectiveBudget[providerId]?.[priority]
+          if (budget === undefined) return false
+          const spentThroughCurrentPriority = spentThroughPriority[providerId]?.[priority] ?? 0
+          return spentThroughCurrentPriority < budget
+        },
+      }
     },
-    track: async (provider: string, costInCent: number) => {
-      const config = tracked.find((item) => item.id === provider)
+    track: async (provider: string, priority: number, costInCent: number) => {
+      const config = tracked.find((item) => item.id === provider && item.budgetPriority === priority)
       if (!config) return
       if (config.budgetContribution === undefined) return
       const cost = centsToMicroCents(costInCent * config.budgetContribution)
       if (cost <= 0) return
+      const redisKey = key(provider, priority, currInterval)
       const pipeline = redis.pipeline()
-      pipeline.incrby(keys[provider], cost)
-      pipeline.expire(keys[provider], 120)
+      pipeline.incrby(redisKey, cost)
+      // Keep two minutes so the previous interval is readable for budget adjustment.
+      pipeline.expire(redisKey, 120)
       await pipeline.exec()
       logger.metric({
-        "provider.budget_usage": budgetUsage[provider] + cost,
-        "model.budget_usage": cost,
+        "provider.budget_usage": cost,
+        "provider.budget_priority": priority,
       })
     },
   }

packages/console/core/src/model.ts

@@ -37,7 +37,7 @@ export namespace ZenData {
         priority: z.number().optional(),
         tpmLimit: z.number().optional(),
         tpsGoal: z.number().optional(),
-        budgetMode: z.enum(["always", "fill"]).optional(),
+        budgetPriority: z.number().optional(),
         budgetContribution: z.number().optional(),
         weight: z.number().optional(),
         disabled: z.boolean().optional(),