zen: budget

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

@@ -284,8 +284,7 @@ export async function handler(
         const costInfo = calculateCost(modelInfo, usageInfo)
         await trialLimiter?.track(usageInfo)
         await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
-        if (providerInfo.budgetPriority !== undefined)
-          await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent)
+        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)
@@ -346,12 +345,7 @@ export async function handler(
                     timestampLastByte,
                     usageInfo,
                   )
-                  if (providerInfo.budgetPriority !== undefined)
-                    await providerBudgetTracker?.track(
-                      providerInfo.id,
-                      providerInfo.budgetPriority,
-                      costInfo.totalCostInCent,
-                    )
+                  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)
@@ -518,7 +512,12 @@ export async function handler(
     stickyProviderId: string | undefined,
     modelTpmLimits: Record<string, number> | undefined,
     modelTpsLimits: Record<string, { qualify: number; unqualify: number }> | undefined,
-    providerBudget: { qualify: (providerId: string, priority: number) => boolean } | undefined,
+    providerBudget:
+      | {
+          qualify: (providerId: string, priority: number) => boolean
+          prefer: (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
@@ -581,15 +580,19 @@ export async function handler(
         const stickProvider = allProviders.find((provider) => provider.id === stickyProviderId)
         if (!stickProvider) return provider
 
-        // stick provider exists + selected provider is API type => use sticky provider
-        if (!provider.tpsGoal) return stickProvider
+        const preferBudgetProvider =
+          provider.budgetPriority !== undefined && providerBudget?.prefer(provider.id, provider.budgetPriority)
 
-        // stick provier exists + selected provider is GPU type + GPU not idle => use selected provider
-        const tps = modelTpsLimits?.[`${provider.id}/${provider.model}/${provider.tpsGoal}`] ?? {
-          qualify: 0,
-          unqualify: 0,
-        }
-        if (tps.qualify <= tps.unqualify * 3) return stickProvider
+        const preferTpsProvider = (() => {
+          if (!provider.tpsGoal) return false
+          const tps = modelTpsLimits?.[`${provider.id}/${provider.model}/${provider.tpsGoal}`] ?? {
+            qualify: 0,
+            unqualify: 0,
+          }
+          return tps.qualify > tps.unqualify * 3
+        })()
+
+        if (!preferBudgetProvider && !preferTpsProvider) return stickProvider
 
         return provider
       }

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

@@ -55,6 +55,7 @@ export function createProviderBudgetTracker(
   let effectiveBudget: Record<string, Record<number, number>> = {}
   // Cumulative current-minute spend through each priority, per provider.
   let spentThroughPriority: Record<string, Record<number, number>> = {}
+  let previousSpentThroughPriority: Record<string, Record<number, number>> = {}
 
   return {
     // Returns whether a provider at a given priority still has budget headroom.
@@ -87,6 +88,7 @@ export function createProviderBudgetTracker(
 
       effectiveBudget = {}
       spentThroughPriority = {}
+      previousSpentThroughPriority = {}
       Object.entries(maxPriorityByProvider).forEach(([providerId, maxPriority]) => {
         const providerBudget = budgetByProvider[providerId]
         if (providerBudget === undefined) return
@@ -96,11 +98,13 @@ export function createProviderBudgetTracker(
         let previousRunning = 0
         effectiveBudget[providerId] = {}
         spentThroughPriority[providerId] = {}
+        previousSpentThroughPriority[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
+          previousSpentThroughPriority[providerId][priority] = previousRunning
         })
       })
 
@@ -114,9 +118,18 @@ export function createProviderBudgetTracker(
           const spentThroughCurrentPriority = spentThroughPriority[providerId]?.[priority] ?? 0
           return spentThroughCurrentPriority < budget
         },
+        prefer: (providerId: string, priority: number) => {
+          const providerBudget = budgetByProvider[providerId]
+          if (providerBudget === undefined) return false
+          const budget = centsToMicroCents(providerBudget * 100)
+          const previousUsage = previousSpentThroughPriority[providerId]?.[priority]
+          if (previousUsage === undefined) return false
+          return previousUsage < budget * 0.8
+        },
       }
     },
-    track: async (provider: string, priority: number, costInCent: number) => {
+    track: async (provider: string, priority: number | undefined, costInCent: number) => {
+      if (priority === undefined) return
       const config = tracked.find((item) => item.id === provider && item.budgetPriority === priority)
       if (!config) return
       if (config.budgetContribution === undefined) return