zen: rate limit (#11735)

This commit is contained in:
Frank
2026-02-02 00:29:52 -05:00
committed by GitHub
parent 43354eeabd
commit 4850ecc419
2 changed files with 20 additions and 4 deletions

View File

@@ -2,13 +2,17 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz
import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js" import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
import { RateLimitError } from "./error" import { RateLimitError } from "./error"
import { logger } from "./logger" import { logger } from "./logger"
import { ZenData } from "@opencode-ai/console-core/model.js"
export function createRateLimiter(limit: number | undefined, rawIp: string) { export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: string) {
if (!limit) return if (!limit) return
const ip = !rawIp.length ? "unknown" : rawIp const ip = !rawIp.length ? "unknown" : rawIp
const now = Date.now() const now = Date.now()
const intervals = [buildYYYYMMDDHH(now), buildYYYYMMDDHH(now - 3_600_000), buildYYYYMMDDHH(now - 7_200_000)] const intervals =
limit.period === "day"
? [buildYYYYMMDD(now)]
: [buildYYYYMMDDHH(now), buildYYYYMMDDHH(now - 3_600_000), buildYYYYMMDDHH(now - 7_200_000)]
return { return {
track: async () => { track: async () => {
@@ -28,11 +32,18 @@ export function createRateLimiter(limit: number | undefined, rawIp: string) {
) )
const total = rows.reduce((sum, r) => sum + r.count, 0) const total = rows.reduce((sum, r) => sum + r.count, 0)
logger.debug(`rate limit total: ${total}`) logger.debug(`rate limit total: ${total}`)
if (total >= limit) throw new RateLimitError(`Rate limit exceeded. Please try again later.`) if (total >= limit.value) throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
}, },
} }
} }
function buildYYYYMMDD(timestamp: number) {
return new Date(timestamp)
.toISOString()
.replace(/[^0-9]/g, "")
.substring(0, 8)
}
function buildYYYYMMDDHH(timestamp: number) { function buildYYYYMMDDHH(timestamp: number) {
return new Date(timestamp) return new Date(timestamp)
.toISOString() .toISOString()

View File

@@ -18,8 +18,13 @@ export namespace ZenData {
}), }),
), ),
}) })
const RateLimitSchema = z.object({
period: z.enum(["day", "rolling"]),
value: z.number().int(),
})
export type Format = z.infer<typeof FormatSchema> export type Format = z.infer<typeof FormatSchema>
export type Trial = z.infer<typeof TrialSchema> export type Trial = z.infer<typeof TrialSchema>
export type RateLimit = z.infer<typeof RateLimitSchema>
const ModelCostSchema = z.object({ const ModelCostSchema = z.object({
input: z.number(), input: z.number(),
@@ -37,7 +42,7 @@ export namespace ZenData {
byokProvider: z.enum(["openai", "anthropic", "google"]).optional(), byokProvider: z.enum(["openai", "anthropic", "google"]).optional(),
stickyProvider: z.enum(["strict", "prefer"]).optional(), stickyProvider: z.enum(["strict", "prefer"]).optional(),
trial: TrialSchema.optional(), trial: TrialSchema.optional(),
rateLimit: z.number().optional(), rateLimit: RateLimitSchema.optional(),
fallbackProvider: z.string().optional(), fallbackProvider: z.string().optional(),
providers: z.array( providers: z.array(
z.object({ z.object({