zen: fix rate limit

2 files changed, 21 insertions, 15 deletions

diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index ac487a341..8981ecac2 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -69,7 +69,7 @@ export async function handler(
     const dataDumper = createDataDumper(sessionId, requestId, projectId)
     const trialLimiter = createTrialLimiter(modelInfo.trial, ip, ocClient)
     const isTrial = await trialLimiter?.isTrial()
-    const rateLimiter = createRateLimiter(modelInfo.id, modelInfo.rateLimit, ip)
+    const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip)
     await rateLimiter?.check()
     const stickyTracker = createStickyTracker(modelInfo.stickyProvider ?? false, sessionId)
     const stickyProvider = await stickyTracker?.get()
diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts
index b3c036815..244db072c 100644
--- a/packages/console/app/src/routes/zen/util/rateLimiter.ts
+++ b/packages/console/app/src/routes/zen/util/rateLimiter.ts
@@ -1,28 +1,34 @@
-import { Resource } from "@opencode-ai/console-resource"
+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 { RateLimitError } from "./error"
 import { logger } from "./logger"
 
-export function createRateLimiter(model: string, limit: number | undefined, ip: string) {
+export function createRateLimiter(limit: number | undefined, rawIp: string) {
   if (!limit) return
 
+  const ip = !rawIp.length ? "unknown" : rawIp
   const now = Date.now()
-  const currKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now)}`
-  const prevKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now - 3_600_000)}`
-  let currRate: number
-  let prevRate: number
+  const intervals = [buildYYYYMMDDHH(now), buildYYYYMMDDHH(now - 3_600_000), buildYYYYMMDDHH(now - 7_200_000)]
 
   return {
     track: async () => {
-      await Resource.GatewayKv.put(currKey, currRate + 1, { expirationTtl: 3600 })
+      await Database.use((tx) =>
+        tx
+          .insert(IpRateLimitTable)
+          .values({ ip, interval: intervals[0], count: 1 })
+          .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }),
+      )
     },
     check: async () => {
-      const values = await Resource.GatewayKv.get([currKey, prevKey])
-      const prevValue = values?.get(prevKey)
-      const currValue = values?.get(currKey)
-      prevRate = prevValue ? parseInt(prevValue) : 0
-      currRate = currValue ? parseInt(currValue) : 0
-      logger.debug(`rate limit ${model} prev/curr: ${prevRate}/${currRate}`)
-      if (prevRate + currRate >= limit) throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
+      const rows = await Database.use((tx) =>
+        tx
+          .select({ count: IpRateLimitTable.count })
+          .from(IpRateLimitTable)
+          .where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, intervals))),
+      )
+      const total = rows.reduce((sum, r) => sum + r.count, 0)
+      logger.debug(`rate limit total: ${total}`)
+      if (total >= limit) throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
     },
   }
 }