add redis-backed ip rate limiting with database fallback

bun.lock

@@ -101,6 +101,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:",
@@ -2467,6 +2468,8 @@
 
     "@ungap/structured-clone": ["@ungap/structured-clone@1.3.0", "", {}, "sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g=="],
 
+    "@upstash/redis": ["@upstash/redis@1.38.0", "", { "dependencies": { "uncrypto": "^0.1.3" } }, "sha512-wu+dZBptlLy0+MCUEoHmzrY/TnmgDey3+c7EbIGwrLqAvkP8yi5MWZHYGIFtAygmL4Bkz2TdFu+eU0vFPncIcg=="],
+
     "@valibot/to-json-schema": ["@valibot/to-json-schema@1.6.0", "", { "peerDependencies": { "valibot": "^1.3.0" } }, "sha512-d6rYyK5KVa2XdqamWgZ4/Nr+cXhxjy7lmpe6Iajw15J/jmU+gyxl2IEd1Otg1d7Rl3gOQL5reulnSypzBtYy1A=="],
 
     "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="],

infra/console.ts

@@ -250,6 +250,8 @@ new sst.cloudflare.x.SolidStart("Console", {
     bucket,
     bucketNew,
     database,
+    SECRET.UpstashRedisRestUrl,
+    SECRET.UpstashRedisRestToken,
     AUTH_API_URL,
     STRIPE_WEBHOOK_SECRET,
     DISCORD_INCIDENT_WEBHOOK_URL,

infra/secret.ts

@@ -8,4 +8,6 @@ export const SECRET = {
   R2AccessKey: new sst.Secret("R2AccessKey", "unknown"),
   R2SecretKey: new sst.Secret("R2SecretKey", "unknown"),
   HoneycombWebhookSecret: new random.RandomPassword("HoneycombWebhookSecret", { length: 24 }),
+  UpstashRedisRestUrl: new sst.Secret("UpstashRedisRestUrl"),
+  UpstashRedisRestToken: new sst.Secret("UpstashRedisRestToken"),
 }

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,7 +2,9 @@ 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 { i18n } from "~/i18n"
+import { getRetryAfterDay, type RateLimiter, type RateLimiterState } from "./rateLimit"
+import { buildRateLimitKey, redis } from "./redis"
+import { i18n, type Dict } from "~/i18n"
 import { localeFromRequest } from "~/lib/language"
 import { Subscription } from "@opencode-ai/console-core/subscription.js"
 
@@ -18,10 +20,23 @@ 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 state = { isNew: false, fallbackDatabase: false }
+  const databaseLimiter = createDatabaseRateLimiter(ip, dailyInterval, dailyLimit, isDefaultModel, dict, retryAfter, state)
+  return createUpstashRateLimiter(ip, dailyInterval, dailyLimit, isDefaultModel, dict, retryAfter, state, databaseLimiter)
+}
 
-  let _isNew: boolean
+function createDatabaseRateLimiter(
+  ip: string,
+  dailyInterval: string,
+  dailyLimit: number,
+  isDefaultModel: boolean,
+  dict: Dict,
+  retryAfter: number,
+  state: RateLimiterState,
+): RateLimiter {
+  const lifetimeInterval = ""
 
   return {
     check: async () => {
@@ -42,10 +57,10 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
       const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0
       logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`)
 
-      _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7
+      state.isNew = isDefaultModel && lifetimeCount < dailyLimit * 7
 
-      if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit))
-        throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now))
+      if ((state.isNew && dailyCount >= dailyLimit * 2) || (!state.isNew && dailyCount >= dailyLimit))
+        throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter)
     },
     track: async () => {
       await Database.use((tx) =>
@@ -53,7 +68,7 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
           .insert(IpRateLimitTable)
           .values([
             { ip, interval: dailyInterval, count: 1 },
-            ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []),
+            ...(state.isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []),
           ])
           .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }),
       )
@@ -61,8 +76,55 @@ export function createRateLimiter(modelId: string, rateLimit: number | undefined
   }
 }
 
-export function getRetryAfterDay(now: number) {
-  return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000)
+function createUpstashRateLimiter(
+  ip: string,
+  dailyInterval: string,
+  dailyLimit: number,
+  isDefaultModel: boolean,
+  dict: Dict,
+  retryAfter: number,
+  state: RateLimiterState,
+  databaseLimiter: RateLimiter,
+): RateLimiter {
+  const lifetimeInterval = ""
+  const lifetimeKey = buildRateLimitKey("ip", ip, lifetimeInterval)
+  const dailyKey = buildRateLimitKey("ip", ip, dailyInterval)
+
+  return {
+    check: async () => {
+      try {
+        const keys = isDefaultModel
+          ? [lifetimeKey, dailyKey]
+          : [dailyKey]
+        const counts = await redis.mget<(string | number | null)[]>(keys)
+        const lifetimeCount = isDefaultModel ? Number(counts[0] ?? 0) : 0
+        const dailyCount = Number(counts[isDefaultModel ? 1 : 0] ?? 0)
+        logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`)
+
+        state.isNew = isDefaultModel && lifetimeCount < dailyLimit * 7
+        if ((state.isNew && dailyCount >= dailyLimit * 2) || (!state.isNew && dailyCount >= dailyLimit))
+          throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], retryAfter)
+      } catch (error) {
+        if (error instanceof FreeUsageLimitError) throw error
+
+        state.fallbackDatabase = true
+        await databaseLimiter.check()
+      }
+    },
+    track: async () => {
+      if (state.fallbackDatabase) return databaseLimiter.track()
+
+      try {
+        const pipeline = redis.pipeline()
+        pipeline.incr(dailyKey)
+        pipeline.expire(dailyKey, retryAfter)
+        if (state.isNew) pipeline.incr(lifetimeKey)
+        await pipeline.exec()
+      } catch {
+        await databaseLimiter.track()
+      }
+    },
+  }
 }
 
 function buildYYYYMMDD(timestamp: number) {

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

@@ -0,0 +1,13 @@
+export type RateLimiter = {
+  check: () => Promise<void>
+  track: () => Promise<void>
+}
+
+export type RateLimiterState = {
+  isNew: boolean
+  fallbackDatabase: boolean
+}
+
+export function getRetryAfterDay(now: number) {
+  return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000)
+}

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

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

packages/console/app/test/rateLimiter.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test"
-import { getRetryAfterDay } from "../src/routes/zen/util/ipRateLimiter"
+import { getRetryAfterDay } from "../src/routes/zen/util/rateLimit"
 
 describe("getRetryAfterDay", () => {
   test("returns full day at midnight UTC", () => {

sst-env.d.ts

@@ -137,6 +137,14 @@ declare module "sst" {
       "type": "sst.cloudflare.SolidStart"
       "url": string
     }
+    "UpstashRedisRestToken": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
+    "UpstashRedisRestUrl": {
+      "type": "sst.sst.Secret"
+      "value": string
+    }
     "Web": {
       "type": "sst.cloudflare.Astro"
       "url": string