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() diff --git a/packages/console/app/test/rateLimiter.test.ts b/packages/console/app/test/rateLimiter.test.ts new file mode 100644 index 000000000..864f907d6 --- /dev/null +++ b/packages/console/app/test/rateLimiter.test.ts @@ -0,0 +1,92 @@ +import { describe, expect, test } from "bun:test" +import { getRetryAfterDay, getRetryAfterHour } from "../src/routes/zen/util/rateLimiter" + +describe("getRetryAfterDay", () => { + test("returns full day at midnight UTC", () => { + const midnight = Date.UTC(2026, 0, 15, 0, 0, 0, 0) + expect(getRetryAfterDay(midnight)).toBe(86_400) + }) + + test("returns remaining seconds until next UTC day", () => { + const noon = Date.UTC(2026, 0, 15, 12, 0, 0, 0) + expect(getRetryAfterDay(noon)).toBe(43_200) + }) + + test("rounds up to nearest second", () => { + const almost = Date.UTC(2026, 0, 15, 23, 59, 59, 500) + expect(getRetryAfterDay(almost)).toBe(1) + }) +}) + +describe("getRetryAfterHour", () => { + // 14:30:00 UTC — 30 minutes into the current hour + const now = Date.UTC(2026, 0, 15, 14, 30, 0, 0) + const intervals = ["2026011514", "2026011513", "2026011512"] + + test("waits 3 hours when all usage is in current hour", () => { + const rows = [{ interval: "2026011514", count: 10 }] + // only current hour has usage — it won't leave the window for 3 hours from hour start + // 3 * 3600 - 1800 = 9000s + expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(9000) + }) + + test("waits 1 hour when dropping oldest interval is sufficient", () => { + const rows = [ + { interval: "2026011514", count: 2 }, + { interval: "2026011512", count: 10 }, + ] + // total=12, drop oldest (-2h, count=10) -> 2 < 10 + // hours = 3 - 2 = 1 -> 1 * 3600 - 1800 = 1800s + expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(1800) + }) + + test("waits 2 hours when usage spans oldest two intervals", () => { + const rows = [ + { interval: "2026011513", count: 8 }, + { interval: "2026011512", count: 5 }, + ] + // total=13, drop -2h (5) -> 8, 8 >= 8, drop -1h (8) -> 0 < 8 + // hours = 3 - 1 = 2 -> 2 * 3600 - 1800 = 5400s + expect(getRetryAfterHour(rows, intervals, 8, now)).toBe(5400) + }) + + test("waits 1 hour when oldest interval alone pushes over limit", () => { + const rows = [ + { interval: "2026011514", count: 1 }, + { interval: "2026011513", count: 1 }, + { interval: "2026011512", count: 10 }, + ] + // total=12, drop -2h (10) -> 2 < 10 + // hours = 3 - 2 = 1 -> 1800s + expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(1800) + }) + + test("waits 2 hours when middle interval keeps total over limit", () => { + const rows = [ + { interval: "2026011514", count: 4 }, + { interval: "2026011513", count: 4 }, + { interval: "2026011512", count: 4 }, + ] + // total=12, drop -2h (4) -> 8, 8 >= 5, drop -1h (4) -> 4 < 5 + // hours = 3 - 1 = 2 -> 5400s + expect(getRetryAfterHour(rows, intervals, 5, now)).toBe(5400) + }) + + test("rounds up to nearest second", () => { + const offset = Date.UTC(2026, 0, 15, 14, 30, 0, 500) + const rows = [ + { interval: "2026011514", count: 2 }, + { interval: "2026011512", count: 10 }, + ] + // hours=1 -> 3_600_000 - 1_800_500 = 1_799_500ms -> ceil(1799.5) = 1800 + expect(getRetryAfterHour(rows, intervals, 10, offset)).toBe(1800) + }) + + test("fallback returns time until next hour when rows are empty", () => { + // edge case: rows empty but function called (shouldn't happen in practice) + // loop drops all zeros, running stays 0 which is < any positive limit on first iteration + const rows: { interval: string; count: number }[] = [] + // drop -2h (0) -> 0 < 1 -> hours = 3 - 2 = 1 -> 1800s + expect(getRetryAfterHour(rows, intervals, 1, now)).toBe(1800) + }) +})