zen: rate limit (#11735)
This commit is contained in:
@@ -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()
|
||||||
|
|||||||
@@ -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({
|
||||||
|
|||||||
Reference in New Issue
Block a user