perf: use redis/upstash for ip rate limits (#28694)

packages/console/app/package.json

@@ -26,6 +26,7 @@
     "@solidjs/router": "catalog:",
     "@solidjs/start": "catalog:",
     "@stripe/stripe-js": "8.6.1",
+    "@upstash/redis": "1.38.0",
     "chart.js": "4.5.1",
     "nitro": "3.0.1-alpha.1",
     "solid-js": "catalog:",

packages/console/app/src/routes/zen/util/ipRateLimiter.ts

@@ -2,6 +2,7 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz
 import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
 import { FreeUsageLimitError } from "./error"
 import { logger } from "./logger"
+import { buildRateLimitKey, getRedis } from "./redis"
 import { i18n } from "~/i18n"
 import { localeFromRequest } from "~/lib/language"
 import { Subscription } from "@opencode-ai/console-core/subscription.js"
@@ -23,43 +24,62 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
   const now = Date.now()
   const lifetimeInterval = ""
   const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now)
-
-  let _isNew: boolean
+  const retryAfter = getRetryAfterDay(now)
+  const redis = getRedis()
+  const lifetimeKey = buildRateLimitKey("ip", ip)
+  const dailyKey = buildRateLimitKey("ip", ip, dailyInterval)
+  let isNew = false
 
   return {
     check: async () => {
-      const rows = await 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 [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 lifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0
-      const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0
+        ),
+      ])
+      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)
       logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`)
 
-      _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7
+      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"], getRetryAfterDay(now))
+      if ((isNew && dailyCount >= dailyLimit * 2) || (!isNew && dailyCount >= dailyLimit))
+        throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter)
     },
     track: async () => {
-      await 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` } }),
-      )
+      const pipeline = redis.pipeline()
+      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` } }),
+        ),
+      ])
     },
   }
 }

packages/console/app/src/routes/zen/util/redis.ts

@@ -0,0 +1,18 @@
+import { Resource } from "@opencode-ai/console-resource"
+import { Redis } from "@upstash/redis/cloudflare"
+
+let redis: Redis | undefined
+
+export function getRedis() {
+  if (redis) return redis
+  redis = new Redis({
+    url: Resource.UpstashRedisRestUrl.value,
+    token: Resource.UpstashRedisRestToken.value,
+    enableTelemetry: false,
+  })
+  return redis
+}
+
+export function buildRateLimitKey(kind: string, identifier: string, interval?: string) {
+  return `${Resource.App.stage}:ratelimit:${kind}:${identifier}${interval ? `:${interval}` : ""}`
+}