zen: rate limiter

This commit is contained in:
Frank 2026-04-14 20:29:21 -04:00
parent 2c36bf9490
commit 8df7ccc304
12 changed files with 7757 additions and 23 deletions

View file

@ -21,6 +21,7 @@ import {
MonthlyLimitError,
UserLimitError,
ModelError,
RateLimitError,
FreeUsageLimitError,
SubscriptionUsageLimitError,
} from "./error"
@ -35,7 +36,8 @@ import { anthropicHelper } from "./provider/anthropic"
import { googleHelper } from "./provider/google"
import { openaiHelper } from "./provider/openai"
import { oaCompatHelper } from "./provider/openai-compatible"
import { createRateLimiter } from "./rateLimiter"
import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter"
import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
import { createDataDumper } from "./dataDumper"
import { createTrialLimiter } from "./trialLimiter"
import { createStickyTracker } from "./stickyProviderTracker"
@ -92,6 +94,8 @@ export async function handler(
const isStream = opts.parseIsStream(url, body)
const rawIp = input.request.headers.get("x-real-ip") ?? ""
const ip = rawIp.includes(":") ? rawIp.split(":").slice(0, 4).join(":") : rawIp
const rawZenApiKey = opts.parseApiKey(input.request.headers)
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
const requestId = input.request.headers.get("x-opencode-request") ?? ""
const projectId = input.request.headers.get("x-opencode-project") ?? ""
@ -108,17 +112,13 @@ export async function handler(
const dataDumper = createDataDumper(sessionId, requestId, projectId)
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
const trialProviders = await trialLimiter?.check()
const rateLimiter = createRateLimiter(
modelInfo.id,
modelInfo.allowAnonymous,
modelInfo.rateLimit,
ip,
input.request,
)
const rateLimiter = modelInfo.allowAnonymous
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
: createKeyRateLimiter(modelInfo.id, zenApiKey, input.request)
await rateLimiter?.check()
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
const stickyProvider = await stickyTracker?.get()
const authInfo = await authenticate(modelInfo)
const authInfo = await authenticate(modelInfo, zenApiKey)
const billingSource = validateBilling(authInfo, modelInfo)
logger.metric({ source: billingSource })
@ -363,7 +363,11 @@ export async function handler(
{ status: 401 },
)
if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) {
if (
error instanceof RateLimitError ||
error instanceof FreeUsageLimitError ||
error instanceof SubscriptionUsageLimitError
) {
const headers = new Headers()
if (error.retryAfter) {
headers.set("retry-after", String(error.retryAfter))
@ -492,9 +496,8 @@ export async function handler(
}
}
async function authenticate(modelInfo: ModelInfo) {
const apiKey = opts.parseApiKey(input.request.headers)
if (!apiKey || apiKey === "public") {
async function authenticate(modelInfo: ModelInfo, zenApiKey?: string) {
if (!zenApiKey) {
if (modelInfo.allowAnonymous) return
throw new AuthError(t("zen.api.error.missingApiKey"))
}
@ -573,7 +576,7 @@ export async function handler(
isNull(LiteTable.timeDeleted),
),
)
.where(and(eq(KeyTable.key, apiKey), isNull(KeyTable.timeDeleted)))
.where(and(eq(KeyTable.key, zenApiKey), isNull(KeyTable.timeDeleted)))
.then((rows) => rows[0]),
)