zen: rate limit (#11735)

2 files changed, 20 insertions, 4 deletions

diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts
index 244db072c..d54bd0306 100644
--- a/packages/console/app/src/routes/zen/util/rateLimiter.ts
+++ b/packages/console/app/src/routes/zen/util/rateLimiter.ts
@@ -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 { RateLimitError } from "./error"
 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
 
   const ip = !rawIp.length ? "unknown" : rawIp
   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 {
     track: async () => {
@@ -28,11 +32,18 @@ export function createRateLimiter(limit: number | undefined, rawIp: string) {
       )
       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.`)
+      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) {
   return new Date(timestamp)
     .toISOString()
diff --git a/packages/console/core/src/model.ts b/packages/console/core/src/model.ts
index 880c63a19..fc9674ced 100644
--- a/packages/console/core/src/model.ts
+++ b/packages/console/core/src/model.ts
@@ -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 Trial = z.infer<typeof TrialSchema>
+  export type RateLimit = z.infer<typeof RateLimitSchema>
 
   const ModelCostSchema = z.object({
     input: z.number(),
@@ -37,7 +42,7 @@ export namespace ZenData {
     byokProvider: z.enum(["openai", "anthropic", "google"]).optional(),
     stickyProvider: z.enum(["strict", "prefer"]).optional(),
     trial: TrialSchema.optional(),
-    rateLimit: z.number().optional(),
+    rateLimit: RateLimitSchema.optional(),
     fallbackProvider: z.string().optional(),
     providers: z.array(
       z.object({