wip: zen

3 files changed, 36 insertions, 5 deletions

diff --git a/packages/console/app/src/routes/zen/util/error.ts b/packages/console/app/src/routes/zen/util/error.ts
index a1393eb7f..a3a93d2ef 100644
--- a/packages/console/app/src/routes/zen/util/error.ts
+++ b/packages/console/app/src/routes/zen/util/error.ts
@@ -3,11 +3,13 @@ export class CreditsError extends Error {}
 export class MonthlyLimitError extends Error {}
 export class UserLimitError extends Error {}
 export class ModelError extends Error {}
-export class FreeUsageLimitError extends Error {}
-export class SubscriptionUsageLimitError extends Error {
+
+class LimitError extends Error {
   retryAfter?: number
   constructor(message: string, retryAfter?: number) {
     super(message)
     this.retryAfter = retryAfter
   }
 }
+export class FreeUsageLimitError extends LimitError {}
+export class SubscriptionUsageLimitError extends LimitError {}
diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index a72435e68..9646cacd0 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -313,7 +313,7 @@ export async function handler(
 
     if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) {
       const headers = new Headers()
-      if (error instanceof SubscriptionUsageLimitError && error.retryAfter) {
+      if (error.retryAfter) {
         headers.set("retry-after", String(error.retryAfter))
       }
       return new Response(
diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts
index fafbc06e9..5e4f31e67 100644
--- a/packages/console/app/src/routes/zen/util/rateLimiter.ts
+++ b/packages/console/app/src/routes/zen/util/rateLimiter.ts
@@ -28,17 +28,46 @@ export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: s
     check: async () => {
       const rows = await Database.use((tx) =>
         tx
-          .select({ count: IpRateLimitTable.count })
+          .select({ interval: IpRateLimitTable.interval, 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 >= limitValue) throw new FreeUsageLimitError(`Rate limit exceeded. Please try again later.`)
+      if (total >= limitValue)
+        throw new FreeUsageLimitError(
+          `Rate limit exceeded. Please try again later.`,
+          limit.period === "day" ? getRetryAfterDay(now) : getRetryAfterHour(rows, intervals, limitValue, now),
+        )
     },
   }
 }
 
+export function getRetryAfterDay(now: number) {
+  return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000)
+}
+
+export function getRetryAfterHour(
+  rows: { interval: string; count: number }[],
+  intervals: string[],
+  limit: number,
+  now: number,
+) {
+  const counts = new Map(rows.map((r) => [r.interval, r.count]))
+  // intervals are ordered newest to oldest: [current, -1h, -2h]
+  // simulate dropping oldest intervals one at a time
+  let running = intervals.reduce((sum, i) => sum + (counts.get(i) ?? 0), 0)
+  for (let i = intervals.length - 1; i >= 0; i--) {
+    running -= counts.get(intervals[i]) ?? 0
+    if (running < limit) {
+      // interval at index i rolls out of the window (intervals.length - i) hours from the current hour start
+      const hours = intervals.length - i
+      return Math.ceil((hours * 3_600_000 - (now % 3_600_000)) / 1000)
+    }
+  }
+  return Math.ceil((3_600_000 - (now % 3_600_000)) / 1000)
+}
+
 function buildYYYYMMDD(timestamp: number) {
   return new Date(timestamp)
     .toISOString()