vrr/opencode
2
0
Fork 0
mirror of https://github.com/anomalyco/opencode.git synced 2026-05-31 05:15:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 18

perf: use redis for api key rate limit (#29242)

Browse source
This commit is contained in:
Victor Navarro 2026-05-27 11:27:06 +02:00 • committed by GitHub
parent e2dc89c6f3
commit 5acc368ef4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 12 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

43
packages/console/app/src/routes/zen/util/ipRateLimiter.ts
View file

@ -1,5 +1,3 @@
import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizzle/index.js"
import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
import { FreeUsageLimitError } from "./error"
import { logger } from "./logger"
import { buildRateLimitKey, getRedis } from "./redis"
@ -22,7 +20,6 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
const ip = !rawIp.length ? "unknown" : rawIp
const now = Date.now()
const lifetimeInterval = ""
const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now)
const retryAfter = getRetryAfterDay(now)
const redis = getRedis()
@ -32,33 +29,12 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
return {
check: async () => {
const [counts, rows] = await Promise.all([
redis.mget<(string | number | null)[]>(isDefaultModel ? [lifetimeKey, dailyKey] : [dailyKey]).catch(() => []),
Database.use((tx) =>
tx
.select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count })
.from(IpRateLimitTable)
.where(
and(
eq(IpRateLimitTable.ip, ip),
isDefaultModel
? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval])
: inArray(IpRateLimitTable.interval, [dailyInterval]),
),
),
),
])
const redisLifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0
const redisDailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0)
const databaseLifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0
const databaseDailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0
const lifetimeCount = Math.max(redisLifetimeCount, databaseLifetimeCount)
const dailyCount = Math.max(redisDailyCount, databaseDailyCount)
const counts = await redis.mget<(string | number | null)[]>(isDefaultModel ? [lifetimeKey, dailyKey] : [dailyKey])
const lifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0
const dailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0)
logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`)
isNew = isDefaultModel && lifetimeCount < dailyLimit * 7
if (isDefaultModel && databaseLifetimeCount > redisLifetimeCount)
await redis.set(lifetimeKey, databaseLifetimeCount).catch(() => {})
if ((isNew && dailyCount >= dailyLimit * 2) || (!isNew && dailyCount >= dailyLimit))
throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter)
@ -68,18 +44,7 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
pipeline.incr(dailyKey)
pipeline.expire(dailyKey, retryAfter)
if (isNew) pipeline.incr(lifetimeKey)
await Promise.all([
pipeline.exec().catch(() => {}),
Database.use((tx) =>
tx
.insert(IpRateLimitTable)
.values([
{ ip, interval: dailyInterval, count: 1 },
...(isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []),
])
.onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }),
),
])
await pipeline.exec()
},
}
}

23
packages/console/app/src/routes/zen/util/keyRateLimiter.ts
View file

@ -1,6 +1,5 @@
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 { buildRateLimitKey, getRedis } from "./redis"
import { i18n } from "~/i18n"
import { localeFromRequest } from "~/lib/language"
@ -19,26 +18,20 @@ export function createRateLimiter(
.replace(/[^0-9]/g, "")
.substring(0, 12)
const interval = `${modelId.substring(0, 27)}-${yyyyMMddHHmm}`
const redis = getRedis()
const key = buildRateLimitKey("key", zenApiKey, interval)
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
const count = Number((await redis.mget<(string | number | null)[]>([key]))[0] ?? 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` } }),
)
const pipeline = redis.pipeline()
pipeline.incr(key)
pipeline.expire(key, 60)
await pipeline.exec()
},
}
}