mirror of
https://fastgit.cc/https://github.com/anomalyco/opencode
synced 2026-04-21 21:31:53 +08:00
zen: rate limiter
This commit is contained in:
@@ -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 {}
|
||||
|
||||
@@ -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]),
|
||||
)
|
||||
|
||||
|
||||
@@ -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()
|
||||
39
packages/console/app/src/routes/zen/util/keyRateLimiter.ts
Normal file
39
packages/console/app/src/routes/zen/util/keyRateLimiter.ts
Normal 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` } }),
|
||||
)
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -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] })],
|
||||
)
|
||||
|
||||
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