wip: zen
This commit is contained in:
@@ -3,11 +3,13 @@ export class CreditsError extends Error {}
|
|||||||
export class MonthlyLimitError extends Error {}
|
export class MonthlyLimitError extends Error {}
|
||||||
export class UserLimitError extends Error {}
|
export class UserLimitError extends Error {}
|
||||||
export class ModelError extends Error {}
|
export class ModelError extends Error {}
|
||||||
export class FreeUsageLimitError extends Error {}
|
|
||||||
export class SubscriptionUsageLimitError extends Error {
|
class LimitError extends Error {
|
||||||
retryAfter?: number
|
retryAfter?: number
|
||||||
constructor(message: string, retryAfter?: number) {
|
constructor(message: string, retryAfter?: number) {
|
||||||
super(message)
|
super(message)
|
||||||
this.retryAfter = retryAfter
|
this.retryAfter = retryAfter
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
export class FreeUsageLimitError extends LimitError {}
|
||||||
|
export class SubscriptionUsageLimitError extends LimitError {}
|
||||||
|
|||||||
@@ -313,7 +313,7 @@ export async function handler(
|
|||||||
|
|
||||||
if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) {
|
if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) {
|
||||||
const headers = new Headers()
|
const headers = new Headers()
|
||||||
if (error instanceof SubscriptionUsageLimitError && error.retryAfter) {
|
if (error.retryAfter) {
|
||||||
headers.set("retry-after", String(error.retryAfter))
|
headers.set("retry-after", String(error.retryAfter))
|
||||||
}
|
}
|
||||||
return new Response(
|
return new Response(
|
||||||
|
|||||||
@@ -28,17 +28,46 @@ export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: s
|
|||||||
check: async () => {
|
check: async () => {
|
||||||
const rows = await Database.use((tx) =>
|
const rows = await Database.use((tx) =>
|
||||||
tx
|
tx
|
||||||
.select({ count: IpRateLimitTable.count })
|
.select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count })
|
||||||
.from(IpRateLimitTable)
|
.from(IpRateLimitTable)
|
||||||
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, intervals))),
|
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, intervals))),
|
||||||
)
|
)
|
||||||
const total = rows.reduce((sum, r) => sum + r.count, 0)
|
const total = rows.reduce((sum, r) => sum + r.count, 0)
|
||||||
logger.debug(`rate limit total: ${total}`)
|
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) {
|
function buildYYYYMMDD(timestamp: number) {
|
||||||
return new Date(timestamp)
|
return new Date(timestamp)
|
||||||
.toISOString()
|
.toISOString()
|
||||||
|
|||||||
92
packages/console/app/test/rateLimiter.test.ts
Normal file
92
packages/console/app/test/rateLimiter.test.ts
Normal file
@@ -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)
|
||||||
|
})
|
||||||
|
})
|
||||||
Reference in New Issue
Block a user