From 078385d38664a8d09e885509245d26837044834d Mon Sep 17 00:00:00 2001 From: Frank Date: Mon, 29 Jun 2026 03:17:49 -0400 Subject: [PATCH] zen: budget --- .../app/src/routes/zen/util/handler.ts | 23 ++-- .../src/routes/zen/util/modelTpsLimiter.ts | 2 +- .../routes/zen/util/providerBudgetTracker.ts | 127 +++++++++++++++--- packages/console/core/src/model.ts | 2 +- 4 files changed, 122 insertions(+), 32 deletions(-) diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 9e0fcfd827..44dc5f0cae 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/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 | undefined, modelTpsLimits: Record | undefined, - providerBudgetUsage: Record | 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 diff --git a/packages/console/app/src/routes/zen/util/modelTpsLimiter.ts b/packages/console/app/src/routes/zen/util/modelTpsLimiter.ts index 3ff63f7d49..fbdab6d319 100644 --- a/packages/console/app/src/routes/zen/util/modelTpsLimiter.ts +++ b/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 () => { diff --git a/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts b/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts index c8b5eb211c..14dc7f31f2 100644 --- a/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts +++ b/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 = {} + const key = (providerId: string, priority: number, withInterval: string) => + buildRateLimitKey("provider-budget", `${providerId}:${priority}`, withInterval) + + const budgetByProvider = tracked.reduce>((acc, provider) => { + acc[provider.id] = provider.budget! + return acc + }, {}) + + const maxPriorityByProvider = tracked.reduce>((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> = {} + // Cumulative current-minute spend through each priority, per provider. + let spentThroughPriority: Record> = {} 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> = {} + const previous: Record> = {} + 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, }) }, } diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts index bb0260df18..84de249f0f 100644 --- a/packages/console/core/src/model.ts +++ b/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(),