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

@@ -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 {}

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]),
)

View File

@@ -6,14 +6,7 @@ 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
export function createRateLimiter(modelId: string, rateLimit: number | undefined, rawIp: string, request: Request) {
const dict = i18n(localeFromRequest(request))
const limits = Subscription.getFreeLimits()

View File

@@ -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` } }),
)
},
}
}

View File

@@ -1,5 +1,5 @@
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", () => {
test("returns full day at midnight UTC", () => {

View File

@@ -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

View File

@@ -0,0 +1 @@
ALTER TABLE `key_rate_limit` MODIFY COLUMN `interval` varchar(20) NOT NULL;

View File

@@ -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

View File

@@ -20,3 +20,13 @@ export const IpRateLimitTable = mysqlTable(
},
(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] })],
)