zen: rate limiter
This commit is contained in:
@@ -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]),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
+1
-8
@@ -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;
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -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] })],
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user