Skip to content

Commit 801e4a8

Browse files
committed
wip: zen
1 parent 3adeed8 commit 801e4a8

3 files changed

Lines changed: 9 additions & 3 deletions

File tree

packages/console/app/src/routes/zen/util/handler.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ export async function handler(
7979
const dataDumper = createDataDumper(sessionId, requestId, projectId)
8080
const trialLimiter = createTrialLimiter(modelInfo.trial, ip, ocClient)
8181
const isTrial = await trialLimiter?.isTrial()
82-
const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip)
82+
const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip, input.request.headers)
8383
await rateLimiter?.check()
8484
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
8585
const stickyProvider = await stickyTracker?.get()

packages/console/app/src/routes/zen/util/rateLimiter.ts

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,13 @@ import { RateLimitError } from "./error"
44
import { logger } from "./logger"
55
import { ZenData } from "@opencode-ai/console-core/model.js"
66

7-
export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: string) {
7+
export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: string, headers: Headers) {
88
if (!limit) return
99

10+
const limitValue = (limit.checkHeader && !headers.get(limit.checkHeader))
11+
? limit.fallbackValue!
12+
: limit.value
13+
1014
const ip = !rawIp.length ? "unknown" : rawIp
1115
const now = Date.now()
1216
const intervals =
@@ -32,7 +36,7 @@ export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: s
3236
)
3337
const total = rows.reduce((sum, r) => sum + r.count, 0)
3438
logger.debug(`rate limit total: ${total}`)
35-
if (total >= limit.value) throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
39+
if (total >= limitValue) throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
3640
},
3741
}
3842
}

packages/console/core/src/model.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ export namespace ZenData {
2121
const RateLimitSchema = z.object({
2222
period: z.enum(["day", "rolling"]),
2323
value: z.number().int(),
24+
checkHeader: z.string().optional(),
25+
fallbackValue: z.number().int().optional(),
2426
})
2527
export type Format = z.infer<typeof FormatSchema>
2628
export type Trial = z.infer<typeof TrialSchema>

0 commit comments

Comments
 (0)