zen: budget

This commit is contained in:
Frank 2026-06-18 14:26:49 +02:00
commit 62c746f2e8
2 changed files with 9 additions and 5 deletions

View file

@ -157,9 +157,6 @@ export async function handler(
logger.metric({ logger.metric({
provider: providerInfo.id, provider: providerInfo.id,
"provider.model": providerInfo.model, "provider.model": providerInfo.model,
...(providerBudgetUsage?.[providerInfo.id]
? { "provider.budget_usage": providerBudgetUsage?.[providerInfo.id] }
: {}),
}) })
const startTimestamp = Date.now() const startTimestamp = Date.now()

View file

@ -1,5 +1,6 @@
import { centsToMicroCents } from "@opencode-ai/console-core/util/price.js" import { centsToMicroCents } from "@opencode-ai/console-core/util/price.js"
import { buildRateLimitKey, getRedis } from "./redis" import { buildRateLimitKey, getRedis } from "./redis"
import { logger } from "./logger"
export function createProviderBudgetTracker( export function createProviderBudgetTracker(
providers: { providers: {
@ -22,13 +23,15 @@ export function createProviderBudgetTracker(
const keys = Object.fromEntries( const keys = Object.fromEntries(
tracked.map((provider) => [provider.id, buildRateLimitKey("provider-budget", provider.id, interval)]), tracked.map((provider) => [provider.id, buildRateLimitKey("provider-budget", provider.id, interval)]),
) )
let budgetUsage: Record<string, number> = {}
return { return {
check: async () => { check: async () => {
const ids = tracked.filter((provider) => provider.budgetMode === "fill").map((provider) => provider.id) const ids = tracked.map((provider) => provider.id)
if (ids.length === 0) return {} if (ids.length === 0) return {}
const values = await redis.mget<(string | number | null)[]>(ids.map((id) => keys[id])) 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)])) budgetUsage = Object.fromEntries(ids.map((id, index) => [id, Number(values[index] ?? 0)]))
return budgetUsage
}, },
track: async (provider: string, costInCent: number) => { track: async (provider: string, costInCent: number) => {
const config = tracked.find((item) => item.id === provider) const config = tracked.find((item) => item.id === provider)
@ -40,6 +43,10 @@ export function createProviderBudgetTracker(
pipeline.incrby(keys[provider], cost) pipeline.incrby(keys[provider], cost)
pipeline.expire(keys[provider], 120) pipeline.expire(keys[provider], 120)
await pipeline.exec() await pipeline.exec()
logger.metric({
"provider.budget_usage": budgetUsage[provider] + cost,
"model.budget_usage": cost,
})
}, },
} }
} }