zen: routing logic

This commit is contained in:
Frank 2026-04-17 09:54:44 -04:00
parent cb425ac927
commit 3707e4a49c
7 changed files with 2723 additions and 6 deletions

View file

@ -45,6 +45,7 @@ import { LiteData } from "@opencode-ai/console-core/lite.js"
import { Resource } from "@opencode-ai/console-resource"
import { i18n, type Key } from "~/i18n"
import { localeFromRequest } from "~/lib/language"
import { createModelTpmLimiter } from "./modelTpmLimiter"
type ZenData = Awaited<ReturnType<typeof ZenData.list>>
type RetryOptions = {
@ -121,6 +122,8 @@ export async function handler(
const authInfo = await authenticate(modelInfo, zenApiKey)
const billingSource = validateBilling(authInfo, modelInfo)
logger.metric({ source: billingSource })
const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers)
const modelTpmLimits = await modelTpmLimiter?.check()
const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
const providerInfo = selectProvider(
@ -133,6 +136,7 @@ export async function handler(
trialProviders,
retry,
stickyProvider,
modelTpmLimits,
)
validateModelSettings(billingSource, authInfo)
updateProviderKey(authInfo, providerInfo)
@ -229,6 +233,7 @@ export async function handler(
const usageInfo = providerInfo.normalizeUsage(json.usage)
const costInfo = calculateCost(modelInfo, usageInfo)
await trialLimiter?.track(usageInfo)
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
await reload(billingSource, authInfo, costInfo)
json.cost = calculateOccurredCost(billingSource, costInfo)
@ -278,6 +283,7 @@ export async function handler(
const usageInfo = providerInfo.normalizeUsage(usage)
const costInfo = calculateCost(modelInfo, usageInfo)
await trialLimiter?.track(usageInfo)
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
await reload(billingSource, authInfo, costInfo)
const cost = calculateOccurredCost(billingSource, costInfo)
@ -433,12 +439,16 @@ export async function handler(
trialProviders: string[] | undefined,
retry: RetryOptions,
stickyProvider: string | undefined,
modelTpmLimits: Record<string, number> | undefined,
) {
const modelProvider = (() => {
// Byok is top priority b/c if user set their own API key, we should use it
// instead of using the sticky provider for the same session
if (authInfo?.provider?.credentials) {
return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
}
// Always use the same provider for the same session
if (stickyProvider) {
const provider = modelInfo.providers.find((provider) => provider.id === stickyProvider)
if (provider) return provider
@ -451,10 +461,20 @@ export async function handler(
}
if (retry.retryCount !== MAX_FAILOVER_RETRIES) {
const providers = modelInfo.providers
const allProviders = modelInfo.providers
.filter((provider) => !provider.disabled)
.filter((provider) => provider.weight !== 0)
.filter((provider) => !retry.excludeProviders.includes(provider.id))
.flatMap((provider) => Array<typeof provider>(provider.weight ?? 1).fill(provider))
.filter((provider) => {
if (!provider.tpmLimit) return true
const usage = modelTpmLimits?.[`${provider.id}/${provider.model}`] ?? 0
return usage < provider.tpmLimit * 1_000_000
})
const topPriority = Math.min(...allProviders.map((p) => p.priority))
const providers = allProviders
.filter((p) => p.priority <= topPriority)
.flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
// Use the last 4 characters of session ID to select a provider
const identifier = sessionId.length ? sessionId : ip