From 7efade2d535132d9057982eaab8773c599d4b06e Mon Sep 17 00:00:00 2001 From: Frank Date: Sun, 14 Jun 2026 01:58:28 -0400 Subject: [PATCH] zen: monitor budget --- .../app/src/routes/zen/util/handler.ts | 18 ++++++++ .../routes/zen/util/providerBudgetTracker.ts | 45 +++++++++++++++++++ packages/console/core/src/model.ts | 3 ++ 3 files changed, 66 insertions(+) create mode 100644 packages/console/app/src/routes/zen/util/providerBudgetTracker.ts diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 0f9c17944..18aeedd78 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -47,6 +47,7 @@ import { i18n, type Key } from "~/i18n" import { localeFromRequest } from "~/lib/language" import { createModelTpmLimiter } from "./modelTpmLimiter" import { createModelTpsLimiter } from "./modelTpsLimiter" +import { createProviderBudgetTracker } from "./providerBudgetTracker" import { accumulateUsage, HOT_WORKSPACES } from "./usageBatcher" type ZenData = Awaited> @@ -132,6 +133,10 @@ export async function handler( const modelTpmLimits = await modelTpmLimiter?.check() const modelTpsLimiter = createModelTpsLimiter(modelInfo.providers) const modelTpsLimits = await modelTpsLimiter?.check() + const providerBudgetTracker = createProviderBudgetTracker( + modelInfo.providers.map((provider) => ({ ...zenData.providers[provider.id], ...provider })), + ) + const providerBudgetUsage = await providerBudgetTracker?.check() const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => { const providerInfo = selectProvider( @@ -145,12 +150,16 @@ export async function handler( stickyProvider, modelTpmLimits, modelTpsLimits, + providerBudgetUsage, ) validateModelSettings(billingSource, authInfo) updateProviderKey(authInfo, providerInfo) logger.metric({ provider: providerInfo.id, "provider.model": providerInfo.model, + ...(providerBudgetUsage?.[providerInfo.id] + ? { "provider.budget_usage": providerBudgetUsage?.[providerInfo.id] } + : {}), }) const startTimestamp = Date.now() @@ -257,6 +266,7 @@ 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) await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo) await reload(billingSource, authInfo, costInfo) json.cost = calculateOccurredCost(billingSource, costInfo) @@ -317,6 +327,7 @@ export async function handler( timestampLastByte, usageInfo, ) + await providerBudgetTracker?.track(providerInfo.id, costInfo.totalCostInCent) await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo) await reload(billingSource, authInfo, costInfo) const cost = calculateOccurredCost(billingSource, costInfo) @@ -483,6 +494,7 @@ export async function handler( stickyProviderId: string | undefined, modelTpmLimits: Record | undefined, modelTpsLimits: Record | undefined, + providerBudgetUsage: Record | undefined, ) { const modelProvider = (() => { // Byok is top priority b/c if user set their own API key, we should use it @@ -505,6 +517,12 @@ export async function handler( const providers = allProviders .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) + }) .filter((provider) => { if (!provider.tpmLimit) return true const usage = modelTpmLimits?.[`${provider.id}/${provider.model}`] ?? 0 diff --git a/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts b/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts new file mode 100644 index 000000000..83cfd7046 --- /dev/null +++ b/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts @@ -0,0 +1,45 @@ +import { centsToMicroCents } from "@opencode-ai/console-core/util/price.js" +import { buildRateLimitKey, getRedis } from "./redis" + +export function createProviderBudgetTracker( + providers: { + id: string + budget?: number + budgetContribution?: number + budgetMode?: "always" | "fill" + }[], +) { + const tracked = providers.filter( + (provider) => provider.budget !== undefined && provider.budgetContribution !== undefined, + ) + if (tracked.length === 0) return undefined + + const interval = new Date() + .toISOString() + .replace(/[^0-9]/g, "") + .substring(0, 12) + const redis = getRedis() + const keys = Object.fromEntries( + tracked.map((provider) => [provider.id, buildRateLimitKey("provider-budget", provider.id, interval)]), + ) + + return { + check: async () => { + const ids = tracked.filter((provider) => provider.budgetMode === "fill").map((provider) => provider.id) + if (ids.length === 0) return {} + const values = await redis.mget<(string | number | null)[]>(ids.map((id) => keys[id])) + return Object.fromEntries(ids.map((id, index) => [id, Number(values[index] ?? 0)])) + }, + track: async (provider: string, costInCent: number) => { + const config = tracked.find((item) => item.id === provider) + if (!config) return + if (config.budgetContribution === undefined) return + const cost = centsToMicroCents(costInCent * config.budgetContribution) + if (cost <= 0) return + const pipeline = redis.pipeline() + pipeline.incrby(keys[provider], cost) + pipeline.expire(keys[provider], 120) + await pipeline.exec() + }, + } +} diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts index b0851c49f..4355c1881 100644 --- a/packages/console/core/src/model.ts +++ b/packages/console/core/src/model.ts @@ -37,6 +37,8 @@ export namespace ZenData { priority: z.number().optional(), tpmLimit: z.number().optional(), tpsGoal: z.number().optional(), + budgetMode: z.enum(["always", "fill"]).optional(), + budgetContribution: z.number().optional(), weight: z.number().optional(), disabled: z.boolean().optional(), storeModel: z.string().optional(), @@ -54,6 +56,7 @@ export namespace ZenData { payloadModifier: z.record(z.string(), z.any()).optional(), payloadMappings: z.record(z.string(), z.string()).optional(), adjustCacheUsage: z.boolean().optional(), + budget: z.number().optional(), }) const ModelsSchema = z.object({