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
@@ -11,5 +11,6 @@ class LimitError extends Error {
this.retryAfter = retryAfter this.retryAfter = retryAfter
} }
} }
export class RateLimitError extends LimitError {}
export class FreeUsageLimitError extends LimitError {} export class FreeUsageLimitError extends LimitError {}
export class SubscriptionUsageLimitError extends LimitError {} export class SubscriptionUsageLimitError extends LimitError {}
@@ -21,6 +21,7 @@ import {
MonthlyLimitError, MonthlyLimitError,
UserLimitError, UserLimitError,
ModelError, ModelError,
RateLimitError,
FreeUsageLimitError, FreeUsageLimitError,
SubscriptionUsageLimitError, SubscriptionUsageLimitError,
} from "./error" } from "./error"
@@ -35,7 +36,8 @@ import { anthropicHelper } from "./provider/anthropic"
import { googleHelper } from "./provider/google" import { googleHelper } from "./provider/google"
import { openaiHelper } from "./provider/openai" import { openaiHelper } from "./provider/openai"
import { oaCompatHelper } from "./provider/openai-compatible" 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 { createDataDumper } from "./dataDumper"
import { createTrialLimiter } from "./trialLimiter" import { createTrialLimiter } from "./trialLimiter"
import { createStickyTracker } from "./stickyProviderTracker" import { createStickyTracker } from "./stickyProviderTracker"
@@ -92,6 +94,8 @@ export async function handler(
const isStream = opts.parseIsStream(url, body) const isStream = opts.parseIsStream(url, body)
const rawIp = input.request.headers.get("x-real-ip") ?? "" const rawIp = input.request.headers.get("x-real-ip") ?? ""
const ip = rawIp.includes(":") ? rawIp.split(":").slice(0, 4).join(":") : rawIp 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 sessionId = input.request.headers.get("x-opencode-session") ?? ""
const requestId = input.request.headers.get("x-opencode-request") ?? "" const requestId = input.request.headers.get("x-opencode-request") ?? ""
const projectId = input.request.headers.get("x-opencode-project") ?? "" const projectId = input.request.headers.get("x-opencode-project") ?? ""
@@ -108,17 +112,13 @@ export async function handler(
const dataDumper = createDataDumper(sessionId, requestId, projectId) const dataDumper = createDataDumper(sessionId, requestId, projectId)
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip) const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
const trialProviders = await trialLimiter?.check() const trialProviders = await trialLimiter?.check()
const rateLimiter = createRateLimiter( const rateLimiter = modelInfo.allowAnonymous
modelInfo.id, ? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
modelInfo.allowAnonymous, : createKeyRateLimiter(modelInfo.id, zenApiKey, input.request)
modelInfo.rateLimit,
ip,
input.request,
)
await rateLimiter?.check() await rateLimiter?.check()
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId) const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
const stickyProvider = await stickyTracker?.get() const stickyProvider = await stickyTracker?.get()
const authInfo = await authenticate(modelInfo) const authInfo = await authenticate(modelInfo, zenApiKey)
const billingSource = validateBilling(authInfo, modelInfo) const billingSource = validateBilling(authInfo, modelInfo)
logger.metric({ source: billingSource }) logger.metric({ source: billingSource })
@@ -363,7 +363,11 @@ export async function handler(
{ status: 401 }, { status: 401 },
) )
if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) { if (
error instanceof RateLimitError ||
error instanceof FreeUsageLimitError ||
error instanceof SubscriptionUsageLimitError
) {
const headers = new Headers() const headers = new Headers()
if (error.retryAfter) { if (error.retryAfter) {
headers.set("retry-after", String(error.retryAfter)) headers.set("retry-after", String(error.retryAfter))
@@ -492,9 +496,8 @@ export async function handler(
} }
} }
async function authenticate(modelInfo: ModelInfo) { async function authenticate(modelInfo: ModelInfo, zenApiKey?: string) {
const apiKey = opts.parseApiKey(input.request.headers) if (!zenApiKey) {
if (!apiKey || apiKey === "public") {
if (modelInfo.allowAnonymous) return if (modelInfo.allowAnonymous) return
throw new AuthError(t("zen.api.error.missingApiKey")) throw new AuthError(t("zen.api.error.missingApiKey"))
} }
@@ -573,7 +576,7 @@ export async function handler(
isNull(LiteTable.timeDeleted), 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]), .then((rows) => rows[0]),
) )
@@ -6,14 +6,7 @@ import { i18n } from "~/i18n"
import { localeFromRequest } from "~/lib/language" import { localeFromRequest } from "~/lib/language"
import { Subscription } from "@opencode-ai/console-core/subscription.js" import { Subscription } from "@opencode-ai/console-core/subscription.js"
export function createRateLimiter( export function createRateLimiter(modelId: string, rateLimit: number | undefined, rawIp: string, request: Request) {
modelId: string,
allowAnonymous: boolean | undefined,
rateLimit: number | undefined,
rawIp: string,
request: Request,
) {
if (!allowAnonymous) return
const dict = i18n(localeFromRequest(request)) const dict = i18n(localeFromRequest(request))
const limits = Subscription.getFreeLimits() const limits = Subscription.getFreeLimits()
@@ -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` } }),
)
},
}
}
@@ -1,5 +1,5 @@
import { describe, expect, test } from "bun:test" import { describe, expect, test } from "bun:test"
import { getRetryAfterDay } from "../src/routes/zen/util/rateLimiter" import { getRetryAfterDay } from "../src/routes/zen/util/ipRateLimiter"
describe("getRetryAfterDay", () => { describe("getRetryAfterDay", () => {
test("returns full day at midnight UTC", () => { test("returns full day at midnight UTC", () => {
@@ -0,0 +1,6 @@
CREATE TABLE `key_rate_limit` (
`key` varchar(255) NOT NULL,
`interval` varchar(12) NOT NULL,
`count` int NOT NULL,
CONSTRAINT PRIMARY KEY(`key`,`interval`)
);
File diff suppressed because it is too large Load Diff
@@ -0,0 +1 @@
ALTER TABLE `key_rate_limit` MODIFY COLUMN `interval` varchar(20) NOT NULL;
@@ -0,0 +1 @@
ALTER TABLE `key_rate_limit` MODIFY COLUMN `interval` varchar(40) NOT NULL;
File diff suppressed because it is too large Load Diff
@@ -20,3 +20,13 @@ export const IpRateLimitTable = mysqlTable(
}, },
(table) => [primaryKey({ columns: [table.ip, table.interval] })], (table) => [primaryKey({ columns: [table.ip, table.interval] })],
) )
export const KeyRateLimitTable = mysqlTable(
"key_rate_limit",
{
key: varchar("key", { length: 255 }).notNull(),
interval: varchar("interval", { length: 40 }).notNull(),
count: int("count").notNull(),
},
(table) => [primaryKey({ columns: [table.key, table.interval] })],
)