From 8df7ccc304273f3642f93543a36c955491be8490 Mon Sep 17 00:00:00 2001 From: Frank Date: Tue, 14 Apr 2026 20:29:21 -0400 Subject: zen: rate limiter --- packages/console/app/src/routes/zen/util/error.ts | 1 + .../console/app/src/routes/zen/util/handler.ts | 31 +++++---- .../app/src/routes/zen/util/ipRateLimiter.ts | 70 ++++++++++++++++++++ .../app/src/routes/zen/util/keyRateLimiter.ts | 39 +++++++++++ .../console/app/src/routes/zen/util/rateLimiter.ts | 77 ---------------------- 5 files changed, 127 insertions(+), 91 deletions(-) create mode 100644 packages/console/app/src/routes/zen/util/ipRateLimiter.ts create mode 100644 packages/console/app/src/routes/zen/util/keyRateLimiter.ts delete mode 100644 packages/console/app/src/routes/zen/util/rateLimiter.ts (limited to 'packages/console/app/src') diff --git a/packages/console/app/src/routes/zen/util/error.ts b/packages/console/app/src/routes/zen/util/error.ts index a3a93d2ef..b2a1d30d0 100644 --- a/packages/console/app/src/routes/zen/util/error.ts +++ b/packages/console/app/src/routes/zen/util/error.ts @@ -11,5 +11,6 @@ class LimitError extends Error { this.retryAfter = retryAfter } } +export class RateLimitError extends LimitError {} export class FreeUsageLimitError extends LimitError {} export class SubscriptionUsageLimitError extends LimitError {} diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 8c391d590..46d843522 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -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]), ) diff --git a/packages/console/app/src/routes/zen/util/ipRateLimiter.ts b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts new file mode 100644 index 000000000..d22ab4ae2 --- /dev/null +++ b/packages/console/app/src/routes/zen/util/ipRateLimiter.ts @@ -0,0 +1,70 @@ +import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizzle/index.js" +import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" +import { FreeUsageLimitError } from "./error" +import { logger } from "./logger" +import { i18n } from "~/i18n" +import { localeFromRequest } from "~/lib/language" +import { Subscription } from "@opencode-ai/console-core/subscription.js" + +export function createRateLimiter(modelId: string, rateLimit: number | undefined, rawIp: string, request: Request) { + const dict = i18n(localeFromRequest(request)) + + const limits = Subscription.getFreeLimits() + const dailyLimit = rateLimit ?? limits.dailyRequests + const isDefaultModel = !rateLimit + + const ip = !rawIp.length ? "unknown" : rawIp + const now = Date.now() + const lifetimeInterval = "" + const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) + + let _isNew: boolean + + return { + check: async () => { + const rows = await Database.use((tx) => + tx + .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) + .from(IpRateLimitTable) + .where( + and( + eq(IpRateLimitTable.ip, ip), + isDefaultModel + ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) + : inArray(IpRateLimitTable.interval, [dailyInterval]), + ), + ), + ) + const lifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 + const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 + logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) + + _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + + if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit)) + throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now)) + }, + track: async () => { + await Database.use((tx) => + tx + .insert(IpRateLimitTable) + .values([ + { ip, interval: dailyInterval, count: 1 }, + ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), + ]) + .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), + ) + }, + } +} + +export function getRetryAfterDay(now: number) { + return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000) +} + +function buildYYYYMMDD(timestamp: number) { + return new Date(timestamp) + .toISOString() + .replace(/[^0-9]/g, "") + .substring(0, 8) +} diff --git a/packages/console/app/src/routes/zen/util/keyRateLimiter.ts b/packages/console/app/src/routes/zen/util/keyRateLimiter.ts new file mode 100644 index 000000000..e3e0fb18f --- /dev/null +++ b/packages/console/app/src/routes/zen/util/keyRateLimiter.ts @@ -0,0 +1,39 @@ +import { Database, eq, and, sql } from "@opencode-ai/console-core/drizzle/index.js" +import { KeyRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" +import { RateLimitError } from "./error" +import { i18n } from "~/i18n" +import { localeFromRequest } from "~/lib/language" + +export function createRateLimiter(modelId: string, zenApiKey: string | undefined, request: Request) { + if (!zenApiKey) return + const dict = i18n(localeFromRequest(request)) + + const LIMIT = 100 + const yyyyMMddHHmm = new Date(Date.now()) + .toISOString() + .replace(/[^0-9]/g, "") + .substring(0, 12) + const interval = `${modelId.substring(0, 27)}-${yyyyMMddHHmm}` + + return { + check: async () => { + const rows = await Database.use((tx) => + tx + .select({ interval: KeyRateLimitTable.interval, count: KeyRateLimitTable.count }) + .from(KeyRateLimitTable) + .where(and(eq(KeyRateLimitTable.key, zenApiKey), eq(KeyRateLimitTable.interval, interval))), + ).then((rows) => rows[0]) + const count = rows?.count ?? 0 + + if (count >= LIMIT) throw new RateLimitError(dict["zen.api.error.rateLimitExceeded"], 60) + }, + track: async () => { + await Database.use((tx) => + tx + .insert(KeyRateLimitTable) + .values({ key: zenApiKey, interval, count: 1 }) + .onDuplicateKeyUpdate({ set: { count: sql`${KeyRateLimitTable.count} + 1` } }), + ) + }, + } +} diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts deleted file mode 100644 index 160633981..000000000 --- a/packages/console/app/src/routes/zen/util/rateLimiter.ts +++ /dev/null @@ -1,77 +0,0 @@ -import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizzle/index.js" -import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" -import { FreeUsageLimitError } from "./error" -import { logger } from "./logger" -import { i18n } from "~/i18n" -import { localeFromRequest } from "~/lib/language" -import { Subscription } from "@opencode-ai/console-core/subscription.js" - -export function createRateLimiter( - modelId: string, - allowAnonymous: boolean | undefined, - rateLimit: number | undefined, - rawIp: string, - request: Request, -) { - if (!allowAnonymous) return - const dict = i18n(localeFromRequest(request)) - - const limits = Subscription.getFreeLimits() - const dailyLimit = rateLimit ?? limits.dailyRequests - const isDefaultModel = !rateLimit - - const ip = !rawIp.length ? "unknown" : rawIp - const now = Date.now() - const lifetimeInterval = "" - const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) - - let _isNew: boolean - - return { - check: async () => { - const rows = await Database.use((tx) => - tx - .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) - .from(IpRateLimitTable) - .where( - and( - eq(IpRateLimitTable.ip, ip), - isDefaultModel - ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) - : inArray(IpRateLimitTable.interval, [dailyInterval]), - ), - ), - ) - const lifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 - const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 - logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) - - _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 - - if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit)) - throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now)) - }, - track: async () => { - await Database.use((tx) => - tx - .insert(IpRateLimitTable) - .values([ - { ip, interval: dailyInterval, count: 1 }, - ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), - ]) - .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), - ) - }, - } -} - -export function getRetryAfterDay(now: number) { - return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000) -} - -function buildYYYYMMDD(timestamp: number) { - return new Date(timestamp) - .toISOString() - .replace(/[^0-9]/g, "") - .substring(0, 8) -} -- cgit v1.2.3